diff --git a/go/tasks/plugins/hive/client/qubole_client.go b/go/tasks/plugins/hive/client/qubole_client.go index 2edf1f8c11..3cf733cdb2 100644 --- 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 index aad0238fe7..6ac2d18c0d 100644 --- 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 index a2a839a13d..e52056c484 100644 --- 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)