Ensure Athena Plugin ClientRequestToken is within allowed limits (fly…

…teorg#243) * Ensure Athena Plugin ClientRequestToken is within allowed limits Signed-off-by: Haytham Abuelfutuh <[email protected]> * Update interface Signed-off-by: Haytham Abuelfutuh <[email protected]> * Unit tests Signed-off-by: Haytham Abuelfutuh <[email protected]>
eapolinario · Feb 15, 2022 · ec275d2 · ec275d2
1 parent 4244f20
commit ec275d2
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 19 deletions.
diff --git a/flyteplugins/go/tasks/logs/logconfig_flags.go b/flyteplugins/go/tasks/logs/logconfig_flags.go
diff --git a/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go b/flyteplugins/go/tasks/pluginmachinery/core/exec_metadata.go
@@ -7,24 +7,32 @@ import (
 	"k8s.io/apimachinery/pkg/types"
 )
 
-// Interface to expose any overrides that have been set for this task (like resource overrides etc)
+// TaskOverrides interface to expose any overrides that have been set for this task (like resource overrides etc)
 type TaskOverrides interface {
 	GetResources() *v1.ResourceRequirements
 	GetConfig() *v1.ConfigMap
 }
 
-// Simple Interface to expose the ExecutionID of the running Task
+// TaskExecutionID is a simple Interface to expose the ExecutionID of the running Task
 type TaskExecutionID interface {
+	// GetGeneratedName returns the computed/generated name for the task id
+	// deprecated: use GetGeneratedNameWithLength
 	GetGeneratedName() string
+
+	// GetGeneratedNameWith returns the generated name within a bounded length. If the name is smaller than minLength,
+	// it'll get right-padded with character '0'. If the name is bigger than maxLength, it'll get hashed to fit within.
+	GetGeneratedNameWith(minLength, maxLength int) (string, error)
+
+	// GetID returns the underlying idl task identifier.
 	GetID() core.TaskExecutionIdentifier
 }
 
-// TaskContext represents any execution information for a Task. It is used to communicate meta information about the
+// TaskExecutionMetadata represents any execution information for a Task. It is used to communicate meta information about the
 // execution or any previously stored information
 type TaskExecutionMetadata interface {
-	// The owning Kubernetes object
+	// GetOwnerID returns the owning Kubernetes object
 	GetOwnerID() types.NamespacedName
-	// A specially generated task execution id, that is guaranteed to be unique and consistent for subsequent calls
+	// GetTaskExecutionID is a specially generated task execution id, that is guaranteed to be unique and consistent for subsequent calls
 	GetTaskExecutionID() TaskExecutionID
 	GetNamespace() string
 	GetOwnerReference() v12.OwnerReference

diff --git a/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_execution_id.go b/flyteplugins/go/tasks/pluginmachinery/core/mocks/task_execution_id.go
diff --git a/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go b/flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go
@@ -221,6 +221,10 @@ func (m mockTaskExecutionIdentifier) GetID() core.TaskExecutionIdentifier {
 	return testTaskExecutionIdentifier
 }
 
+func (m mockTaskExecutionIdentifier) GetGeneratedNameWith(minLength, maxLength int) (string, error) {
+	return "task-exec-name", nil
+}
+
 func (m mockTaskExecutionIdentifier) GetGeneratedName() string {
 	return "task-exec-name"
 }

diff --git a/flyteplugins/go/tasks/plugins/array/awsbatch/mocks/client.go b/flyteplugins/go/tasks/plugins/array/awsbatch/mocks/client.go
diff --git a/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go b/flyteplugins/go/tasks/plugins/webapi/athena/plugin.go
@@ -73,9 +73,17 @@ func (p Plugin) Create(ctx context.Context, tCtx webapi.TaskExecutionContextRead
 		return nil, nil, errors.Errorf(errors2.BadTaskSpecification, "Database must not be empty.")
 	}
 
-	execID := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetID().NodeExecutionId.GetExecutionId()
+	// https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html dictates that the length
+	// must be within the range [32, 128].
+	clientRequestToken, err := tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedNameWith(32, 128)
+	if err != nil {
+		return nil, nil, errors.Wrapf(errors2.BadTaskSpecification, err,
+			"Generated Name [%v] couldn't be converted to a ClientRequestToken",
+			tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName())
+	}
+
 	resp, err := p.client.StartQueryExecution(ctx, &athena.StartQueryExecutionInput{
-		ClientRequestToken: awsSdk.String(fmt.Sprintf("%v-%v-%v", execID.Project, execID.Domain, tCtx.TaskExecutionMetadata().GetTaskExecutionID().GetGeneratedName())),
+		ClientRequestToken: awsSdk.String(clientRequestToken),
 		QueryExecutionContext: &athenaTypes.QueryExecutionContext{
 			Database: awsSdk.String(queryInfo.Database),
 			Catalog:  awsSdk.String(queryInfo.Catalog),