feat: add tagging support (#1419)

.gitignore

@@ -12,3 +12,4 @@ system-test/*key.json
 .DS_Store
 package-lock.json
 __pycache__
+.idea

src/database.ts

@@ -2662,6 +2662,10 @@ class Database extends common.GrpcServiceObject {
     while (true) {
       try {
         const [session, transaction] = await promisify(getWriteSession)();
+        transaction.requestOptions = Object.assign(
+          transaction.requestOptions || {},
+          options.requestOptions
+        );
         const runner = new AsyncTransactionRunner<T>(
           session,
           transaction,

src/session.ts

@@ -38,6 +38,7 @@ import {
 import {ServiceObjectConfig} from '@google-cloud/common';
 import {NormalCallback, CLOUD_RESOURCE_HEADER} from './common';
 import {grpc, CallOptions} from 'google-gax';
+import IRequestOptions = google.spanner.v1.IRequestOptions;
 
 export type GetSessionResponse = [Session, r.Response];
 
@@ -450,9 +451,10 @@ export class Session extends common.GrpcServiceObject {
    * const transaction = session.transaction();
    */
   transaction(
-    queryOptions?: google.spanner.v1.ExecuteSqlRequest.IQueryOptions
+    queryOptions?: google.spanner.v1.ExecuteSqlRequest.IQueryOptions,
+    requestOptions?: Pick<IRequestOptions, 'transactionTag'>
   ) {
-    return new Transaction(this, undefined, queryOptions);
+    return new Transaction(this, undefined, queryOptions, requestOptions);
   }
   /**
    * Format the session name to include the parent database's name.

src/table.ts

@@ -30,6 +30,7 @@ import {
 } from './transaction';
 import {google as databaseAdmin} from '../protos/protos';
 import {Schema, LongRunningCallback} from './common';
+import IRequestOptions = databaseAdmin.spanner.v1.IRequestOptions;
 
 export type Key = string | string[];
 
@@ -43,25 +44,30 @@ export type CreateTableCallback = LongRunningCallback<Table>;
 export type DropTableResponse = UpdateSchemaResponse;
 export type DropTableCallback = UpdateSchemaCallback;
 
+interface MutateRowsOptions extends CommitOptions {
+  requestOptions?: Omit<IRequestOptions, 'requestTag'>;
+}
+
 export type DeleteRowsCallback = CommitCallback;
 export type DeleteRowsResponse = CommitResponse;
-export type DeleteRowsOptions = CommitOptions;
+export type DeleteRowsOptions = MutateRowsOptions;
 
 export type InsertRowsCallback = CommitCallback;
 export type InsertRowsResponse = CommitResponse;
-export type InsertRowsOptions = CommitOptions;
+export type InsertRowsOptions = MutateRowsOptions;
 
 export type ReplaceRowsCallback = CommitCallback;
 export type ReplaceRowsResponse = CommitResponse;
-export type ReplaceRowsOptions = CommitOptions;
+export type ReplaceRowsOptions = MutateRowsOptions;
 
 export type UpdateRowsCallback = CommitCallback;
 export type UpdateRowsResponse = CommitResponse;
-export type UpdateRowsOptions = CommitOptions;
+export type UpdateRowsOptions = MutateRowsOptions;
 
 export type UpsertRowsCallback = CommitCallback;
 export type UpsertRowsResponse = CommitResponse;
-export type UpsertRowsOptions = CommitOptions;
+export type UpsertRowsOptions = MutateRowsOptions;
+
 /**
  * Create a Table object to interact with a table in a Cloud Spanner
  * database.
@@ -1012,10 +1018,12 @@ class Table {
   private _mutate(
     method: 'deleteRows' | 'insert' | 'replace' | 'update' | 'upsert',
     rows: object | object[],
-    options: CommitOptions | CallOptions,
+    options: MutateRowsOptions | CallOptions = {},
     callback: CommitCallback
   ): void {
-    this.database.runTransaction((err, transaction) => {
+    const requestOptions =
+      'requestOptions' in options ? options.requestOptions : {};
+    this.database.runTransaction({requestOptions}, (err, transaction) => {
       if (err) {
         callback(err);
         return;

src/transaction-runner.ts

@@ -24,6 +24,8 @@ import {Transaction} from './transaction';
 import {NormalCallback} from './common';
 import {isSessionNotFoundError} from './session-pool';
 import {Database} from './database';
+import {google} from '../protos/protos';
+import IRequestOptions = google.spanner.v1.IRequestOptions;
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const jsonProtos = require('../protos/protos.json');
@@ -41,6 +43,7 @@ const RetryInfo = Root.fromJSON(jsonProtos).lookup('google.rpc.RetryInfo');
  */
 export interface RunTransactionOptions {
   timeout?: number;
+  requestOptions?: Pick<IRequestOptions, 'transactionTag'>;
 }
 
 /**

src/transaction.ts

@@ -125,6 +125,10 @@ export type RunResponse = [
 ];
 export type RunUpdateResponse = [number];
 
+export interface BatchUpdateOptions {
+  requestOptions?: Omit<IRequestOptions, 'transactionTag'>;
+  gaxOptions?: CallOptions;
+}
 export interface BatchUpdateCallback {
   (
     err: null | BatchUpdateError,
@@ -220,6 +224,7 @@ export class Snapshot extends EventEmitter {
   session: Session;
   queryOptions?: IQueryOptions;
   resourceHeader_: {[k: string]: string};
+  requestOptions?: Pick<IRequestOptions, 'transactionTag'>;
 
   /**
    * The transaction ID.
@@ -344,6 +349,14 @@ export class Snapshot extends EventEmitter {
       options,
     };
 
+    // Only hand crafted read-write transactions will be able to set a
+    // transaction tag for the BeginTransaction RPC. Also, this.requestOptions
+    // is only set in the constructor of Transaction, which is the constructor
+    // for read/write transactions.
+    if (this.requestOptions) {
+      reqOpts.requestOptions = this.requestOptions;
+    }
+
     this.request(
       {
         client: 'SpannerClient',
@@ -538,7 +551,8 @@ export class Snapshot extends EventEmitter {
     table: string,
     request = {} as ReadRequest
   ): PartialResultStream {
-    const {gaxOptions, json, jsonOptions, maxResumeRetries} = request;
+    const {gaxOptions, json, jsonOptions, maxResumeRetries, requestOptions} =
+      request;
     const keySet = Snapshot.encodeKeySet(request);
     const transaction: spannerClient.spanner.v1.ITransactionSelector = {};
 
@@ -556,13 +570,22 @@ export class Snapshot extends EventEmitter {
     delete request.maxResumeRetries;
     delete request.keys;
     delete request.ranges;
+    delete request.requestOptions;
 
-    const reqOpts: ReadRequest = Object.assign(request, {
-      session: this.session.formattedName_!,
-      transaction,
-      table,
-      keySet,
-    });
+    const reqOpts: spannerClient.spanner.v1.IReadRequest = Object.assign(
+      request,
+      {
+        session: this.session.formattedName_!,
+        requestOptions: this.configureTagOptions(
+          typeof transaction.singleUse !== 'undefined',
+          this.requestOptions?.transactionTag!,
+          requestOptions
+        ),
+        transaction,
+        table,
+        keySet,
+      }
+    );
 
     const makeRequest = (resumeToken?: ResumeToken): Readable => {
       return this.requestStream({
@@ -965,7 +988,8 @@ export class Snapshot extends EventEmitter {
       query.queryOptions
     );
 
-    const {gaxOptions, json, jsonOptions, maxResumeRetries} = query;
+    const {gaxOptions, json, jsonOptions, maxResumeRetries, requestOptions} =
+      query;
     let reqOpts;
 
     const sanitizeRequest = () => {
@@ -981,10 +1005,17 @@ export class Snapshot extends EventEmitter {
       delete query.json;
       delete query.jsonOptions;
       delete query.maxResumeRetries;
+      delete query.requestOptions;
       delete query.types;
+
       reqOpts = Object.assign(query, {
         session: this.session.formattedName_!,
         seqno: this._seqno++,
+        requestOptions: this.configureTagOptions(
+          typeof transaction.singleUse !== 'undefined',
+          this.requestOptions?.transactionTag!,
+          requestOptions
+        ),
         transaction,
         params,
         paramTypes,
@@ -1018,6 +1049,22 @@ export class Snapshot extends EventEmitter {
     });
   }
 
+  /**
+   *
+   * @private
+   */
+  configureTagOptions(
+    singleUse?: boolean,
+    transactionTag?: string,
+    requestOptions = {}
+  ): IRequestOptions | null {
+    if (!singleUse && transactionTag) {
+      (requestOptions as IRequestOptions).transactionTag = transactionTag;
+    }
+
+    return requestOptions!;
+  }
+
   /**
    * Transforms convenience options `keys` and `ranges` into a KeySet object.
    *
@@ -1151,7 +1198,7 @@ export class Snapshot extends EventEmitter {
  * that a callback is omitted.
  */
 promisifyAll(Snapshot, {
-  exclude: ['end'],
+  exclude: ['configureTagOptions', 'end'],
 });
 
 /**
@@ -1312,12 +1359,14 @@ export class Transaction extends Dml {
   constructor(
     session: Session,
     options = {} as spannerClient.spanner.v1.TransactionOptions.ReadWrite,
-    queryOptions?: IQueryOptions
+    queryOptions?: IQueryOptions,
+    requestOptions?: Pick<IRequestOptions, 'transactionTag'>
   ) {
     super(session, undefined, queryOptions);
 
     this._queuedMutations = [];
     this._options = {readWrite: options};
+    this.requestOptions = requestOptions;
   }
 
   batchUpdate(
@@ -1435,9 +1484,13 @@ export class Transaction extends Dml {
 
     const reqOpts: spannerClient.spanner.v1.ExecuteBatchDmlRequest = {
       session: this.session.formattedName_!,
+      requestOptions: this.configureTagOptions(
+        false,
+        this.requestOptions?.transactionTag!,
+        (options as BatchUpdateOptions).requestOptions
+      ),
       transaction: {id: this.id!},
       seqno: this._seqno++,
-      requestOptions: (options as BatchUpdateOptions).requestOptions,
       statements,
     } as spannerClient.spanner.v1.ExecuteBatchDmlRequest;
 
@@ -1595,6 +1648,10 @@ export class Transaction extends Dml {
     ) {
       reqOpts.returnCommitStats = (options as CommitOptions).returnCommitStats;
     }
+    reqOpts.requestOptions = Object.assign(
+      requestOptions || {},
+      this.requestOptions
+    );
 
     this.request(
       {

test/mockserver/mockspanner.ts

@@ -870,24 +870,33 @@ export class MockSpanner {
         }
         const session = this.sessions.get(call.request!.session);
         if (session) {
-          const buffer = Buffer.from(call.request!.transactionId as string);
-          const transactionId = buffer.toString();
-          const fullTransactionId =
-            session.name + '/transactions/' + transactionId;
-          const transaction = this.transactions.get(fullTransactionId);
-          if (transaction) {
-            this.transactions.delete(fullTransactionId);
-            this.transactionOptions.delete(fullTransactionId);
+          if (call.request!.transactionId) {
+            const buffer = Buffer.from(call.request!.transactionId as string);
+            const transactionId = buffer.toString();
+            const fullTransactionId =
+              session.name + '/transactions/' + transactionId;
+            const transaction = this.transactions.get(fullTransactionId);
+            if (transaction) {
+              this.transactions.delete(fullTransactionId);
+              this.transactionOptions.delete(fullTransactionId);
+              callback(
+                null,
+                protobuf.CommitResponse.create({
+                  commitTimestamp: now(),
+                })
+              );
+            } else {
+              callback(
+                MockSpanner.createTransactionNotFoundError(fullTransactionId)
+              );
+            }
+          } else if (call.request!.singleUseTransaction) {
             callback(
               null,
               protobuf.CommitResponse.create({
                 commitTimestamp: now(),
               })
             );
-          } else {
-            callback(
-              MockSpanner.createTransactionNotFoundError(fullTransactionId)
-            );
           }
         } else {
           callback(