From 49b9c77737b51b1a3d1fc751813d852507d981ee Mon Sep 17 00:00:00 2001
From: Haytham AbuelFutuh <habuelfutuh@lyft.com>
Date: Mon, 11 May 2020 16:21:17 -0700
Subject: [PATCH] Make default cluster label configurable (#86)

* Make default cluster label configurable

* Unit tests
---
 .../go/tasks/plugins/hive/config/config.go    |  2 ++
 .../go/tasks/plugins/hive/execution_state.go  | 19 +++++++++++--------
 .../plugins/hive/execution_state_test.go      |  1 +
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/flyteplugins/go/tasks/plugins/hive/config/config.go b/flyteplugins/go/tasks/plugins/hive/config/config.go
index 25228525c7..da7e8576ab 100644
--- a/flyteplugins/go/tasks/plugins/hive/config/config.go
+++ b/flyteplugins/go/tasks/plugins/hive/config/config.go
@@ -47,6 +47,7 @@ var (
 		TokenKey:                  "FLYTE_QUBOLE_CLIENT_TOKEN",
 		LruCacheSize:              2000,
 		Workers:                   15,
+		DefaultClusterLabel:       "default",
 		ClusterConfigs:            []ClusterConfig{{PrimaryLabel: "default", Labels: []string{"default"}, Limit: 100, ProjectScopeQuotaProportionCap: 0.7, NamespaceScopeQuotaProportionCap: 0.7}},
 		DestinationClusterConfigs: []DestinationClusterConfig{},
 	}
@@ -62,6 +63,7 @@ type Config struct {
 	TokenKey                  string                     `json:"quboleTokenKey" pflag:",Name of the key where to find Qubole token in the secret manager."`
 	LruCacheSize              int                        `json:"lruCacheSize" pflag:",Size of the AutoRefreshCache"`
 	Workers                   int                        `json:"workers" pflag:",Number of parallel workers to refresh the cache"`
+	DefaultClusterLabel       string                     `json:"defaultClusterConfig" pflag:",The default cluster label. This will be used if label is not specified on the hive job."`
 	ClusterConfigs            []ClusterConfig            `json:"clusterConfigs" pflag:"-,A list of cluster configs. Each of the configs corresponds to a service cluster"`
 	DestinationClusterConfigs []DestinationClusterConfig `json:"destinationClusterConfigs" pflag:"-,A list configs specifying the destination service cluster for (project, domain)"`
 }
diff --git a/flyteplugins/go/tasks/plugins/hive/execution_state.go b/flyteplugins/go/tasks/plugins/hive/execution_state.go
index f4f511bfbb..c0caa5c15a 100644
--- a/flyteplugins/go/tasks/plugins/hive/execution_state.go
+++ b/flyteplugins/go/tasks/plugins/hive/execution_state.go
@@ -259,13 +259,13 @@ func GetQueryInfo(ctx context.Context, tCtx core.TaskExecutionContext) (
 	return
 }
 
-func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label string) (string, bool) {
-	primaryLabel := DefaultClusterPrimaryLabel
-	found := false
+func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label string) (primaryLabel string, found bool) {
+	primaryLabel = quboleCfg.DefaultClusterLabel
+	found = false
 
 	if label == "" {
-		logger.Debugf(ctx, "Input cluster label is an empty string; falling back to using the default primary label [%v]", label, DefaultClusterPrimaryLabel)
-		return primaryLabel, found
+		logger.Debugf(ctx, "Input cluster label is an empty string; falling back to using the default primary label [%v]", label, primaryLabel)
+		return
 	}
 
 	// Using a linear search because N is small and because of ClusterConfig's struct definition
@@ -280,8 +280,11 @@ func mapLabelToPrimaryLabel(ctx context.Context, quboleCfg *config.Config, label
 		}
 	}
 
-	logger.Debugf(ctx, "Cannot find the primary cluster label for label [%v] in configmap; "+
-		"falling back to using the default primary label [%v]", label, DefaultClusterPrimaryLabel)
+	if !found {
+		logger.Debugf(ctx, "Cannot find the primary cluster label for label [%v] in configmap; "+
+			"falling back to using the default primary label [%v]", label, primaryLabel)
+	}
+
 	return primaryLabel, found
 }
 
@@ -319,7 +322,7 @@ func getClusterPrimaryLabel(ctx context.Context, tCtx core.TaskExecutionContext,
 	}
 
 	// Else we return the default primary label
-	return DefaultClusterPrimaryLabel
+	return cfg.DefaultClusterLabel
 }
 
 func KickOffQuery(ctx context.Context, tCtx core.TaskExecutionContext, currentState ExecutionState, quboleClient client.QuboleClient,
diff --git a/flyteplugins/go/tasks/plugins/hive/execution_state_test.go b/flyteplugins/go/tasks/plugins/hive/execution_state_test.go
index 2acdd6d9bb..9ebd6a9b37 100644
--- a/flyteplugins/go/tasks/plugins/hive/execution_state_test.go
+++ b/flyteplugins/go/tasks/plugins/hive/execution_state_test.go
@@ -348,6 +348,7 @@ func TestKickOffQuery(t *testing.T) {
 
 func createMockQuboleCfg() *config.Config {
 	return &config.Config{
+		DefaultClusterLabel: "default",
 		ClusterConfigs: []config.ClusterConfig{
 			{PrimaryLabel: "primary A", Labels: []string{"primary A", "A", "label A", "A-prod"}, Limit: 10},
 			{PrimaryLabel: "primary B", Labels: []string{"B"}, Limit: 10},