From c1c0def689991c13a160806324aa798277863d34 Mon Sep 17 00:00:00 2001
From: Katrina Rogan <krogan@lyft.com>
Date: Wed, 18 Mar 2020 10:24:46 -0700
Subject: [PATCH] Add task type as a filterable field (#80)

---
 flyteadmin/pkg/repositories/config/migrations.go       | 10 ++++++++++
 flyteadmin/pkg/repositories/gormimpl/task_repo_test.go |  9 +++++++++
 flyteadmin/pkg/repositories/models/task.go             |  2 ++
 flyteadmin/pkg/repositories/transformers/task.go       |  5 +++++
 flyteadmin/pkg/repositories/transformers/task_test.go  |  1 +
 5 files changed, 27 insertions(+)

diff --git a/flyteadmin/pkg/repositories/config/migrations.go b/flyteadmin/pkg/repositories/config/migrations.go
index 1ca5857677..73bab81808 100644
--- a/flyteadmin/pkg/repositories/config/migrations.go
+++ b/flyteadmin/pkg/repositories/config/migrations.go
@@ -159,4 +159,14 @@ var Migrations = []*gormigrate.Migration{
 			return tx.DropTable("resources").Error
 		},
 	},
+	// Add Type to Task model.
+	{
+		ID: "2020-03-17-task-type",
+		Migrate: func(tx *gorm.DB) error {
+			return tx.AutoMigrate(&models.Task{}).Error
+		},
+		Rollback: func(tx *gorm.DB) error {
+			return tx.Exec("ALTER TABLE tasks DROP COLUMN IF EXISTS type").Error
+		},
+	},
 }
diff --git a/flyteadmin/pkg/repositories/gormimpl/task_repo_test.go b/flyteadmin/pkg/repositories/gormimpl/task_repo_test.go
index f3f9f337cb..690cf1aa23 100644
--- a/flyteadmin/pkg/repositories/gormimpl/task_repo_test.go
+++ b/flyteadmin/pkg/repositories/gormimpl/task_repo_test.go
@@ -16,6 +16,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
+const pythonTestTaskType = "python-task"
+
 func TestCreateTask(t *testing.T) {
 	taskRepo := NewTaskRepo(GetDbForTest(t), errors.NewTestErrorTransformer(), mockScope.NewTestScope())
 	err := taskRepo.Create(context.Background(), models.Task{
@@ -26,6 +28,7 @@ func TestCreateTask(t *testing.T) {
 			Version: version,
 		},
 		Closure: []byte{1, 2},
+		Type:    pythonTestTaskType,
 	})
 	assert.NoError(t, err)
 }
@@ -37,6 +40,7 @@ func getMockTaskResponseFromDb(version string, spec []byte) map[string]interface
 	task["name"] = name
 	task["version"] = version
 	task["closure"] = spec
+	task["type"] = pythonTestTaskType
 	return task
 }
 
@@ -66,6 +70,7 @@ func TestGetTask(t *testing.T) {
 	assert.Equal(t, name, output.Name)
 	assert.Equal(t, version, output.Version)
 	assert.Equal(t, []byte{1, 2}, output.Closure)
+	assert.Equal(t, pythonTestTaskType, output.Type)
 }
 
 func TestListTasks(t *testing.T) {
@@ -100,6 +105,7 @@ func TestListTasks(t *testing.T) {
 		assert.Equal(t, name, task.Name)
 		assert.Contains(t, versions, task.Version)
 		assert.Equal(t, spec, task.Closure)
+		assert.Equal(t, pythonTestTaskType, task.Type)
 	}
 }
 
@@ -135,6 +141,7 @@ func TestListTasks_Pagination(t *testing.T) {
 		assert.Equal(t, name, task.Name)
 		assert.Equal(t, versions[idx], task.Version)
 		assert.Equal(t, spec, task.Closure)
+		assert.Equal(t, pythonTestTaskType, task.Type)
 	}
 }
 
@@ -167,6 +174,7 @@ func TestListTasks_Filters(t *testing.T) {
 	assert.Equal(t, name, collection.Tasks[0].Name)
 	assert.Equal(t, "ABC", collection.Tasks[0].Version)
 	assert.Equal(t, []byte{1, 2}, collection.Tasks[0].Closure)
+	assert.Equal(t, pythonTestTaskType, collection.Tasks[0].Type)
 }
 
 func TestListTasks_Order(t *testing.T) {
@@ -246,6 +254,7 @@ func TestListTaskIds(t *testing.T) {
 		assert.Equal(t, name, task.Name)
 		assert.Equal(t, versions[idx], task.Version)
 		assert.Equal(t, spec, task.Closure)
+		assert.Equal(t, pythonTestTaskType, task.Type)
 	}
 }
 
diff --git a/flyteadmin/pkg/repositories/models/task.go b/flyteadmin/pkg/repositories/models/task.go
index 3ffe4e6072..06ff4a5e28 100644
--- a/flyteadmin/pkg/repositories/models/task.go
+++ b/flyteadmin/pkg/repositories/models/task.go
@@ -18,4 +18,6 @@ type Task struct {
 	Closure []byte `gorm:"not null"`
 	// Hash of the compiled task closure
 	Digest []byte
+	// Task type (also stored in the closure put promoted as a column for filtering).
+	Type string
 }
diff --git a/flyteadmin/pkg/repositories/transformers/task.go b/flyteadmin/pkg/repositories/transformers/task.go
index 0541c87b14..b43a2f0968 100644
--- a/flyteadmin/pkg/repositories/transformers/task.go
+++ b/flyteadmin/pkg/repositories/transformers/task.go
@@ -20,6 +20,10 @@ func CreateTaskModel(
 	if err != nil {
 		return models.Task{}, errors.NewFlyteAdminError(codes.Internal, "Failed to serialize task closure")
 	}
+	var taskType string
+	if taskClosure.CompiledTask != nil && taskClosure.CompiledTask.Template != nil {
+		taskType = taskClosure.CompiledTask.Template.Type
+	}
 	return models.Task{
 		TaskKey: models.TaskKey{
 			Project: request.Id.Project,
@@ -29,6 +33,7 @@ func CreateTaskModel(
 		},
 		Closure: closureBytes,
 		Digest:  digest,
+		Type:    taskType,
 	}, nil
 }
 
diff --git a/flyteadmin/pkg/repositories/transformers/task_test.go b/flyteadmin/pkg/repositories/transformers/task_test.go
index 1759e90f78..300b714da9 100644
--- a/flyteadmin/pkg/repositories/transformers/task_test.go
+++ b/flyteadmin/pkg/repositories/transformers/task_test.go
@@ -27,6 +27,7 @@ func TestCreateTask(t *testing.T) {
 	assert.Equal(t, "version", task.Version)
 	assert.Equal(t, testutils.GetTaskClosureBytes(), task.Closure)
 	assert.Equal(t, taskDigest, task.Digest)
+	assert.Equal(t, "type", task.Type)
 }
 
 func TestFromTaskModel(t *testing.T) {