feat(spanner): add support for snapshot isolation (#2245)

alkatrivedi · web-flow · commit b60a683c0e1d · 2025-03-20T10:14:21.000Z
This PR contains code changes to add support for option IsolationLevel at the client level and at the transaction level.
supported methods are(RW and Blind Write):

```
- writeAtLeastOnce
- runTransactionAsync
- runTransaction
- getTransaction
- async getTransaction(from transaction runner class)
```
diff --git a/observability-test/database.ts b/observability-test/database.ts
@@ -48,6 +48,7 @@ import * as pfy from '@google-cloud/promisify';
 import {grpc} from 'google-gax';
 import {MockError} from '../test/mockserver/mockspanner';
 import {FakeSessionFactory} from '../test/database';
+import {RunTransactionOptions} from '../src/transaction-runner';
 const {generateWithAllSpansHaveDBName} = require('./helper');
 
 const fakePfy = extend({}, pfy, {
@@ -147,6 +148,7 @@ class FakeTransaction extends EventEmitter {
   setQueuedMutations(mutation) {
     this._queuedMutations = mutation;
   }
+  setReadWriteTransactionOptions(options: RunTransactionOptions) {}
   commit(
     options?: CommitOptions,
     callback?: CommitCallback
diff --git a/src/database.ts b/src/database.ts
@@ -42,6 +42,7 @@ import {
   google,
   google as spannerClient,
 } from '../protos/protos';
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
 import {
   CreateDatabaseCallback,
   CreateDatabaseOptions,
@@ -316,6 +317,10 @@ export interface RestoreOptions {
   gaxOptions?: CallOptions;
 }
 
+export interface WriteAtLeastOnceOptions extends CallOptions {
+  isolationLevel?: IsolationLevel;
+}
+
 /**
  * Create a Database object to interact with a Cloud Spanner database.
  *
@@ -2221,12 +2226,9 @@ class Database extends common.GrpcServiceObject {
             options.requestOptions
           );
         }
-        if (options.optimisticLock) {
-          transaction!.useOptimisticLock();
-        }
-        if (options.excludeTxnFromChangeStreams) {
-          transaction!.excludeTxnFromChangeStreams();
-        }
+        transaction?.setReadWriteTransactionOptions(
+          options as RunTransactionOptions
+        );
 
         if (!err) {
           span.addEvent('Using Session', {'session.id': session?.id});
@@ -3274,12 +3276,10 @@ class Database extends common.GrpcServiceObject {
         }
 
         transaction!._observabilityOptions = this._observabilityOptions;
-        if (options.optimisticLock) {
-          transaction!.useOptimisticLock();
-        }
-        if (options.excludeTxnFromChangeStreams) {
-          transaction!.excludeTxnFromChangeStreams();
-        }
+
+        transaction!.setReadWriteTransactionOptions(
+          options as RunTransactionOptions
+        );
 
         const release = () => {
           this.pool_.release(session!);
@@ -3406,12 +3406,9 @@ class Database extends common.GrpcServiceObject {
               transaction.requestOptions || {},
               options.requestOptions
             );
-            if (options.optimisticLock) {
-              transaction.useOptimisticLock();
-            }
-            if (options.excludeTxnFromChangeStreams) {
-              transaction.excludeTxnFromChangeStreams();
-            }
+            transaction!.setReadWriteTransactionOptions(
+              options as RunTransactionOptions
+            );
             sessionId = session?.id;
             span.addEvent('Using Session', {'session.id': sessionId});
             const runner = new AsyncTransactionRunner<T>(
@@ -3638,17 +3635,17 @@ class Database extends common.GrpcServiceObject {
   writeAtLeastOnce(mutations: MutationSet): Promise<CommitResponse>;
   writeAtLeastOnce(
     mutations: MutationSet,
-    options: CallOptions
+    options: WriteAtLeastOnceOptions
   ): Promise<CommitResponse>;
   writeAtLeastOnce(mutations: MutationSet, callback: CommitCallback): void;
   writeAtLeastOnce(
     mutations: MutationSet,
-    options: CallOptions,
+    options: WriteAtLeastOnceOptions,
     callback: CommitCallback
   ): void;
   writeAtLeastOnce(
     mutations: MutationSet,
-    optionsOrCallback?: CallOptions | CommitCallback,
+    optionsOrCallback?: WriteAtLeastOnceOptions | CommitCallback,
     callback?: CommitCallback
   ): void | Promise<CommitResponse> {
     const cb =
@@ -3657,7 +3654,7 @@ class Database extends common.GrpcServiceObject {
         : callback;
     const options =
       typeof optionsOrCallback === 'object' && optionsOrCallback
-        ? (optionsOrCallback as CallOptions)
+        ? (optionsOrCallback as WriteAtLeastOnceOptions)
         : {};
 
     return startTrace('Database.writeAtLeastOnce', this._traceConfig, span => {
@@ -3683,6 +3680,9 @@ class Database extends common.GrpcServiceObject {
         span.addEvent('Using Session', {'session.id': session?.id});
         this._releaseOnEnd(session!, transaction!, span);
         try {
+          transaction!.setReadWriteTransactionOptions(
+            options as RunTransactionOptions
+          );
           transaction?.setQueuedMutations(mutations.proto());
           return transaction?.commit(options, (err, resp) => {
             if (err) {
diff --git a/src/index.ts b/src/index.ts
@@ -60,6 +60,7 @@ import {
   ClientOptions,
 } from 'google-gax';
 import {google, google as instanceAdmin} from '../protos/protos';
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
 import {
   PagedOptions,
   PagedResponse,
@@ -145,6 +146,7 @@ export interface SpannerOptions extends GrpcClientOptions {
   sslCreds?: grpc.ChannelCredentials;
   routeToLeaderEnabled?: boolean;
   directedReadOptions?: google.spanner.v1.IDirectedReadOptions | null;
+  defaultTransactionOptions?: Pick<RunTransactionOptions, 'isolationLevel'>;
   observabilityOptions?: ObservabilityOptions;
 }
 export interface RequestConfig {
@@ -247,6 +249,7 @@ class Spanner extends GrpcService {
   commonHeaders_: {[k: string]: string};
   routeToLeaderEnabled = true;
   directedReadOptions: google.spanner.v1.IDirectedReadOptions | null;
+  defaultTransactionOptions: RunTransactionOptions;
   _observabilityOptions: ObservabilityOptions | undefined;
 
   /**
@@ -331,6 +334,13 @@ class Spanner extends GrpcService {
       : null;
     delete options.directedReadOptions;
 
+    const defaultTransactionOptions = options.defaultTransactionOptions
+      ? options.defaultTransactionOptions
+      : {
+          isolationLevel: IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
+        };
+    delete options.defaultTransactionOptions;
+
     const emulatorHost = Spanner.getSpannerEmulatorHost();
     if (
       emulatorHost &&
@@ -371,6 +381,7 @@ class Spanner extends GrpcService {
     this.projectIdReplaced_ = false;
     this.projectFormattedName_ = 'projects/' + this.projectId;
     this.directedReadOptions = directedReadOptions;
+    this.defaultTransactionOptions = defaultTransactionOptions;
     this._observabilityOptions = options.observabilityOptions;
     this.commonHeaders_ = getCommonHeaders(
       this.projectFormattedName_,
@@ -2068,6 +2079,7 @@ export {MutationSet};
  */
 import * as protos from '../protos/protos';
 import IInstanceConfig = instanceAdmin.spanner.admin.instance.v1.IInstanceConfig;
+import {RunTransactionOptions} from './transaction-runner';
 export {v1, protos};
 export default {Spanner};
 export {Float32, Float, Int, Struct, Numeric, PGNumeric, SpannerDate};
diff --git a/src/table.ts b/src/table.ts
@@ -37,6 +37,8 @@ import {
   setSpanError,
   traceConfig,
 } from './instrument';
+import {google} from '../protos/protos';
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
 
 export type Key = string | string[];
 
@@ -53,6 +55,7 @@ export type DropTableCallback = UpdateSchemaCallback;
 interface MutateRowsOptions extends CommitOptions {
   requestOptions?: Omit<IRequestOptions, 'requestTag'>;
   excludeTxnFromChangeStreams?: boolean;
+  isolationLevel?: IsolationLevel;
 }
 
 export type DeleteRowsCallback = CommitCallback;
@@ -1100,10 +1103,16 @@ class Table {
           ? options.excludeTxnFromChangeStreams
           : false;
 
+      const isolationLevel =
+        'isolationLevel' in options
+          ? options.isolationLevel
+          : IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED;
+
       this.database.runTransaction(
         {
           requestOptions: requestOptions,
           excludeTxnFromChangeStreams: excludeTxnFromChangeStreams,
+          isolationLevel: isolationLevel,
         },
         (err, transaction) => {
           if (err) {
diff --git a/src/transaction-runner.ts b/src/transaction-runner.ts
@@ -26,6 +26,7 @@ import {isSessionNotFoundError} from './session-pool';
 import {Database} from './database';
 import {google} from '../protos/protos';
 import IRequestOptions = google.spanner.v1.IRequestOptions;
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const jsonProtos = require('../protos/protos.json');
@@ -46,6 +47,7 @@ export interface RunTransactionOptions {
   requestOptions?: Pick<IRequestOptions, 'transactionTag'>;
   optimisticLock?: boolean;
   excludeTxnFromChangeStreams?: boolean;
+  isolationLevel?: IsolationLevel;
 }
 
 /**
@@ -124,7 +126,10 @@ export abstract class Runner<T> {
     this.transaction = transaction;
     this.transaction.useInRunner();
 
-    const defaults = {timeout: 3600000};
+    const defaults = {
+      timeout: 3600000,
+      isolationLevel: IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED,
+    };
 
     this.options = Object.assign(defaults, options);
   }
@@ -202,12 +207,9 @@ export abstract class Runner<T> {
     const transaction = this.session.transaction(
       (this.session.parent as Database).queryOptions_
     );
-    if (this.options.optimisticLock) {
-      transaction.useOptimisticLock();
-    }
-    if (this.options.excludeTxnFromChangeStreams) {
-      transaction.excludeTxnFromChangeStreams();
-    }
+    transaction!.setReadWriteTransactionOptions(
+      this.options as RunTransactionOptions
+    );
     if (this.attempts > 0) {
       await transaction.begin();
     }
diff --git a/src/transaction.ts b/src/transaction.ts
@@ -41,6 +41,7 @@ import {
   getCommonHeaders,
 } from './common';
 import {google} from '../protos/protos';
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
 import IAny = google.protobuf.IAny;
 import IQueryOptions = google.spanner.v1.ExecuteSqlRequest.IQueryOptions;
 import IRequestOptions = google.spanner.v1.IRequestOptions;
@@ -52,6 +53,7 @@ import {
   setSpanError,
   setSpanErrorAndException,
 } from './instrument';
+import {RunTransactionOptions} from './transaction-runner';
 
 export type Rows = Array<Row | Json>;
 const RETRY_INFO_TYPE = 'type.googleapis.com/google.rpc.retryinfo';
@@ -1822,6 +1824,7 @@ export class Transaction extends Dml {
 
     this._queuedMutations = [];
     this._options = {readWrite: options};
+    this._options.isolationLevel = IsolationLevel.ISOLATION_LEVEL_UNSPECIFIED;
     this.requestOptions = requestOptions;
   }
 
@@ -2700,6 +2703,28 @@ export class Transaction extends Dml {
   excludeTxnFromChangeStreams(): void {
     this._options.excludeTxnFromChangeStreams = true;
   }
+
+  setReadWriteTransactionOptions(options: RunTransactionOptions) {
+    /**
+     * Set optimistic concurrency control for the transaction.
+     */
+    if (options?.optimisticLock) {
+      this._options.readWrite!.readLockMode = ReadLockMode.OPTIMISTIC;
+    }
+    /**
+     * Set option excludeTxnFromChangeStreams=true to exclude read/write transactions
+     * from being tracked in change streams.
+     */
+    if (options?.excludeTxnFromChangeStreams) {
+      this._options.excludeTxnFromChangeStreams = true;
+    }
+    /**
+     * Set isolation level .
+     */
+    this._options.isolationLevel = options?.isolationLevel
+      ? options?.isolationLevel
+      : this._getSpanner().defaultTransactionOptions.isolationLevel;
+  }
 }
 
 /*! Developer Documentation
diff --git a/test/database.ts b/test/database.ts
@@ -39,6 +39,8 @@ import {google} from '../protos/protos';
 import {protos} from '../src';
 import * as inst from '../src/instance';
 import RequestOptions = google.spanner.v1.RequestOptions;
+import IsolationLevel = google.spanner.v1.TransactionOptions.IsolationLevel;
+import ReadLockMode = google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
 import EncryptionType = google.spanner.admin.database.v1.RestoreDatabaseEncryptionConfig.EncryptionType;
 import {
   BatchWriteOptions,
@@ -47,6 +49,7 @@ import {
   MutationSet,
 } from '../src/transaction';
 import {SessionFactory} from '../src/session-factory';
+import {RunTransactionOptions} from '../src/transaction-runner';
 let promisified = false;
 const fakePfy = extend({}, pfy, {
   promisifyAll(klass, options) {
@@ -177,6 +180,7 @@ class FakeTransaction extends EventEmitter {
   setQueuedMutations(mutation) {
     this._queuedMutations = mutation;
   }
+  setReadWriteTransactionOptions(options: RunTransactionOptions) {}
   commit(
     options?: CommitOptions,
     callback?: CommitCallback
@@ -3166,6 +3170,17 @@ describe('Database', () => {
       assert.strictEqual(options, fakeOptions);
     });
 
+    it('should optionally accept runner `option` isolationLevel', async () => {
+      const fakeOptions = {
+        isolationLevel: IsolationLevel.REPEATABLE_READ,
+      };
+
+      await database.runTransaction(fakeOptions, assert.ifError);
+
+      const options = fakeTransactionRunner.calledWith_[3];
+      assert.strictEqual(options, fakeOptions);
+    });
+
     it('should release the session when finished', done => {
       const releaseStub = (
         sandbox.stub(pool, 'release') as sinon.SinonStub
@@ -3237,6 +3252,17 @@ describe('Database', () => {
       assert.strictEqual(options, fakeOptions);
     });
 
+    it('should optionally accept runner `option` isolationLevel', async () => {
+      const fakeOptions = {
+        isolationLevel: IsolationLevel.REPEATABLE_READ,
+      };
+
+      await database.runTransactionAsync(fakeOptions, assert.ifError);
+
+      const options = fakeAsyncTransactionRunner.calledWith_[3];
+      assert.strictEqual(options, fakeOptions);
+    });
+
     it('should return the runners resolved value', async () => {
       const fakeValue = {};
 
diff --git a/test/index.ts b/test/index.ts
@@ -38,6 +38,7 @@ import {
   GetInstancesOptions,
 } from '../src';
 import {CLOUD_RESOURCE_HEADER} from '../src/common';
+import IsolationLevel = protos.google.spanner.v1.TransactionOptions.IsolationLevel;
 const singer = require('./data/singer');
 const music = singer.examples.spanner.music;
 
@@ -327,6 +328,20 @@ describe('Spanner', () => {
       assert.strictEqual(spanner.directedReadOptions, fakeDirectedReadOptions);
     });
 
+    it('should optionally accept defaultTransactionOptions', () => {
+      const fakeDefaultTxnOptions = {
+        defaultTransactionOptions: {
+          isolationLevel: IsolationLevel.REPEATABLE_READ,
+        },
+      };
+
+      const spanner = new Spanner(fakeDefaultTxnOptions);
+      assert.strictEqual(
+        spanner.defaultTransactionOptions,
+        fakeDefaultTxnOptions.defaultTransactionOptions
+      );
+    });
+
     it('should set projectFormattedName_', () => {
       assert.strictEqual(
         spanner.projectFormattedName_,
diff --git a/test/spanner.ts b/test/spanner.ts
diff --git a/test/table.ts b/test/table.ts
diff --git a/test/transaction-runner.ts b/test/transaction-runner.ts
diff --git a/test/transaction.ts b/test/transaction.ts
