Specify a dataset location for vizier metrics

Summary: This ensures that our vizier metrics are colocated with cloud metrics. Test Plan: Deploy to testing cluster Reviewers: philkuz, jamesbartlett, michelle, nserrino Reviewed By: jamesbartlett JIRA Issues: PC-1512 Signed-off-by: Vihang Mehta <vihang@pixielabs.ai> Differential Revision: https://phab.corp.pixielabs.ai/D11382 GitOrigin-RevId: 5fe2e62
sgewirtz1 · May 6, 2022 · 32e8293 · 32e8293
1 parent 1607651
commit 32e8293
Show file tree

Hide file tree

Showing 6 changed files with 13 additions and 5 deletions.
diff --git a/k8s/cloud/base/bq_config.yaml b/k8s/cloud/base/bq_config.yaml
@@ -6,4 +6,5 @@ metadata:
 data:
   PL_BQ_PROJECT: ""
   PL_BQ_DATASET: ""
+  PL_BQ_DATASET_LOC: ""
   PL_BQ_SA_KEY_PATH: ""
diff --git a/k8s/cloud/dev/bq_config.yaml b/k8s/cloud/dev/bq_config.yaml
@@ -5,5 +5,6 @@ metadata:
   name: pl-bq-config
 data:
   PL_BQ_PROJECT: "pl-pixies"
-  PL_BQ_DATASET: "pixie_viziers_dev"
+  PL_BQ_DATASET: "pixie_vizier_dev"
+  PL_BQ_DATASET_LOC: "us-west1"
   PL_BQ_SA_KEY_PATH: "/creds/bq.client.default.credentials_file"
diff --git a/k8s/cloud/prod/bq_config.yaml b/k8s/cloud/prod/bq_config.yaml
@@ -5,5 +5,6 @@ metadata:
   name: pl-bq-config
 data:
   PL_BQ_PROJECT: "pixie-prod"
-  PL_BQ_DATASET: "pixie_viziers_prod"
+  PL_BQ_DATASET: "pixie_vizier_prod"
+  PL_BQ_DATASET_LOC: "us-west1"
   PL_BQ_SA_KEY_PATH: "/creds/bq.client.default.credentials_file"
diff --git a/k8s/cloud/staging/bq_config.yaml b/k8s/cloud/staging/bq_config.yaml
@@ -5,5 +5,6 @@ metadata:
   name: pl-bq-config
 data:
   PL_BQ_PROJECT: "pixie-prod"
-  PL_BQ_DATASET: "pixie_viziers_staging"
+  PL_BQ_DATASET: "pixie_vizier_staging"
+  PL_BQ_DATASET_LOC: "us-west1"
   PL_BQ_SA_KEY_PATH: "/creds/bq.client.default.credentials_file"
diff --git a/k8s/cloud/testing/bq_config.yaml b/k8s/cloud/testing/bq_config.yaml
@@ -5,5 +5,6 @@ metadata:
   name: pl-bq-config
 data:
   PL_BQ_PROJECT: "pl-pixies"
-  PL_BQ_DATASET: "pixie_viziers_testing"
+  PL_BQ_DATASET: "pixie_vizier_testing"
+  PL_BQ_DATASET_LOC: "us-west1"
   PL_BQ_SA_KEY_PATH: "/creds/bq.client.default.credentials_file"
diff --git a/src/cloud/metrics/metrics_server.go b/src/cloud/metrics/metrics_server.go
@@ -44,6 +44,7 @@ func init() {
 	pflag.String("bq_sa_key_path", "", "The service account for the BigQuery instance that should be used.")
 
 	pflag.String("bq_dataset", "vizier_metrics", "The BigQuery dataset to write metrics to.")
+	pflag.String("bq_dataset_loc", "", "The location for the BigQuery dataset. Used during creation.")
 }
 
 func main() {
@@ -77,8 +78,10 @@ func main() {
 			log.WithError(err).Fatal("Missing a BigQuery dataset name.")
 		}
 
+		dsLoc := viper.GetString("bq_dataset_loc")
+
 		dataset := client.Dataset(dsName)
-		err = dataset.Create(context.Background(), nil)
+		err = dataset.Create(context.Background(), &bigquery.DatasetMetadata{Location: dsLoc})
 		apiError, ok := err.(*googleapi.Error)
 		if !ok {
 			log.WithError(err).Fatal("Problem with BigQuery dataset")