feat: make name of secret consistent for bq (#39)

raystack · Jul 6, 2022 · d565549 · d565549
1 parent ea1de4f
commit d565549
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 10 deletions.
diff --git a/task/bq2bq/executor/bumblebee/bigquery_service.py b/task/bq2bq/executor/bumblebee/bigquery_service.py
@@ -16,7 +16,7 @@
 
 logger = get_logger(__name__)
 
-SERVICE_ACCOUNT_VAR = "TASK_BQ2BQ"
+SERVICE_ACCOUNT_VAR = "BQ_SERVICE_ACCOUNT"
 SERVICE_ACCOUNT_TYPE = "service_account"
 
 
@@ -166,7 +166,7 @@ def create_bigquery_service(task_config: TaskConfigFromEnv, labels, writer):
 
 
 def _get_bigquery_credentials():
-    """Gets credentials from the TASK_BQ2BQ environment var else GOOGLE_APPLICATION_CREDENTIALS for file path."""
+    """Gets credentials from the BQ_SERVICE_ACCOUNT environment var else GOOGLE_APPLICATION_CREDENTIALS for file path."""
     scope = ('https://www.googleapis.com/auth/bigquery',
              'https://www.googleapis.com/auth/cloud-platform',
              'https://www.googleapis.com/auth/drive')

diff --git a/task/bq2bq/main.go b/task/bq2bq/main.go
@@ -60,8 +60,9 @@ var (
 
 	QueryFileName = "query.sql"
 
-	// Required secret
-	SecretName = "TASK_BQ2BQ"
+	// Deprecated
+	SecretName       = "TASK_BQ2BQ"
+	BqServiceAccount = "BQ_SERVICE_ACCOUNT"
 
 	MaxBQApiRetries = 3
 	FakeSelectStmt  = "SELECT * from `%s` WHERE FALSE LIMIT 1"
@@ -405,9 +406,9 @@ func (b *BQ2BQ) GenerateDependencies(ctx context.Context, request models.Generat
 	}
 
 	var svcAcc string
-	accConfig, ok := request.Config.Get(SecretName)
+	accConfig, ok := request.Config.Get(BqServiceAccount)
 	if !ok || len(accConfig.Value) == 0 {
-		span.AddEvent("Required secret not found in config")
+		span.AddEvent("Required secret BQ_SERVICE_ACCOUNT not found in config")
 		// Fallback to project for getting the secret
 		svcAcc, ok = request.Project.Secret.GetByName(SecretName)
 		if !ok || len(svcAcc) == 0 {

diff --git a/task/bq2bq/main_test.go b/task/bq2bq/main_test.go
@@ -630,6 +630,10 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 						Name:  SecretName,
 						Value: "some_secret",
 					},
+					{
+						Name:  BqServiceAccount,
+						Value: "BQ_ACCOUNT_SECRET",
+					},
 				}),
 			}
 
@@ -661,7 +665,7 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 			defer client.AssertExpectations(t)
 
 			bqClientFac := new(bqClientFactoryMock)
-			bqClientFac.On("New", mock.Anything, "some_secret").Return(client, nil)
+			bqClientFac.On("New", mock.Anything, "BQ_ACCOUNT_SECRET").Return(client, nil)
 			defer bqClientFac.AssertExpectations(t)
 
 			b := &BQ2BQ{
@@ -702,6 +706,10 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 						Name:  SecretName,
 						Value: "some_secret",
 					},
+					{
+						Name:  BqServiceAccount,
+						Value: "BQ_ACCOUNT_SECRET",
+					},
 				}),
 			}
 
@@ -733,7 +741,7 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 			defer client.AssertExpectations(t)
 
 			bqClientFac := new(bqClientFactoryMock)
-			bqClientFac.On("New", mock.Anything, "some_secret").Return(client, nil)
+			bqClientFac.On("New", mock.Anything, "BQ_ACCOUNT_SECRET").Return(client, nil)
 			defer bqClientFac.AssertExpectations(t)
 
 			b := &BQ2BQ{
@@ -774,6 +782,10 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 						Name:  SecretName,
 						Value: "some_secret",
 					},
+					{
+						Name:  BqServiceAccount,
+						Value: "BQ_ACCOUNT_SECRET",
+					},
 				}),
 			}
 
@@ -805,7 +817,7 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 			defer client.AssertExpectations(t)
 
 			bqClientFac := new(bqClientFactoryMock)
-			bqClientFac.On("New", mock.Anything, "some_secret").Return(client, nil).Once()
+			bqClientFac.On("New", mock.Anything, "BQ_ACCOUNT_SECRET").Return(client, nil).Once()
 			defer bqClientFac.AssertExpectations(t)
 
 			b := &BQ2BQ{
@@ -864,6 +876,10 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 						Name:  SecretName,
 						Value: "some_secret",
 					},
+					{
+						Name:  BqServiceAccount,
+						Value: "BQ_ACCOUNT_SECRET",
+					},
 				}),
 			}
 
@@ -937,7 +953,7 @@ Select * from table where ts > "2021-01-16T00:00:00Z"`
 			defer client.AssertExpectations(t)
 
 			bqClientFac := new(bqClientFactoryMock)
-			bqClientFac.On("New", mock.Anything, "some_secret").Return(client, nil)
+			bqClientFac.On("New", mock.Anything, "BQ_ACCOUNT_SECRET").Return(client, nil)
 			defer bqClientFac.AssertExpectations(t)
 
 			b := &BQ2BQ{