Handle GCP multi-universe

cli/BUILD

@@ -74,6 +74,7 @@ ts_library(
 ts_test_suite(
     name = "tests",
     srcs = [
+        "credentials_test.ts",
         "index_help_test.ts",
         "index_init_test.ts",
         "index_project_test.ts",

cli/api/BUILD

@@ -62,6 +62,7 @@ ts_test_suite(
     srcs = [
         "tasks_test.ts",
         "utils_test.ts",
+        "commands/credentials_test.ts",
         "commands/jit/rpc_test.ts",
         "commands/prune_test.ts",
         "dbadapters/bigquery_test.ts",

cli/api/commands/credentials_test.ts

@@ -0,0 +1,47 @@
+import { expect } from "chai";
+import * as fs from "fs-extra";
+import * as path from "path";
+
+import { read } from "df/cli/api/commands/credentials";
+import { suite, test } from "df/testing";
+import { TmpDirFixture } from "df/testing/fixtures";
+
+suite("credentials", ({ afterEach }) => {
+  const tmpDirFixture = new TmpDirFixture(afterEach);
+
+  function writeCredentials(contents: object): string {
+    const credentialsPath = path.join(tmpDirFixture.createNewTmpDir(), ".df-credentials.json");
+    fs.writeFileSync(credentialsPath, JSON.stringify(contents));
+    return credentialsPath;
+  }
+
+  test("read maps universeDomain when present", () => {
+    const credentialsPath = writeCredentials({
+      projectId: "my-project",
+      location: "US",
+      universeDomain: "my-universe.example.com"
+    });
+
+    const credentials = read(credentialsPath);
+
+    expect(credentials.projectId).to.equal("my-project");
+    expect(credentials.location).to.equal("US");
+    expect(credentials.universeDomain).to.equal("my-universe.example.com");
+  });
+
+  test("read leaves universeDomain unset when omitted", () => {
+    const credentialsPath = writeCredentials({ projectId: "my-project", location: "US" });
+
+    const credentials = read(credentialsPath);
+
+    expect(credentials.universeDomain).to.satisfy(
+      (value: string) => value === "" || value === undefined
+    );
+  });
+
+  test("read rejects unknown fields", () => {
+    const credentialsPath = writeCredentials({ projectId: "my-project", notARealField: "x" });
+
+    expect(() => read(credentialsPath)).to.throw(/notARealField/);
+  });
+});

cli/api/dbadapters/bigquery.ts

@@ -52,7 +52,8 @@ export function createBigQueryClientProvider(
           projectId,
           scopes: EXTRA_GOOGLE_SCOPES,
           location: credentials.location,
-          credentials: credentials.credentials && JSON.parse(credentials.credentials)
+          credentials: credentials.credentials && JSON.parse(credentials.credentials),
+          universeDomain: credentials.universeDomain || undefined
         })
       );
     }

cli/api/dbadapters/bigquery_test.ts

@@ -2,7 +2,7 @@ import { Dataset, Table } from "@google-cloud/bigquery";
 import { expect } from "chai";
 import { anything, instance, mock, verify, when } from "ts-mockito";
 
-import { BigQueryDbAdapter } from "df/cli/api/dbadapters/bigquery";
+import { BigQueryDbAdapter, createBigQueryClientProvider } from "df/cli/api/dbadapters/bigquery";
 import { dataform } from "df/protos/ts";
 import { suite, test } from "df/testing";
 
@@ -146,4 +146,26 @@ suite("BigQueryDbAdapter", () => {
 
     await adapter.setMetadata(action);
   });
+
+  suite("createBigQueryClientProvider", () => {
+    test("passes universeDomain to the BigQuery client when set", () => {
+      const credentials = dataform.BigQuery.create({
+        projectId: "project1",
+        location: "US",
+        universeDomain: "my-universe.example.com"
+      });
+
+      const client = createBigQueryClientProvider(credentials)();
+
+      expect(client.universeDomain).to.equal("my-universe.example.com");
+    });
+
+    test("defaults to googleapis.com when universeDomain is unset", () => {
+      const credentials = dataform.BigQuery.create({ projectId: "project1", location: "US" });
+
+      const client = createBigQueryClientProvider(credentials)();
+
+      expect(client.universeDomain).to.equal("googleapis.com");
+    });
+  });
 });

cli/api/lineage/emitter.ts

@@ -42,7 +42,10 @@ export function createLineageClientProvider(
           apiEndpoint: endpoint,
           credentials: credentials.credentials && JSON.parse(credentials.credentials),
           libName: DATAFORM_CLI_LIB_NAME,
-          libVersion: version
+          libVersion: version,
+          // Omit empty universeDomain: gax uses ??, so "" would skip the default. (BigQuery
+          // falsy-checks options.universeDomain, so this guard is not needed there.)
+          ...(credentials.universeDomain ? { universeDomain: credentials.universeDomain } : {})
         })
       );
     }

cli/credentials.ts

