Presto - add option to use namespace as the user (flyteorg#96)

flyteplugins/go/tasks/plugins/presto/config/config.go

@@ -76,6 +76,7 @@ var (
 		Environment:         URLMustParse(""),
 		DefaultRoutingGroup: "adhoc",
 		DefaultUser:         "flyte-default-user",
+		UseNamespaceAsUser:  true,
 		RoutingGroupConfigs: []RoutingGroupConfig{{Name: "adhoc", Limit: 100}, {Name: "etl", Limit: 25}},
 		RefreshCacheConfig: RefreshCacheConfig{
 			Name:         "presto",
@@ -101,6 +102,7 @@ type Config struct {
 	Environment            config.URL           `json:"environment" pflag:",Environment endpoint for Presto to use"`
 	DefaultRoutingGroup    string               `json:"defaultRoutingGroup" pflag:",Default Presto routing group"`
 	DefaultUser            string               `json:"defaultUser" pflag:",Default Presto user"`
+	UseNamespaceAsUser     bool                 `json:"useNamespaceAsUser" pflag:",Use the K8s namespace as the user"`
 	RoutingGroupConfigs    []RoutingGroupConfig `json:"routingGroupConfigs" pflag:"-,A list of cluster configs. Each of the configs corresponds to a service cluster"`
 	RefreshCacheConfig     RefreshCacheConfig   `json:"refreshCacheConfig" pflag:"Refresh cache config"`
 	ReadRateLimiterConfig  RateLimiterConfig    `json:"readRateLimiterConfig" pflag:"Rate limiter config for read requests going to Presto"`

flyteplugins/go/tasks/plugins/presto/execution_state.go

@@ -93,6 +93,8 @@ type Query struct {
 	ExternalLocation  string                   `json:"externalLocation"`
 }
 
+const PrestoSource = "flyte"
+
 // This is the main state iteration
 func HandleExecutionState(
 	ctx context.Context,
@@ -296,6 +298,11 @@ func GetNextQuery(
 		if err != nil {
 			return Query{}, err
 		}
+		var user = getUser(ctx, prestoCfg.DefaultUser)
+
+		if prestoCfg.UseNamespaceAsUser {
+			user = tCtx.TaskExecutionMetadata().GetNamespace()
+		}
 
 		statement = fmt.Sprintf(`CREATE TABLE hive.flyte_temporary_tables."%s_temp" AS %s`, tempTableName, statement)
 
@@ -305,8 +312,8 @@ func GetNextQuery(
 				RoutingGroup: resolveRoutingGroup(ctx, routingGroup, prestoCfg),
 				Catalog:      catalog,
 				Schema:       schema,
-				Source:       "flyte",
-				User:         getUser(ctx, prestoCfg.DefaultUser),
+				Source:       PrestoSource,
+				User:         user,
 			},
 			TempTableName:     tempTableName + "_temp",
 			ExternalTableName: tempTableName + "_external",