Sets the lables to the qubole CommandMetadata (flyteorg#125)

eapolinario · Oct 8, 2020 · b5bbe10 · b5bbe10
1 parent 069476e
commit b5bbe10
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 5 deletions.
diff --git a/go/tasks/plugins/hive/client/qubole_client.go b/go/tasks/plugins/hive/client/qubole_client.go
@@ -44,9 +44,11 @@ type QuboleCommandDetails struct {
 }
 
 type CommandMetadata struct {
-	TaskName string
-	Domain   string
-	Project  string
+	TaskName     string
+	Domain       string
+	Project      string
+	Labels       map[string]string
+	RetryAttempt uint32
 }
 
 // QuboleClient API Request Body, meant to be passed into JSON.marshal

diff --git a/go/tasks/plugins/hive/execution_state.go b/go/tasks/plugins/hive/execution_state.go
@@ -344,8 +344,11 @@ func KickOffQuery(ctx context.Context, tCtx core.TaskExecutionContext, currentSt
 
 	taskExecutionIdentifier := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetID()
 	commandMetadata := client.CommandMetadata{TaskName: taskName,
-		Domain:  taskExecutionIdentifier.GetTaskId().GetDomain(),
-		Project: taskExecutionIdentifier.GetNodeExecutionId().GetExecutionId().GetProject()}
+		Domain:       taskExecutionIdentifier.GetTaskId().GetDomain(),
+		Project:      taskExecutionIdentifier.GetNodeExecutionId().GetExecutionId().GetProject(),
+		Labels:       tCtx.TaskExecutionMetadata().GetLabels(),
+		RetryAttempt: taskExecutionIdentifier.GetRetryAttempt(),
+	}
 
 	cmdDetails, err := quboleClient.ExecuteHiveCommand(ctx, query, timeoutSec,
 		clusterPrimaryLabel, apiKey, tags, commandMetadata)

diff --git a/go/tasks/plugins/hive/test_helpers.go b/go/tasks/plugins/hive/test_helpers.go
@@ -87,6 +87,7 @@ func GetMockTaskExecutionMetadata() core.TaskExecutionMetadata {
 				Domain:  "my_wf_exec_domain",
 			},
 		},
+		RetryAttempt: 1,
 	})
 	tID.On("GetGeneratedName").Return("my_wf_exec_project:my_wf_exec_domain:my_wf_exec_name")
 	taskMetadata.On("GetTaskExecutionID").Return(tID)