@@ -8,38 +8,51 @@ export function getBigQueryCredentials(): dataform.IBigQuery {
   const locationIndex = selectionQuestion("Enter the location of your datasets:", [
     "US (default)",
     "EU",
-    "other"
+    "other",
   ]);
   let location = locationIndex === 0 ? "US" : "EU";
   if (locationIndex === 2) {
     location = question("Enter the location's region name (e.g. 'asia-south1'):");
   }
   const isApplicationDefaultOrJSONKeyIndex = selectionQuestion(
     "Do you wish to use Application Default Credentials or JSON Key:",
-    ["ADC (default)", "JSON Key"]
+    ["ADC (default)", "JSON Key"],
   );
   if (isApplicationDefaultOrJSONKeyIndex === 0) {
     const projectId = question("Enter your billing project ID:");
     return {
       projectId,
-      location
+      location,
     };
   }
   const cloudCredentialsPath = actuallyResolve(
     question(
       "Please follow the instructions at https://docs.dataform.co/dataform-cli#create-a-credentials-file/\n" +
         "to create and download a private key from the Google Cloud Console in JSON format.\n" +
         "(You can delete this file after credential initialization is complete.)\n\n" +
-        "Enter the path to your Google Cloud private key file:"
-    )
+        "Enter the path to your Google Cloud private key file:",
+    ),
   );
   if (!fs.existsSync(cloudCredentialsPath)) {
     throw new Error(`Google Cloud private key file "${cloudCredentialsPath}" does not exist!`);
   }
-  const cloudCredentials = JSON.parse(fs.readFileSync(cloudCredentialsPath, "utf8"));
+  return credentialsFromServiceAccountJson(fs.readFileSync(cloudCredentialsPath, "utf8"), location);
+}
+
+// Copy universe_domain from the key so BigQuery and google-auth target the same universe.
+export function credentialsFromServiceAccountJson(
+  keyJson: string,
+  location: string,
+): dataform.IBigQuery {
+  const cloudCredentials = JSON.parse(keyJson);
+  const universeDomain =
+    typeof cloudCredentials.universe_domain === "string"
+      ? cloudCredentials.universe_domain.trim()
+      : "";
   return {
     projectId: cloudCredentials.project_id,
-    credentials: fs.readFileSync(cloudCredentialsPath, "utf8"),
-    location
+    credentials: keyJson,
+    location,
+    ...(universeDomain ? { universeDomain } : {}),
   };
 }

cli/credentials_test.ts

@@ -0,0 +1,39 @@
+import { expect } from "chai";
+
+import { credentialsFromServiceAccountJson } from "df/cli/credentials";
+import { suite, test } from "df/testing";
+
+suite("credentialsFromServiceAccountJson", () => {
+  function keyJson(overrides: object = {}): string {
+    return JSON.stringify({
+      type: "service_account",
+      project_id: "my-project",
+      private_key: "fake-key",
+      ...overrides,
+    });
+  }
+
+  test("copies universe_domain from the key so BigQuery matches google-auth", () => {
+    const credentials = credentialsFromServiceAccountJson(
+      keyJson({ universe_domain: "my-universe.example.com" }),
+      "EU",
+    );
+
+    expect(credentials.projectId).to.equal("my-project");
+    expect(credentials.location).to.equal("EU");
+    expect(credentials.universeDomain).to.equal("my-universe.example.com");
+    expect(JSON.parse(credentials.credentials).universe_domain).to.equal("my-universe.example.com");
+  });
+
+  test("omits universeDomain when the key has no universe_domain", () => {
+    const credentials = credentialsFromServiceAccountJson(keyJson(), "US");
+
+    expect(credentials).to.not.have.property("universeDomain");
+  });
+
+  test("omits universeDomain when universe_domain is blank", () => {
+    const credentials = credentialsFromServiceAccountJson(keyJson({ universe_domain: "  " }), "US");
+
+    expect(credentials).to.not.have.property("universeDomain");
+  });
+});

contributing.md

@@ -57,6 +57,7 @@ To run the CLI integration test against your own GCP project:
    - `projectId`: your GCP project id
    - `credentials`: the entire content of your GCP service account key JSON file as a single string (you can generate it with `jq -Rsa < path/to/key.json`).
    - `location`: location to use in your project
+   - `universeDomain` (optional): the universe domain to connect to (e.g. `googleapis.com`). Leave unset to use the default Google Default Universe (GDU). Set this only when targeting a non-default universe such as a Trusted Partner Cloud (TPC).
 
    Example:
 

protos/profiles.proto

@@ -11,6 +11,10 @@ message BigQuery {
   string credentials = 3;
   // Options are listed here: https://cloud.google.com/bigquery/docs/locations
   string location = 4;
+  // The universe domain to connect to (e.g. "googleapis.com"). Leave unset to
+  // use the default Google Default Universe (GDU). Set this when targeting a
+  // Trusted Partner Cloud (TPC) or another non-default universe.
+  string universe_domain = 5;
 
   reserved 2;
 }