feat: add support for spanner client resource base routing

src/batch-transaction.ts

@@ -20,7 +20,7 @@ import * as extend from 'extend';
 import * as is from 'is';
 import {Snapshot} from './transaction';
 import {google} from '../proto/spanner';
-import {Session} from '.';
+import {Session, Instance} from '.';
 
 export interface TransactionIdentifier {
   session: string | Session;
@@ -130,6 +130,7 @@ class BatchTransaction extends Snapshot {
       {
         client: 'SpannerClient',
         method: 'partitionQuery',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
         gaxOpts: query.gaxOptions,
       },
@@ -218,6 +219,7 @@ class BatchTransaction extends Snapshot {
       {
         client: 'SpannerClient',
         method: 'partitionRead',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
         gaxOpts: options.gaxOptions,
       },

src/common.ts

@@ -68,3 +68,11 @@ export type PagedRequest<P> = P & {
   maxApiCalls?: number;
   gaxOptions?: CallOptions;
 };
+
+export const PERMISSION_DENIED_WARNING_MESSAGE = `The client library attempted to connect to 
+an endpoint closer to your Cloud Spanner data but was unable to 
+do so. The client library will fallback to the API endpoint given 
+in the client options, which may result in increased latency.
+We recommend including the scope 
+https://www.googleapis.com/auth/spanner.admin 
+so that the client library can get an instance-specific endpoint and efficiently route requests.`;

src/database.ts

@@ -360,6 +360,7 @@ class Database extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'batchCreateSessions',
+        instanceId: ((this.parent as {}) as Instance).id,
         reqOpts,
       },
       (err, resp) => {
@@ -602,6 +603,7 @@ class Database extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'createSession',
+        instanceId: ((this.parent as {}) as Instance).id,
         reqOpts,
         gaxOpts,
       },
@@ -1131,6 +1133,7 @@ class Database extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'listSessions',
+        instanceId: ((this.parent as {}) as Instance).id,
         reqOpts,
         gaxOpts,
       },

src/index.ts

@@ -21,7 +21,7 @@ import {PreciseDate} from '@google-cloud/precise-date';
 import {replaceProjectIdToken} from '@google-cloud/projectify';
 import {promisifyAll} from '@google-cloud/promisify';
 import * as extend from 'extend';
-import {GoogleAuth, GoogleAuthOptions} from 'google-auth-library';
+import {GoogleAuth} from 'google-auth-library';
 import * as is from 'is';
 import * as path from 'path';
 import {common as p} from 'protobufjs';
@@ -35,13 +35,13 @@ import {SessionPool} from './session-pool';
 import {Table} from './table';
 import {PartitionedDml, Snapshot, Transaction} from './transaction';
 import {GrpcClientOptions} from 'google-gax';
-import {ChannelCredentials} from 'grpc';
+import {ChannelCredentials, status} from 'grpc';
 import {
   createGcpApiConfig,
   gcpCallInvocationTransformer,
   gcpChannelFactoryOverride,
 } from 'grpc-gcp';
-import {google} from '../proto/spanner';
+import {PERMISSION_DENIED_WARNING_MESSAGE} from './common';
 
 const grpc = require('grpc');
 
@@ -57,10 +57,12 @@ export interface SpannerOptions extends GrpcClientOptions {
   servicePath?: string;
   port?: number;
   sslCreds?: ChannelCredentials;
+  enableResourceBasedRouting?: boolean;
 }
 export interface RequestConfig {
   client: string;
   method: string;
+  instanceId?: string;
   // tslint:disable-next-line: no-any
   reqOpts: any;
   gaxOpts?: {};
@@ -197,7 +199,7 @@ export interface RequestConfig {
  * @param {ClientConfig} [options] Configuration options.
  */
 class Spanner extends GrpcService {
-  options: GoogleAuthOptions;
+  options: SpannerOptions;
   auth: GoogleAuth;
   clients_: Map<string, {}>;
   instances_: Map<string, Instance>;
@@ -302,7 +304,11 @@ class Spanner extends GrpcService {
     this.auth = new GoogleAuth(this.options);
     this.clients_ = new Map();
     this.instances_ = new Map();
-
+    this.options.enableResourceBasedRouting =
+      typeof options!.enableResourceBasedRouting === 'boolean'
+        ? options.enableResourceBasedRouting
+        : process.env.GOOGLE_CLOUD_SPANNER_ENABLE_RESOURCE_BASED_ROUTING ===
+          'true';
     /**
      * Get a list of {@link Instance} objects as a readable object stream.
      *
@@ -771,20 +777,95 @@ class Spanner extends GrpcService {
         callback(err);
         return;
       }
-      const clientName = config.client;
+      this.getGaxClient_(config, (err, gaxClient) => {
+        if (err) {
+          callback(err);
+          return;
+        }
+        let reqOpts = extend(true, {}, config.reqOpts);
+        reqOpts = replaceProjectIdToken(reqOpts, projectId!);
+        const requestFn = gaxClient[config.method].bind(
+          gaxClient,
+          reqOpts,
+          config.gaxOpts
+        );
+        callback(null, requestFn);
+      });
+    });
+  }
+
+  /**
+   * get GAX client. This will retrieve the end point uris for specific instance and cache the client accordingly.
+   *
+   * @private
+   *
+   * @param {string} projectId Request config
+   * @param {object} config Request config
+   * @param {function} callback Callback function
+   */
+  getGaxClient_(config, callback) {
+    let clientName = config.client;
+    if (
+      clientName !== 'SpannerClient' ||
+      this.options.apiEndpoint ||
+      !this.options.enableResourceBasedRouting
+    ) {
       if (!this.clients_.has(clientName)) {
-        this.clients_.set(clientName, new gapic.v1[clientName](this.options));
+        this.setSpannerClient(clientName, config, this.options);
       }
       const gaxClient = this.clients_.get(clientName)!;
-      let reqOpts = extend(true, {}, config.reqOpts);
-      reqOpts = replaceProjectIdToken(reqOpts, projectId!);
-      const requestFn = gaxClient[config.method].bind(
-        gaxClient,
-        reqOpts,
-        config.gaxOpts
+      callback(null, gaxClient);
+      return;
+    }
+
+    if (!config.instanceId) {
+      const error = new Error('instanceId is required.');
+      callback(error);
+      return;
+    }
+    const instanceId = config.instanceId;
+    clientName = `${clientName}-${instanceId}`;
+
+    if (!this.clients_.has(clientName)) {
+      this.instances_.get(instanceId)!.getMetadata(
+        {
+          fieldNames: ['endpointUris'],
+        },
+        (err, instance) => {
+          if (err) {
+            if (err.code === status.PERMISSION_DENIED) {
+              process.emitWarning(PERMISSION_DENIED_WARNING_MESSAGE);
+              this.setSpannerClient(clientName, config, this.options);
+              callback(null, this.clients_.get(clientName)!);
+              return;
+            } else {
+              callback(err);
+              return;
+            }
+          }
+          const options = Object.assign({}, this.options);
+          if ((instance as Instance)!.endpointUris!.length) {
+            options.apiEndpoint = (instance as Instance)!.endpointUris![0];
+          }
+          this.setSpannerClient(clientName, config, options);
+          callback(null, this.clients_.get(clientName)!);
+        }
       );
-      callback(null, requestFn);
-    });
+    } else {
+      callback(null, this.clients_.get(clientName)!);
+    }
+  }
+  /**
+   * Cache the GAX client with given name and configuration.
+   *
+   * @private
+   *
+   * @param {string} clientName Client name to cache.
+   * @param {object} config Request config
+   * @param {object} options Spanner options
+   */
+  setSpannerClient(clientName: string, config, options: SpannerOptions) {
+    this.clients_.set(clientName, new gapic.v1[config.client](options));
   }
 
   /**

src/session.ts

@@ -36,6 +36,7 @@ import {
 } from './database';
 import {ServiceObjectConfig, DeleteCallback} from '@google-cloud/common';
 import {NormalCallback} from './common';
+import {Instance} from '.';
 import {ServiceError} from 'grpc';
 
 export type GetSessionResponse = [Session, r.Response];
@@ -275,6 +276,7 @@ export class Session extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'deleteSession',
+        instanceId: ((this.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       callback!
@@ -325,6 +327,7 @@ export class Session extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'getSession',
+        instanceId: ((this.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       callback!
@@ -354,6 +357,7 @@ export class Session extends GrpcServiceObject {
       {
         client: 'SpannerClient',
         method: 'executeSql',
+        instanceId: ((this.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       callback!

src/transaction.ts

@@ -36,6 +36,7 @@ import {Key} from './table';
 import {SpannerClient as s} from './v1';
 import {google as spannerClient} from '../proto/spanner';
 import {NormalCallback} from './common';
+import {Instance} from '.';
 
 export type Rows = Array<Row | Json>;
 
@@ -299,6 +300,7 @@ export class Snapshot extends EventEmitter {
       {
         client: 'SpannerClient',
         method: 'beginTransaction',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       (
@@ -498,6 +500,7 @@ export class Snapshot extends EventEmitter {
       return this.requestStream({
         client: 'SpannerClient',
         method: 'streamingRead',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts: Object.assign({}, reqOpts, {resumeToken}),
         gaxOpts: gaxOptions,
       });
@@ -887,6 +890,7 @@ export class Snapshot extends EventEmitter {
       return this.requestStream({
         client: 'SpannerClient',
         method: 'executeStreamingSql',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts: Object.assign({}, reqOpts, {resumeToken}),
         gaxOpts: gaxOptions,
       });
@@ -1280,6 +1284,7 @@ export class Transaction extends Dml {
       {
         client: 'SpannerClient',
         method: 'executeBatchDml',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       (err: null | ServiceError, resp: s.ExecuteBatchDmlResponse) => {
@@ -1374,6 +1379,7 @@ export class Transaction extends Dml {
       {
         client: 'SpannerClient',
         method: 'commit',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       (err: null | Error, resp: spannerClient.spanner.v1.ICommitResponse) => {
@@ -1585,6 +1591,7 @@ export class Transaction extends Dml {
       {
         client: 'SpannerClient',
         method: 'rollback',
+        instanceId: ((this.session.parent.parent as {}) as Instance).id,
         reqOpts,
       },
       (err: null | ServiceError) => {