observability: trace BatchTransaction and Table

This change is part of a series of changes to add
OpenTelemetry traces, focused on BatchTransaction and Table.

While here, made the tests for sessionPool spans much more
precise to avoid flakes.

Updates googleapis#2079
Built from PR googleapis#2087
Updates googleapis#2114

src/batch-transaction.ts

@@ -25,6 +25,7 @@ import {
   CLOUD_RESOURCE_HEADER,
   addLeaderAwareRoutingHeader,
 } from '../src/common';
+import {startTrace, setSpanError, traceConfig} from './instrument';
 
 export interface TransactionIdentifier {
   session: string | Session;
@@ -136,20 +137,37 @@ class BatchTransaction extends Snapshot {
     delete reqOpts.gaxOptions;
     delete reqOpts.types;
 
-    const headers: {[k: string]: string} = {};
-    if (this._getSpanner().routeToLeaderEnabled) {
-      addLeaderAwareRoutingHeader(headers);
-    }
+    const traceConfig: traceConfig = {
+      sql: query,
+      opts: this.observabilityOptions,
+    };
+    return startTrace(
+      'BatchTransaction.createQueryPartitions',
+      traceConfig,
+      span => {
+        const headers: {[k: string]: string} = {};
+        if (this._getSpanner().routeToLeaderEnabled) {
+          addLeaderAwareRoutingHeader(headers);
+        }
 
-    this.createPartitions_(
-      {
-        client: 'SpannerClient',
-        method: 'partitionQuery',
-        reqOpts,
-        gaxOpts: query.gaxOptions,
-        headers: headers,
-      },
-      callback
+        this.createPartitions_(
+          {
+            client: 'SpannerClient',
+            method: 'partitionQuery',
+            reqOpts,
+            gaxOpts: query.gaxOptions,
+            headers: headers,
+          },
+          (err, partitions, resp) => {
+            if (err) {
+              setSpanError(span, err);
+            }
+
+            span.end();
+            callback(err, partitions, resp);
+          }
+        );
+      }
     );
   }
   /**
@@ -163,38 +181,52 @@ class BatchTransaction extends Snapshot {
    * @param {function} callback Callback function.
    */
   createPartitions_(config, callback) {
-    const query = extend({}, config.reqOpts, {
-      session: this.session.formattedName_,
-      transaction: {id: this.id},
-    });
-    config.reqOpts = extend({}, query);
-    config.headers = {
-      [CLOUD_RESOURCE_HEADER]: (this.session.parent as Database).formattedName_,
+    const traceConfig: traceConfig = {
+      opts: this.observabilityOptions,
     };
-    delete query.partitionOptions;
-    this.session.request(config, (err, resp) => {
-      if (err) {
-        callback(err, null, resp);
-        return;
-      }
 
-      const partitions = resp.partitions.map(partition => {
-        return extend({}, query, partition);
-      });
+    return startTrace(
+      'BatchTransaction.createPartitions_',
+      traceConfig,
+      span => {
+        const query = extend({}, config.reqOpts, {
+          session: this.session.formattedName_,
+          transaction: {id: this.id},
+        });
+        config.reqOpts = extend({}, query);
+        config.headers = {
+          [CLOUD_RESOURCE_HEADER]: (this.session.parent as Database)
+            .formattedName_,
+        };
+        delete query.partitionOptions;
+        this.session.request(config, (err, resp) => {
+          if (err) {
+            setSpanError(span, err);
+            span.end();
+            callback(err, null, resp);
+            return;
+          }
 
-      if (resp.transaction) {
-        const {id, readTimestamp} = resp.transaction;
+          const partitions = resp.partitions.map(partition => {
+            return extend({}, query, partition);
+          });
 
-        this.id = id;
+          if (resp.transaction) {
+            const {id, readTimestamp} = resp.transaction;
 
-        if (readTimestamp) {
-          this.readTimestampProto = readTimestamp;
-          this.readTimestamp = new PreciseDate(readTimestamp);
-        }
-      }
+            this.id = id;
+
+            if (readTimestamp) {
+              this.readTimestampProto = readTimestamp;
+              this.readTimestamp = new PreciseDate(readTimestamp);
+            }
+          }
 
-      callback(null, partitions, resp);
-    });
+          span.end();
+          callback(null, partitions, resp);
+        });
+      }
+    );
   }
   /**
    * @typedef {object} ReadPartition
@@ -226,28 +258,45 @@ class BatchTransaction extends Snapshot {
    * @returns {Promise<CreateReadPartitionsResponse>}
    */
   createReadPartitions(options, callback) {
-    const reqOpts = Object.assign({}, options, {
-      keySet: Snapshot.encodeKeySet(options),
-    });
+    const traceConfig: traceConfig = {
+      opts: this.observabilityOptions,
+    };
 
-    delete reqOpts.gaxOptions;
-    delete reqOpts.keys;
-    delete reqOpts.ranges;
+    return startTrace(
+      'BatchTransaction.createReadPartitions',
+      traceConfig,
+      span => {
+        const reqOpts = Object.assign({}, options, {
+          keySet: Snapshot.encodeKeySet(options),
+        });
 
-    const headers: {[k: string]: string} = {};
-    if (this._getSpanner().routeToLeaderEnabled) {
-      addLeaderAwareRoutingHeader(headers);
-    }
+        delete reqOpts.gaxOptions;
+        delete reqOpts.keys;
+        delete reqOpts.ranges;
 
-    this.createPartitions_(
-      {
-        client: 'SpannerClient',
-        method: 'partitionRead',
-        reqOpts,
-        gaxOpts: options.gaxOptions,
-        headers: headers,
-      },
-      callback
+        const headers: {[k: string]: string} = {};
+        if (this._getSpanner().routeToLeaderEnabled) {
+          addLeaderAwareRoutingHeader(headers);
+        }
+
+        this.createPartitions_(
+          {
+            client: 'SpannerClient',
+            method: 'partitionRead',
+            reqOpts,
+            gaxOpts: options.gaxOptions,
+            headers: headers,
+          },
+          (err, partitions, resp) => {
+            if (err) {
+              setSpanError(span, err);
+            }
+
+            span.end();
+            callback(err, partitions, resp);
+          }
+        );
+      }
     );
   }
   /**
@@ -322,6 +371,7 @@ class BatchTransaction extends Snapshot {
    * ```
    */
   executeStream(partition) {
+    // TODO: Instrument the streams with Otel.
     if (is.string(partition.table)) {
       return this.createReadStream(partition.table, partition);
     }

src/instrument.ts

@@ -47,7 +47,7 @@ interface SQLStatement {
 
 interface observabilityOptions {
   tracerProvider: TracerProvider;
-  enableExtendedTracing: boolean;
+  enableExtendedTracing?: boolean;
 }
 
 export type {observabilityOptions as ObservabilityOptions};
@@ -81,7 +81,7 @@ interface traceConfig {
 }
 
 const SPAN_NAMESPACE_PREFIX = 'CloudSpanner'; // TODO: discuss & standardize this prefix.
-export {SPAN_NAMESPACE_PREFIX};
+export {SPAN_NAMESPACE_PREFIX, traceConfig};
 
 /**
  * startTrace begins an active span in the current active context

src/table.ts

@@ -31,6 +31,12 @@ import {
 import {google as databaseAdmin} from '../protos/protos';
 import {Schema, LongRunningCallback} from './common';
 import IRequestOptions = databaseAdmin.spanner.v1.IRequestOptions;
+import {
+  ObservabilityOptions,
+  startTrace,
+  setSpanError,
+  traceConfig,
+} from './instrument';
 
 export type Key = string | string[];
 
@@ -93,6 +99,7 @@ const POSTGRESQL = 'POSTGRESQL';
 class Table {
   database: Database;
   name: string;
+  observabilityOptions?: ObservabilityOptions;
   constructor(database: Database, name: string) {
     /**
      * The {@link Database} instance of this {@link Table} instance.
@@ -1072,29 +1079,44 @@ class Table {
     options: MutateRowsOptions | CallOptions = {},
     callback: CommitCallback
   ): void {
-    const requestOptions =
-      'requestOptions' in options ? options.requestOptions : {};
+    const traceConfig: traceConfig = {
+      opts: this.observabilityOptions,
+    };
 
-    const excludeTxnFromChangeStreams =
-      'excludeTxnFromChangeStreams' in options
-        ? options.excludeTxnFromChangeStreams
-        : false;
+    startTrace('Table.' + method, traceConfig, span => {
+      const requestOptions =
+        'requestOptions' in options ? options.requestOptions : {};
 
-    this.database.runTransaction(
-      {
-        requestOptions: requestOptions,
-        excludeTxnFromChangeStreams: excludeTxnFromChangeStreams,
-      },
-      (err, transaction) => {
-        if (err) {
-          callback(err);
-          return;
-        }
+      const excludeTxnFromChangeStreams =
+        'excludeTxnFromChangeStreams' in options
+          ? options.excludeTxnFromChangeStreams
+          : false;
 
-        transaction![method](this.name, rows as Key[]);
-        transaction!.commit(options, callback);
-      }
-    );
+      this.database.runTransaction(
+        {
+          requestOptions: requestOptions,
+          excludeTxnFromChangeStreams: excludeTxnFromChangeStreams,
+        },
+        (err, transaction) => {
+          if (err) {
+            setSpanError(span, err);
+            span.end();
+            callback(err);
+            return;
+          }
+
+          span.end();
+          transaction![method](this.name, rows as Key[]);
+          transaction!.commit(options, (err, resp) => {
+            if (err) {
+              setSpanError(span, err);
+            }
+            span.end();
+            callback(err, resp);
+          });
+        }
+      );
+    });
   }
 }
 

src/transaction.ts

@@ -46,6 +46,7 @@ import IQueryOptions = google.spanner.v1.ExecuteSqlRequest.IQueryOptions;
 import IRequestOptions = google.spanner.v1.IRequestOptions;
 import {Database, Spanner} from '.';
 import ReadLockMode = google.spanner.v1.TransactionOptions.ReadWrite.ReadLockMode;
+import {ObservabilityOptions} from './instrument';
 
 export type Rows = Array<Row | Json>;
 const RETRY_INFO_TYPE = 'type.googleapis.com/google.rpc.retryinfo';
@@ -285,6 +286,7 @@ export class Snapshot extends EventEmitter {
   queryOptions?: IQueryOptions;
   resourceHeader_: {[k: string]: string};
   requestOptions?: Pick<IRequestOptions, 'transactionTag'>;
+  observabilityOptions?: ObservabilityOptions;
 
   /**
    * The transaction ID.