feat: Leader aware routing (#1783)

* feat: leader aware routing * routetoleaderenabled property and test cases * feat: Added unit test cases * test: for LAR disabled * fix: review comments * fix: review comments * fix: review comments * fix : test cases * fix: lint * feat: Default disable LAR feature --------- Co-authored-by: Astha Mohta <[email protected]>
googleapis · Apr 26, 2023 · 0703f41 · 0703f41
1 parent 23ec060
commit 0703f41
Show file tree

Hide file tree

Showing 12 changed files with 423 additions and 35 deletions.
diff --git a/src/batch-transaction.ts b/src/batch-transaction.ts
@@ -20,8 +20,11 @@ import * as extend from 'extend';
 import * as is from 'is';
 import {Snapshot} from './transaction';
 import {google} from '../protos/protos';
-import {Session, Database} from '.';
-import {CLOUD_RESOURCE_HEADER} from '../src/common';
+import {Session, Database, Spanner} from '.';
+import {
+  CLOUD_RESOURCE_HEADER,
+  addLeaderAwareRoutingHeader,
+} from '../src/common';
 
 export interface TransactionIdentifier {
   session: string | Session;
@@ -133,12 +136,18 @@ class BatchTransaction extends Snapshot {
     delete reqOpts.gaxOptions;
     delete reqOpts.types;
 
+    const headers: {[k: string]: string} = {};
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(headers);
+    }
+
     this.createPartitions_(
       {
         client: 'SpannerClient',
         method: 'partitionQuery',
         reqOpts,
         gaxOpts: query.gaxOptions,
+        headers: headers,
       },
       callback
     );
@@ -225,12 +234,18 @@ class BatchTransaction extends Snapshot {
     delete reqOpts.keys;
     delete reqOpts.ranges;
 
+    const headers: {[k: string]: string} = {};
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(headers);
+    }
+
     this.createPartitions_(
       {
         client: 'SpannerClient',
         method: 'partitionRead',
         reqOpts,
         gaxOpts: options.gaxOptions,
+        headers: headers,
       },
       callback
     );

diff --git a/src/common.ts b/src/common.ts
@@ -75,3 +75,16 @@ export interface PagedOptionsWithFilter extends PagedOptions {
  * by the backend.
  */
 export const CLOUD_RESOURCE_HEADER = 'google-cloud-resource-prefix';
+
+/*!
+ * HTTP header to route the requests at Leader
+ */
+export const LEADER_AWARE_ROUTING_HEADER = 'x-goog-spanner-route-to-leader';
+
+/**
+ * Add Leader aware routing header to existing header list.
+ * @param headers Existing header list.
+ */
+export function addLeaderAwareRoutingHeader(headers: {[k: string]: string}) {
+  headers[LEADER_AWARE_ROUTING_HEADER] = 'true';
+}
diff --git a/src/database.ts b/src/database.ts
@@ -85,10 +85,11 @@ import {
   RequestCallback,
   ResourceCallback,
   Schema,
+  addLeaderAwareRoutingHeader,
 } from './common';
 import {Duplex, Readable, Transform} from 'stream';
 import {PreciseDate} from '@google-cloud/precise-date';
-import {EnumKey, RequestConfig, TranslateEnumKeys} from '.';
+import {EnumKey, RequestConfig, TranslateEnumKeys, Spanner} from '.';
 import arrify = require('arrify');
 import {ServiceError} from 'google-gax';
 import IPolicy = google.iam.v1.IPolicy;
@@ -523,13 +524,18 @@ class Database extends common.GrpcServiceObject {
       sessionCount: count,
     };
 
+    const headers = this.resourceHeader_;
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(headers);
+    }
+
     this.request<google.spanner.v1.IBatchCreateSessionsResponse>(
       {
         client: 'SpannerClient',
         method: 'batchCreateSessions',
         reqOpts,
         gaxOpts: options.gaxOptions,
-        headers: this.resourceHeader_,
+        headers: headers,
       },
       (err, resp) => {
         if (err) {
@@ -791,13 +797,18 @@ class Database extends common.GrpcServiceObject {
     reqOpts.session.creatorRole =
       options.databaseRole || this.databaseRole || null;
 
+    const headers = this.resourceHeader_;
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(headers);
+    }
+
     this.request<google.spanner.v1.ISession>(
       {
         client: 'SpannerClient',
         method: 'createSession',
         reqOpts,
         gaxOpts: options.gaxOptions,
-        headers: this.resourceHeader_,
+        headers: headers,
       },
       (err, resp) => {
         if (err) {
@@ -3261,6 +3272,17 @@ class Database extends common.GrpcServiceObject {
     const databaseName = name.split('/').pop();
     return instanceName + '/databases/' + databaseName;
   }
+
+  /**
+   * Gets the Spanner object
+   *
+   * @private
+   *
+   * @returns {Spanner}
+   */
+  private _getSpanner(): Spanner {
+    return this.instance.parent as Spanner;
+  }
 }
 
 /*! Developer Documentation

diff --git a/src/index.ts b/src/index.ts
@@ -106,11 +106,17 @@ export type GetInstanceConfigOperationsCallback = PagedCallback<
   instanceAdmin.spanner.admin.instance.v1.IListInstanceConfigOperationsResponse
 >;
 
+/**
+ * Session pool configuration options.
+ * @property {boolean} [routeToLeaderEnabled=True] If set to false leader aware routing will be disabled.
+ * Disabling leader aware routing would route all requests in RW/PDML transactions to any region.
+ */
 export interface SpannerOptions extends GrpcClientOptions {
   apiEndpoint?: string;
   servicePath?: string;
   port?: number;
   sslCreds?: grpc.ChannelCredentials;
+  routeToLeaderEnabled?: boolean;
 }
 export interface RequestConfig {
   client: string;
@@ -210,6 +216,7 @@ class Spanner extends GrpcService {
   projectIdReplaced_: boolean;
   projectFormattedName_: string;
   resourceHeader_: {[k: string]: string};
+  routeToLeaderEnabled = false;
 
   /**
    * Placeholder used to auto populate a column with the commit timestamp.
@@ -310,6 +317,11 @@ class Spanner extends GrpcService {
       packageJson: require('../../package.json'),
     } as {} as GrpcServiceConfig;
     super(config, options);
+
+    if (options.routeToLeaderEnabled === true) {
+      this.routeToLeaderEnabled = true;
+    }
+
     this.options = options;
     this.auth = new GoogleAuth(this.options);
     this.clients_ = new Map();

diff --git a/src/session.ts b/src/session.ts
@@ -36,9 +36,14 @@ import {
   CreateSessionOptions,
 } from './database';
 import {ServiceObjectConfig} from '@google-cloud/common';
-import {NormalCallback, CLOUD_RESOURCE_HEADER} from './common';
+import {
+  NormalCallback,
+  CLOUD_RESOURCE_HEADER,
+  addLeaderAwareRoutingHeader,
+} from './common';
 import {grpc, CallOptions} from 'google-gax';
 import IRequestOptions = google.spanner.v1.IRequestOptions;
+import {Spanner} from '.';
 
 export type GetSessionResponse = [Session, r.Response];
 
@@ -378,13 +383,18 @@ export class Session extends common.GrpcServiceObject {
     const reqOpts = {
       name: this.formattedName_,
     };
+
+    const headers = this.resourceHeader_;
+    if (this._getSpanner().routeToLeaderEnabled) {
+      addLeaderAwareRoutingHeader(headers);
+    }
     return this.request(
       {
         client: 'SpannerClient',
         method: 'getSession',
         reqOpts,
         gaxOpts,
-        headers: this.resourceHeader_,
+        headers: headers,
       },
       (err, resp) => {
         if (resp) {
@@ -512,6 +522,17 @@ export class Session extends common.GrpcServiceObject {
     const sessionName = name.split('/').pop();
     return databaseName + '/sessions/' + sessionName;
   }
+
+  /**
+   * Gets the Spanner object
+   *
+   * @private
+   *
+   * @returns {Spanner}
+   */
+  private _getSpanner(): Spanner {
+    return this.parent.parent.parent as Spanner;
+  }
 }
 
 /*! Developer Documentation