Cursor/bigquery universe domain clean 3190

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),
+          ...(credentials.universeDomain ? { universeDomain: credentials.universeDomain } : {})
         })
       );
     }

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,8 @@ export function createLineageClientProvider(
           apiEndpoint: endpoint,
           credentials: credentials.credentials && JSON.parse(credentials.credentials),
           libName: DATAFORM_CLI_LIB_NAME,
-          libVersion: version
+          libVersion: version,
+          ...(credentials.universeDomain ? { universeDomain: credentials.universeDomain } : {})
         })
       );
     }

cli/credentials.ts

@@ -14,6 +14,11 @@ export function getBigQueryCredentials(): dataform.IBigQuery {
   if (locationIndex === 2) {
     location = question("Enter the location's region name (e.g. 'asia-south1'):");
   }
+  const universeDomain = question(
+    "Enter the universe domain to connect to, or leave blank to use the default " +
+      "('googleapis.com'). Set this only when targeting a non-default universe such as a " +
+      "Trusted Partner Cloud (TPC):"
+  ).trim();
   const isApplicationDefaultOrJSONKeyIndex = selectionQuestion(
     "Do you wish to use Application Default Credentials or JSON Key:",
     ["ADC (default)", "JSON Key"]
@@ -22,7 +27,8 @@ export function getBigQueryCredentials(): dataform.IBigQuery {
     const projectId = question("Enter your billing project ID:");
     return {
       projectId,
-      location
+      location,
+      ...(universeDomain ? { universeDomain } : {})
     };
   }
   const cloudCredentialsPath = actuallyResolve(
@@ -40,6 +46,7 @@ export function getBigQueryCredentials(): dataform.IBigQuery {
   return {
     projectId: cloudCredentials.project_id,
     credentials: fs.readFileSync(cloudCredentialsPath, "utf8"),
-    location
+    location,
+    ...(universeDomain ? { 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;
 }