get proper start time (#11)

* get proper start time * fix type * fix test * add undefined
eapolinario · Oct 9, 2019 · 922af83 · 922af83
1 parent cb881dd
commit 922af83
Show file tree

Hide file tree

Showing 3 changed files with 120 additions and 7 deletions.
diff --git a/Gopkg.lock b/Gopkg.lock
diff --git a/pkg/repositories/transformers/task_execution.go b/pkg/repositories/transformers/task_execution.go
@@ -18,7 +18,7 @@ type CreateTaskExecutionModelInput struct {
 	Request *admin.TaskExecutionEventRequest
 }
 
-func addTaskSubmittedState(request *admin.TaskExecutionEventRequest, taskExecutionModel *models.TaskExecution,
+func addTaskStartedState(request *admin.TaskExecutionEventRequest, taskExecutionModel *models.TaskExecution,
 	closure *admin.TaskExecutionClosure) error {
 	occurredAt, err := ptypes.Timestamp(request.Event.OccurredAt)
 	if err != nil {
@@ -99,8 +99,8 @@ func CreateTaskExecutionModel(input CreateTaskExecutionModelInput) (*models.Task
 	// Different tasks may report different phases as their first event.
 	// If the first event we receive for this execution is a valid
 	// non-terminal phase, mark the execution start time.
-	if eventPhase == core.TaskExecution_QUEUED || eventPhase == core.TaskExecution_RUNNING {
-		err := addTaskSubmittedState(input.Request, taskExecution, closure)
+	if eventPhase == core.TaskExecution_RUNNING {
+		err := addTaskStartedState(input.Request, taskExecution, closure)
 		if err != nil {
 			return nil, err
 		}
@@ -136,11 +136,19 @@ func UpdateTaskExecutionModel(request *admin.TaskExecutionEventRequest, taskExec
 		return errors.NewFlyteAdminErrorf(codes.Internal,
 			"failed to unmarshal task execution closure with error: %+v", err)
 	}
+	existingTaskPhase := taskExecutionModel.Phase
 	taskExecutionModel.Phase = request.Event.Phase.String()
 	taskExecutionModel.PhaseVersion = request.Event.PhaseVersion
 	taskExecutionClosure.Phase = request.Event.Phase
 	taskExecutionClosure.UpdatedAt = request.Event.OccurredAt
 	taskExecutionClosure.Logs = request.Event.Logs
+	if (existingTaskPhase == core.TaskExecution_QUEUED.String() || existingTaskPhase == core.TaskExecution_UNDEFINED.String()) && taskExecutionModel.Phase == core.TaskExecution_RUNNING.String() {
+		err = addTaskStartedState(request, taskExecutionModel, &taskExecutionClosure)
+		if err != nil {
+			return err
+		}
+	}
+
 	if common.IsTaskExecutionTerminal(request.Event.Phase) {
 		err := addTaskTerminalState(request, taskExecutionModel, &taskExecutionClosure)
 		if err != nil {

diff --git a/pkg/repositories/transformers/task_execution_test.go b/pkg/repositories/transformers/task_execution_test.go
@@ -47,7 +47,7 @@ var customInfo = ptypesStruct.Struct{
 	},
 }
 
-func TestAddTaskSubmittedState(t *testing.T) {
+func TestAddTaskStartedState(t *testing.T) {
 	var startedAt = time.Now().UTC()
 	var startedAtProto, _ = ptypes.TimestampProto(startedAt)
 	request := admin.TaskExecutionEventRequest{
@@ -58,7 +58,7 @@ func TestAddTaskSubmittedState(t *testing.T) {
 	}
 	taskExecutionModel := models.TaskExecution{}
 	closure := &admin.TaskExecutionClosure{}
-	err := addTaskSubmittedState(&request, &taskExecutionModel, closure)
+	err := addTaskStartedState(&request, &taskExecutionModel, closure)
 	assert.Nil(t, err)
 
 	timestamp, err := ptypes.Timestamp(closure.StartedAt)
@@ -140,7 +140,7 @@ func TestCreateTaskExecutionModelQueued(t *testing.T) {
 
 	expectedClosure := &admin.TaskExecutionClosure{
 		Phase:     core.TaskExecution_QUEUED,
-		StartedAt: taskEventOccurredAtProto,
+		StartedAt: nil,
 		CreatedAt: taskEventOccurredAtProto,
 		UpdatedAt: taskEventOccurredAtProto,
 	}
@@ -169,7 +169,7 @@ func TestCreateTaskExecutionModelQueued(t *testing.T) {
 		Phase:                  "QUEUED",
 		InputURI:               "input uri",
 		Closure:                expectedClosureBytes,
-		StartedAt:              &taskEventOccurredAt,
+		StartedAt:              nil,
 		TaskExecutionCreatedAt: &taskEventOccurredAt,
 		TaskExecutionUpdatedAt: &taskEventOccurredAt,
 	}, taskExecutionModel)