diff --git a/cli/api/BUILD b/cli/api/BUILD index 9b817c962..0d439e2fc 100644 --- a/cli/api/BUILD +++ b/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", diff --git a/cli/api/commands/credentials_test.ts b/cli/api/commands/credentials_test.ts new file mode 100644 index 000000000..bc205d3a2 --- /dev/null +++ b/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/); + }); +}); diff --git a/cli/api/dbadapters/bigquery.ts b/cli/api/dbadapters/bigquery.ts index 303e08024..102973a86 100644 --- a/cli/api/dbadapters/bigquery.ts +++ b/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 } : {}) }) ); } diff --git a/cli/api/dbadapters/bigquery_test.ts b/cli/api/dbadapters/bigquery_test.ts index 6659941e8..3915b61d7 100644 --- a/cli/api/dbadapters/bigquery_test.ts +++ b/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"); + }); + }); }); diff --git a/cli/api/lineage/emitter.ts b/cli/api/lineage/emitter.ts index eab5700e6..d06f48e05 100644 --- a/cli/api/lineage/emitter.ts +++ b/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 } : {}) }) ); } diff --git a/cli/credentials.ts b/cli/credentials.ts index 12e55c5e2..ef8b54fd5 100644 --- a/cli/credentials.ts +++ b/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 } : {}) }; } diff --git a/contributing.md b/contributing.md index 96096a514..47901c578 100644 --- a/contributing.md +++ b/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: diff --git a/protos/profiles.proto b/protos/profiles.proto index 09eabcdb8..08d89cef4 100644 --- a/protos/profiles.proto +++ b/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; }