From f67a841c29f1f546187e4b670f5e5a4df50b4926 Mon Sep 17 00:00:00 2001
From: Alvaro Viebrantz <aviebrantz@google.com>
Date: Mon, 12 Feb 2024 11:29:39 -0800
Subject: [PATCH] feat: add support for TPC Universes (#1333)

---
 samples/package.json |  2 +-
 src/bigquery.ts      | 25 ++++++++++++++++++++++++-
 test/bigquery.ts     | 16 ++++++++++++++++
 3 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/samples/package.json b/samples/package.json
index aa0a1421..34790994 100644
--- a/samples/package.json
+++ b/samples/package.json
@@ -19,7 +19,7 @@
   "dependencies": {
     "@google-cloud/bigquery": "^7.4.0",
     "@google-cloud/storage": "^7.0.0",
-    "google-auth-library": "^9.0.0",
+    "google-auth-library": "^9.6.0",
     "readline-promise": "^1.0.4",
     "yargs": "^17.0.0"
   },
diff --git a/src/bigquery.ts b/src/bigquery.ts
index c2a2983d..f3e14113 100644
--- a/src/bigquery.ts
+++ b/src/bigquery.ts
@@ -227,6 +227,12 @@ export interface BigQueryOptions extends GoogleAuthOptions {
    * Defaults to `bigquery.googleapis.com`.
    */
   apiEndpoint?: string;
+
+  /**
+   * The Trusted Cloud Domain (TPC) DNS of the service used to make requests.
+   * Defaults to `googleapis.com`.
+   */
+  universeDomain?: string;
 }
 
 export interface IntegerTypeCastOptions {
@@ -311,6 +317,7 @@ export const PROTOCOL_REGEX = /^(\w*):\/\//;
  */
 export class BigQuery extends Service {
   location?: string;
+  private _universeDomain: string;
 
   createQueryStream(options?: Query | string): ResourceStream<RowMetadata> {
     // placeholder body, overwritten in constructor
@@ -328,10 +335,17 @@ export class BigQuery extends Service {
   }
 
   constructor(options: BigQueryOptions = {}) {
-    let apiEndpoint = 'https://bigquery.googleapis.com';
+    let universeDomain = 'googleapis.com';
+    const servicePath = 'bigquery';
+
+    if (options.universeDomain) {
+      universeDomain = BigQuery.sanitizeDomain(options.universeDomain);
+    }
 
     const EMULATOR_HOST = process.env.BIGQUERY_EMULATOR_HOST;
 
+    let apiEndpoint = `https://${servicePath}.${universeDomain}`;
+
     if (typeof EMULATOR_HOST === 'string') {
       apiEndpoint = BigQuery.sanitizeEndpoint(EMULATOR_HOST);
     }
@@ -361,6 +375,7 @@ export class BigQuery extends Service {
 
     super(config, options);
 
+    this._universeDomain = universeDomain;
     this.location = options.location;
     /**
      * Run a query scoped to your project as a readable object stream.
@@ -473,10 +488,18 @@ export class BigQuery extends Service {
     });
   }
 
+  get universeDomain() {
+    return this._universeDomain;
+  }
+
   private static sanitizeEndpoint(url: string) {
     if (!PROTOCOL_REGEX.test(url)) {
       url = `https://${url}`;
     }
+    return this.sanitizeDomain(url);
+  }
+
+  private static sanitizeDomain(url: string) {
     return url.replace(/\/+$/, ''); // Remove trailing slashes
   }
 
diff --git a/test/bigquery.ts b/test/bigquery.ts
index 1a6f51b9..34669d1b 100644
--- a/test/bigquery.ts
+++ b/test/bigquery.ts
@@ -263,6 +263,22 @@ describe('BigQuery', () => {
       assert.strictEqual(calledWith.apiEndpoint, 'https://some.fake.endpoint');
     });
 
+    it('should allow overriding TPC universe', () => {
+      const universeDomain = 'fake-tpc-env.example.com/';
+      bq = new BigQuery({
+        universeDomain: universeDomain,
+      });
+      const calledWith = bq.calledWith_[0];
+      assert.strictEqual(
+        calledWith.baseUrl,
+        'https://bigquery.fake-tpc-env.example.com/bigquery/v2'
+      );
+      assert.strictEqual(
+        calledWith.apiEndpoint,
+        'https://bigquery.fake-tpc-env.example.com'
+      );
+    });
+
     it('should capture any user specified location', () => {
       const bq = new BigQuery({
         projectId: PROJECT_ID,