From 32e82931f9428da2349e0d213294e236ecc3d77c Mon Sep 17 00:00:00 2001
From: Vihang Mehta <vihang@pixielabs.ai>
Date: Fri, 6 May 2022 10:59:40 -0700
Subject: [PATCH] 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: 5fe2e62620576c60b837de8f93888c55244508e5
---
 k8s/cloud/base/bq_config.yaml       | 1 +
 k8s/cloud/dev/bq_config.yaml        | 3 ++-
 k8s/cloud/prod/bq_config.yaml       | 3 ++-
 k8s/cloud/staging/bq_config.yaml    | 3 ++-
 k8s/cloud/testing/bq_config.yaml    | 3 ++-
 src/cloud/metrics/metrics_server.go | 5 ++++-
 6 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/k8s/cloud/base/bq_config.yaml b/k8s/cloud/base/bq_config.yaml
index 8a31829cc97..ea7c2391d23 100644
--- 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
index b0694a5e363..5b9f90b4f76 100644
--- 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
index 590ad3d1df7..0f6cadcef06 100644
--- 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
index 9ac134902f5..b02a732a6b5 100644
--- 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
index becd5afbc68..30597b3e6ac 100644
--- 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
index 30ec94ba702..660b82989cc 100644
--- 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")