feat: new transaction feature

src/request.ts

@@ -27,11 +27,17 @@
 
 // eslint-disable-next-line @typescript-eslint/no-var-requires
 const streamEvents = require('stream-events');
+export const transactionExpiredError = 'This transaction has already expired.';
 
 export interface AbortableDuplex extends Duplex {
   abort(): void;
 }
 
+interface TransactionRequestOptions {
+  readOnly?: {};
+  readWrite?: {previousTransaction?: string | Uint8Array | null};
+}
+
 // Import the clients for each version supported by this package.
 const gapic = Object.freeze({
   v1: require('./v1'),
@@ -55,9 +61,10 @@
   RunQueryResponse,
   RunQueryCallback,
 } from './query';
-import {Datastore} from '.';
+import {Datastore, Transaction} from '.';
 import ITimestamp = google.protobuf.ITimestamp;
 import {AggregateQuery} from './aggregate';
+import {RunOptions} from './transaction';
 import * as protos from '../protos/protos';
 import {serializer} from 'google-gax';
 import * as gax from 'google-gax';
@@ -153,6 +160,16 @@
   return {};
 }
 
+const readTimeAndConsistencyError =
+  'Read time and read consistency cannot both be specified.';
+
+// Write function to check for readTime and readConsistency.
+function throwOnReadTimeAndConsistency(options: RunQueryStreamOptions) {
+  if (options.readTime && options.consistency) {
+    throw new Error(readTimeAndConsistencyError);
+  }
+}
+
 /**
  * A map of read consistency values to proto codes.
  *
@@ -164,6 +181,19 @@
   strong: 1,
 };
 
+/**
+ * By default a DatastoreRequest is in the NOT_TRANSACTION state. If the
+ * DatastoreRequest is a Transaction object, then initially it will be in
+ * the NOT_STARTED state, but then the state will become IN_PROGRESS after the
+ * transaction has started.
+ */
+export enum TransactionState {
+  NOT_TRANSACTION,
+  NOT_STARTED,
+  IN_PROGRESS,
+  EXPIRED,
+}
+
 /**
  * Handle logic for Datastore API operations. Handles request logic for
  * Datastore.
@@ -184,6 +214,7 @@
     | Array<(err: Error | null, resp: Entity | null) => void>
     | Entity;
   datastore!: Datastore;
+  protected state: TransactionState = TransactionState.NOT_TRANSACTION;
   [key: string]: Entity;
 
   /**
@@ -335,6 +366,15 @@
     );
   }
 
+  /* This throws an error if the transaction has already expired.
+   *
+   */
+  protected checkExpired() {
+    if (this.state === TransactionState.EXPIRED) {
+      throw Error(transactionExpiredError);
+    }
+  }
+
   /**
    * Retrieve the entities as a readable object stream.
    *
@@ -365,12 +405,19 @@
     keys: Entities,
     options: CreateReadStreamOptions = {}
   ): Transform {
+    this.checkExpired();
     keys = arrify(keys).map(entity.keyToKeyProto);
     if (keys.length === 0) {
       throw new Error('At least one Key object is required.');
     }
 
     const makeRequest = (keys: entity.Key[] | KeyProto[]) => {
+      try {
+        throwOnReadTimeAndConsistency(options);
+      } catch (error: any) {
+        stream.destroy(error);
+        return;
+      }
       const reqOpts = this.getRequestOptions(options);
       Object.assign(reqOpts, {keys});
       this.request_(
@@ -381,6 +428,7 @@
           gaxOpts: options.gaxOptions,
         },
         (err, resp) => {
+          this.parseTransactionResponse(resp);
           if (err) {
             stream.destroy(err);
             return;
@@ -631,6 +679,10 @@
     const callback =
       typeof optionsOrCallback === 'function' ? optionsOrCallback : cb!;
 
+    if (this.state === TransactionState.EXPIRED) {
+      callback(new Error(transactionExpiredError));
+      return;
+    }
     this.createReadStream(keys, options)
       .on('error', callback)
       .pipe(
@@ -641,6 +693,22 @@
       );
   }
 
+  /**
+   * This function saves results from a successful beginTransaction call.
+   *
+   * @param {BeginAsyncResponse} [response] The response from a call to
+   * begin a transaction that completed successfully.
+   *
+   **/
+  protected parseTransactionResponse(resp?: {
+    transaction?: Uint8Array | string | undefined | null;
+  }): void {
+    if (resp && resp.transaction && Buffer.byteLength(resp.transaction) > 0) {
+      this.id = resp!.transaction;
+      this.state = TransactionState.IN_PROGRESS;
+    }
+  }
+
   /**
    * Datastore allows you to run aggregate queries by supplying aggregate fields
    * which will determine the type of aggregation that is performed.
@@ -677,6 +745,14 @@
     const callback =
       typeof optionsOrCallback === 'function' ? optionsOrCallback : cb!;
 
+    if (this.state === TransactionState.EXPIRED) {
+      callback(new Error(transactionExpiredError));
+      return;
+    }
+    if (options.readTime && options.consistency) {
+      callback(new Error(readTimeAndConsistencyError));
+      return;
+    }
     query.query = extend(true, new Query(), query.query);
     let queryProto: QueryProto;
     try {
@@ -687,7 +763,13 @@
       setImmediate(callback, e as Error);
       return;
     }
-    const sharedQueryOpts = this.getQueryOptions(query.query, options);
+    let sharedQueryOpts;
+    try {
+      sharedQueryOpts = this.getQueryOptions(query.query, options);
+    } catch (error: any) {
+      callback(error);
+      return;
+    }
     const aggregationQueryOptions: AggregationQueryOptions = {
       nestedQuery: queryProto,
       aggregations: query.toProto(),
@@ -704,6 +786,7 @@
       },
       (err, res) => {
         const info = getInfoFromStats(res);
+        this.parseTransactionResponse(res);
         if (res && res.batch) {
           const results = res.batch.aggregationResults;
           const finalResults = results
@@ -846,6 +929,10 @@
 
     let info: RunQueryInfo;
 
+    if (this.state === TransactionState.EXPIRED) {
+      callback(new Error(transactionExpiredError));
+      return;
+    }
     this.runQueryStream(query, options)
       .on('error', callback)
       .on('info', info_ => {
@@ -892,11 +979,13 @@
    * ```
    */
   runQueryStream(query: Query, options: RunQueryStreamOptions = {}): Transform {
+    this.checkExpired();
     query = extend(true, new Query(), query);
     const makeRequest = (query: Query) => {
       let queryProto: QueryProto;
       try {
         queryProto = entity.queryToQueryProto(query);
+        throwOnReadTimeAndConsistency(options);
       } catch (e) {
         // using setImmediate here to make sure this doesn't throw a
         // synchronous error
@@ -918,7 +1007,8 @@
       );
     };
 
-    function onResultSet(err?: Error | null, resp?: Entity) {
+    const onResultSet = (err?: Error | null, resp?: Entity) => {
+      this.parseTransactionResponse(resp);
       if (err) {
         stream.destroy(err);
         return;
@@ -977,7 +1067,7 @@
 
         makeRequest(query);
       });
-    }
+    };
 
     const stream = streamEvents(new Transform({objectMode: true}));
     stream.once('reading', () => {
@@ -990,11 +1080,24 @@
     options: RunQueryStreamOptions
   ): SharedQueryOptions {
     const sharedQueryOpts = {} as SharedQueryOptions;
+    if (isTransaction(this)) {
+      if (this.state === TransactionState.NOT_STARTED) {
+        if (sharedQueryOpts.readOptions === undefined) {
+          sharedQueryOpts.readOptions = {};
+        }
+        sharedQueryOpts.readOptions.newTransaction = getTransactionRequest(
+          this,
+          {}
+        );
+        sharedQueryOpts.readOptions.consistencyType = 'newTransaction';
+      }
+    }
     if (options.consistency) {
       const code = CONSISTENCY_PROTO_CODE[options.consistency.toLowerCase()];
-      sharedQueryOpts.readOptions = {
-        readConsistency: code,
-      };
+      if (sharedQueryOpts.readOptions === undefined) {
+        sharedQueryOpts.readOptions = {};
+      }
+      sharedQueryOpts.readOptions.readConsistency = code;
     }
     if (options.readTime) {
       if (sharedQueryOpts.readOptions === undefined) {
@@ -1122,18 +1225,27 @@
       reqOpts.transaction = this.id;
     }
 
+    if (
+      isTransaction ||
+      (reqOpts.readOptions && reqOpts.readOptions.newTransaction)
+    ) {
+      if (reqOpts.readOptions && reqOpts.readOptions.readConsistency) {
+        callback(
+          new Error('Read consistency cannot be specified in a transaction.')
+        );
+        return;
+      }
+      if (reqOpts.readOptions && reqOpts.readOptions.readTime) {
+        callback(new Error('Read time cannot be specified in a transaction.'));
+        return;
+      }
+    }
     if (
       isTransaction &&
       (method === 'lookup' ||
         method === 'runQuery' ||
         method === 'runAggregationQuery')
     ) {
-      if (reqOpts.readOptions && reqOpts.readOptions.readConsistency) {
-        throw new Error(
-          'Read consistency cannot be specified in a transaction.'
-        );
-      }
-
       if (reqOpts.readOptions) {
         Object.assign(reqOpts.readOptions, {transaction: this.id});
       } else {
@@ -1229,6 +1341,36 @@
   }
 }
 
+function isTransaction(request: DatastoreRequest): request is Transaction {
+  return request instanceof Transaction;
+}
+
+export function getTransactionRequest(
+  transaction: Transaction,
+  options: RunOptions
+): TransactionRequestOptions {
+  let reqOpts: TransactionRequestOptions = {};
+  if (options.readOnly || transaction.readOnly) {
+    reqOpts.readOnly = {};
+  }
+  if (options.transactionId || transaction.id) {
+    reqOpts.readWrite = {
+      previousTransaction: options.transactionId || transaction.id,
+    };
+  }
+  if (options.transactionOptions) {
+    reqOpts = {};
+    if (options.transactionOptions.readOnly) {
+      reqOpts.readOnly = {};
+    }
+    const id = options.transactionOptions.id;
+    if (id) {
+      reqOpts.readWrite = {previousTransaction: id};
+    }
+  }
+  return reqOpts;
+}
+
 export interface ConsistencyProtoCode {
   [key: string]: number;
 }
@@ -1294,15 +1436,18 @@
     readConsistency?: number;
     transaction?: string | Uint8Array | null;
     readTime?: ITimestamp;
+    newTransaction?: TransactionRequestOptions;
+    consistencyType?:
+      | 'readConsistency'
+      | 'transaction'
+      | 'newTransaction'
+      | 'readTime';
   };
 }
 export interface RequestOptions extends SharedQueryOptions {
   mutations?: google.datastore.v1.IMutation[];
   keys?: Entity;
-  transactionOptions?: {
-    readOnly?: {};
-    readWrite?: {previousTransaction?: string | Uint8Array | null};
-  } | null;
+  transactionOptions?: TransactionRequestOptions | null;
   transaction?: string | null | Uint8Array;
   mode?: string;
   query?: QueryProto;
@@ -1333,7 +1478,7 @@
  * that a callback is omitted.
  */
 promisifyAll(DatastoreRequest, {
-  exclude: ['getQueryOptions', 'getRequestOptions'],
+  exclude: ['checkExpired', 'getQueryOptions', 'getRequestOptions'],
 });
 
 /**