Add custom marshaler support for lists

database/boltdb/tes.go

@@ -115,6 +115,13 @@ func (taskBolt *BoltDB) GetTask(ctx context.Context, req *tes.GetTaskRequest) (*
 	var err error
 
 	err = taskBolt.db.View(func(tx *bolt.Tx) error {
+		if req.View == "" {
+			req.View = tes.View_MINIMAL.String()
+		}
+		_, ok := tes.View_value[req.View]
+		if !ok {
+			return fmt.Errorf("Unknown view: %s", req.View)
+		}
 		task, err = getTaskView(tx, req.Id, tes.View(tes.View_value[req.View]))
 		return err
 	})
@@ -123,6 +130,7 @@ func (taskBolt *BoltDB) GetTask(ctx context.Context, req *tes.GetTaskRequest) (*
 
 func getTaskView(tx *bolt.Tx, id string, view tes.View) (*tes.Task, error) {
 	var err error
+
 	task := &tes.Task{}
 
 	switch {

go.sum

@@ -47,8 +47,6 @@ github.com/bradfitz/gomemcache v0.0.0-20170208213004-1952afaa557d/go.mod h1:PmM6
 github.com/cenkalti/backoff v2.2.1+incompatible h1:tNowT99t7UNflLxfYYSlKYsBpXdEet03Pg2g16Swow4=
 github.com/cenkalti/backoff v2.2.1+incompatible/go.mod h1:90ReRw6GdpyfrHakVjL/QHaoyV4aDUVVkXQJJJ3NXXM=
 github.com/census-instrumentation/opencensus-proto v0.2.1/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
-github.com/census-instrumentation/opencensus-proto v0.3.0/go.mod h1:f6KPmirojxKA12rnyqOA5BBL4O983OfeGPqjHWSTneU=
-github.com/census-instrumentation/opencensus-proto v0.4.1/go.mod h1:4T9NM4+4Vw91VeyqjLS6ao50K5bOcLKN6Q42XnYaRYw=
 github.com/cespare/xxhash v1.1.0 h1:a6HrQnmkObjyL+Gs60czilIUGqrzKutQD6XZog3p+ko=
 github.com/cespare/xxhash v1.1.0/go.mod h1:XrSqR1VqqWfGrhpAt58auRo0WTKS1nRRg3ghfAqPWnc=
 github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
@@ -113,7 +111,6 @@ github.com/frankban/quicktest v1.14.4 h1:g2rn0vABPOOXmZUj+vbmUp0lPoXEMuhTpIluN0X
 github.com/frankban/quicktest v1.14.4/go.mod h1:4ptaffx2x8+WTWXmUCuVU6aPUX1/Mz7zb5vbUoiM6w0=
 github.com/fsnotify/fsnotify v1.4.3-0.20170329110642-4da3e2cfbabc/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
-github.com/fsnotify/fsnotify v1.4.9/go.mod h1:znqG4EE+3YCdAaPaxE2ZRY/06pZUdp0tY4IgpuI1SZQ=
 github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622 h1:lxbhOGZ9pU3Kf8P6lFluUcE82yVZn2EqEf4+mWRNPV0=
 github.com/gammazero/deque v0.0.0-20190521012701-46e4ffb7a622/go.mod h1:D90+MBHVc9Sk1lJAbEVgws0eYEurY4mv2TDso3Nxh3w=
 github.com/gammazero/workerpool v0.0.0-20200206003619-019d125201ab h1:BLCWHqLiTSzLZriw46b3HJLswyqY0E6B9y9tk4+7eFI=
@@ -268,10 +265,6 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
 github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
 github.com/logrusorgru/aurora v0.0.0-20200102142835-e9ef32dff381 h1:bqDmpDG49ZRnB5PcgP0RXtQvnMSgIF14M7CBd2shtXs=
 github.com/logrusorgru/aurora v0.0.0-20200102142835-e9ef32dff381/go.mod h1:7rIyQOR62GCctdiQpZ/zOJlFyk6y+94wXzv6RNZgaR4=
-github.com/lyft/protoc-gen-star v0.6.0/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=
-github.com/lyft/protoc-gen-star v0.6.1/go.mod h1:TGAoBVkt8w7MPG72TrKIu85MIdXwDuzJYeZuUPFPNwA=
-github.com/lyft/protoc-gen-star/v2 v2.0.1/go.mod h1:RcCdONR2ScXaYnQC5tUzxzlpA3WVYF7/opLeUgcQs/o=
-github.com/lyft/protoc-gen-star/v2 v2.0.3/go.mod h1:amey7yeodaJhXSbf/TlLvWiqQfLOSpEk//mLlc+axEk=
 github.com/magiconair/properties v1.7.4-0.20170902060319-8d7837e64d3c/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/magiconair/properties v1.8.0/go.mod h1:PppfXfuXeibc/6YijjN8zIbojt8czPbwD3XqdrwzmxQ=
 github.com/mailru/easyjson v0.0.0-20180730094502-03f2033d19d5/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=

server/marshal.go

@@ -39,21 +39,83 @@ func (mclean *MarshalNew) ContentType(i interface{}) string {
 // itself. This is mainly to get around a weird behavior of the GRPC gateway
 // streaming output
 func (mclean *MarshalNew) Marshal(v interface{}) ([]byte, error) {
-    // Type assertion to get the underlying *tes.Task
-    task, ok := v.(*tes.Task)
-    if !ok {
-        // v is not of type *tes.Task
-		return mclean.m.Marshal(v)
+
+	list, ok := v.(*tes.ListTasksResponse)
+    if ok {
+        // v is of type *tes.ListTasksResponse
+		return mclean.MarshalList(list)
     }
+
+	task, ok := v.(*tes.Task)
+	if ok {
+		// v is of type *tes.Task
+		return mclean.MarshalTask(task)
+	}
+
+	return mclean.m.Marshal(v)
+}
+
+func (mclean *MarshalNew) MarshalTask(task *tes.Task) ([]byte, error) {
+	view, _ := mclean.DetectView(task)
+	newTask := mclean.TranslateTask(task, view)
+	return mclean.m.Marshal(newTask)
+}
+
+func (mclean *MarshalNew) MarshalList(list *tes.ListTasksResponse) ([]byte, error) {
+	if len(list.Tasks) == 0 {
+		return mclean.m.Marshal(list)
+	}
+
+	task := list.Tasks[0]
+	view, _ := mclean.DetectView(task)
+
+	if view == tes.View_MINIMAL {
+		minList := &tes.ListTasksResponseMin{}
+		for _, task := range list.Tasks {
+			minTask := mclean.TranslateTask(task, view).(*tes.TaskMin)
+			minList.Tasks = append(minList.Tasks, minTask)
+		}
+		return mclean.m.Marshal(minList)
+	}
+
+	if view == tes.View_BASIC {
+		basicList := &tes.ListTasksResponseBasic{}
+		for _, task := range list.Tasks {
+			basicTask := mclean.TranslateTask(task, view).(*tes.TaskBasic)
+			basicList.Tasks = append(basicList.Tasks, basicTask)
+		}
+		return mclean.m.Marshal(basicList)
+	}
+
+	return mclean.m.Marshal(list)
+}
+
+func (mclean *MarshalNew) DetectView(task *tes.Task) (tes.View, error) {
 	if task.CreationTime == "" {
-		// view = "MINIMAL"
+		// return a MINIMAL view
+		return tes.View_MINIMAL, nil
+	}
+
+	if len(task.Logs[0].SystemLogs) == 0 {
+		return tes.View_BASIC, nil
+	} 
+
+	// view = "FULL"
+	return tes.View_FULL, nil
+}
+
+func (mclean *MarshalNew) TranslateTask(task *tes.Task, view tes.View) interface{} {
+	// view = "MINIMAL"
+	if view == tes.View_MINIMAL {
 		min := &tes.TaskMin{
 			Id:    task.Id,
 			State: task.State,
 		}
-		return mclean.m.Marshal(min)
-	} else if len(task.Logs[0].SystemLogs) == 0 {
-		// view = "BASIC"
+		return min
+	}
+
+	// view = "BASIC"
+	if view == tes.View_BASIC {
 		executors := []*tes.ExecutorBasic{}
 		for _, executor := range task.Executors {
 			executors = append(executors, &tes.ExecutorBasic{
@@ -89,7 +151,7 @@ func (mclean *MarshalNew) Marshal(v interface{}) ([]byte, error) {
 			})
 		}
 
-		basic := &tes.TaskBasic{
+		basic := &tes.TaskBasic {
 			CreationTime: task.CreationTime,
 			Description:  task.Description,
 			Executors:    executors,
@@ -104,11 +166,11 @@ func (mclean *MarshalNew) Marshal(v interface{}) ([]byte, error) {
 			Volumes: 	  task.Volumes,
 		}
 
-		return mclean.m.Marshal(basic)
-	} else {
-		// view = "FULL"
-		return mclean.m.Marshal(v)
+		return basic
 	}
+
+	// view = "FULL"
+	return task
 }
 
 // NewDecoder shims runtime.Marshaler.NewDecoder