Initialize baseline _traceConfig then extend it as needed

googleapis · Oct 10, 2024 · 0c902ad · 0c902ad
1 parent 8f5ea2b
commit 0c902ad
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 66 deletions.
diff --git a/observability-test/database.ts b/observability-test/database.ts
@@ -243,13 +243,13 @@ describe('Database', () => {
  beforeEach(() => {
  fakeCodec.encode = util.noop;
  extend(Database, DatabaseCached);
- database = new Database(INSTANCE, NAME, POOL_OPTIONS);
- database.parent = INSTANCE;
- database.databaseRole = 'parent_role';
- database._observabilityOptions = {
+ INSTANCE._observabilityOptions = {
  tracerProvider: provider,
  enableExtendedTracing: false,
  };
+ database = new Database(INSTANCE, NAME, POOL_OPTIONS);
+ database.parent = INSTANCE;
+ database.databaseRole = 'parent_role';
  const gaxOpts = {};
  const options: {
  a: string;

diff --git a/src/database.ts b/src/database.ts
@@ -109,6 +109,7 @@ import {
  startTrace,
  setSpanError,
  setSpanErrorAndException,
+ traceConfig,
 } from './instrument';
 
 export type GetDatabaseRolesCallback = RequestCallback<
@@ -344,7 +345,8 @@ class Database extends common.GrpcServiceObject {
  databaseDialect?: EnumKey<
  typeof databaseAdmin.spanner.admin.database.v1.DatabaseDialect
  > | null;
- _observabilityOptions?: ObservabilityOptions;
+ _observabilityOptions?: ObservabilityOptions; // TODO: exmaine if we can remove it
+ private _traceConfig: traceConfig;
  constructor(
  instance: Instance,
  name: string,
@@ -455,6 +457,12 @@ class Database extends common.GrpcServiceObject {
  }
  this.formattedName_ = formattedName_;
  this.instance = instance;
+ this._observabilityOptions = instance._observabilityOptions;
+ this._traceConfig = {
+ opts: this._observabilityOptions,
+ dbName: this.formattedName_,
+ };
+
  this.resourceHeader_ = {
  [CLOUD_RESOURCE_HEADER]: this.formattedName_,
  };
@@ -467,7 +475,6 @@ class Database extends common.GrpcServiceObject {
  Object.assign({}, queryOptions),
  Database.getEnvironmentQueryOptions()
  );
- this._observabilityOptions = instance._observabilityOptions;
  }
  /**
  * @typedef {array} SetDatabaseMetadataResponse
@@ -693,7 +700,7 @@ class Database extends common.GrpcServiceObject {
 
  const sessions = (resp!.session || []).map(metadata => {
  const session = this.session(metadata.name!);
- session._observabilityOptions = this._observabilityOptions;
+ session._observabilityOptions = this._traceConfig!.opts;
  session.metadata = metadata;
  return session;
  });
@@ -738,7 +745,7 @@ class Database extends common.GrpcServiceObject {
  const id = identifier.transaction;
  const transaction = new BatchTransaction(session, options);
  transaction.id = id;
- transaction._observabilityOptions = this._observabilityOptions;
+ transaction._observabilityOptions = this._traceConfig!.opts;
  transaction.readTimestamp = identifier.timestamp as PreciseDate;
  return transaction;
  }
@@ -827,39 +834,41 @@ class Database extends common.GrpcServiceObject {
  typeof optionsOrCallback === 'object'
  ? (optionsOrCallback as TimestampBounds)
  : {};
-
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- return startTrace('Database.createBatchTransaction', traceConfig, span => {
- this.pool_.getSession((err, session) => {
- if (err) {
- setSpanError(span, err);
- span.end();
- callback!(err as ServiceError, null, undefined);
- return;
- }
- const transaction = this.batchTransaction({session: session!}, options);
- this._releaseOnEnd(session!, transaction, span);
- transaction.begin((err, resp) => {
+ return startTrace(
+ 'Database.createBatchTransaction',
+ this._traceConfig,
+ span => {
+ this.pool_.getSession((err, session) => {
  if (err) {
  setSpanError(span, err);
- if (isSessionNotFoundError(err)) {
- span.addEvent('No session available', {
- 'session.id': session?.id,
- });
- }
  span.end();
- callback!(err, null, resp!);
+ callback!(err as ServiceError, null, undefined);
  return;
  }
- span.addEvent('Using Session', {'session.id': session?.id});
- span.end();
- callback!(null, transaction, resp!);
+ const transaction = this.batchTransaction(
+ {session: session!},
+ options
+ );
+ this._releaseOnEnd(session!, transaction, span);
+ transaction.begin((err, resp) => {
+ if (err) {
+ setSpanError(span, err);
+ if (isSessionNotFoundError(err)) {
+ span.addEvent('No session available', {
+ 'session.id': session?.id,
+ });
+ }
+ span.end();
+ callback!(err, null, resp!);
+ return;
+ }
+ span.addEvent('Using Session', {'session.id': session?.id});
+ span.end();
+ callback!(null, transaction, resp!);
+ });
  });
- });
- });
+ }
+ );
  }
  /**
  * Create a new session.
@@ -1089,7 +1098,7 @@ class Database extends common.GrpcServiceObject {
  /CREATE TABLE `*([^\s`(]+)/
  )![1];
  const table = this.table(tableName!);
- table._observabilityOptions = this._observabilityOptions;
+ table._observabilityOptions = this._traceConfig!.opts;
  callback!(null, table, operation!, resp!);
  });
  }
@@ -1878,11 +1887,7 @@ class Database extends common.GrpcServiceObject {
  delete (gaxOpts as GetSessionsOptions).pageToken;
  }
 
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- return startTrace('Database.getSessions', traceConfig, span => {
+ return startTrace('Database.getSessions', this._traceConfig, span => {
  this.request<
  google.spanner.v1.ISession,
  google.spanner.v1.IListSessionsResponse
@@ -1903,7 +1908,7 @@ class Database extends common.GrpcServiceObject {
  sessionInstances = sessions.map(metadata => {
  const session = self.session(metadata.name!);
  session.metadata = metadata;
- session._observabilityOptions = this._observabilityOptions;
+ session._observabilityOptions = this._traceConfig!.opts;
  return session;
  });
  }
@@ -2064,11 +2069,7 @@ class Database extends common.GrpcServiceObject {
  ? (optionsOrCallback as TimestampBounds)
  : {};
 
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- return startTrace('Database.getSnapshot', traceConfig, span => {
+ return startTrace('Database.getSnapshot', this._traceConfig, span => {
  this.pool_.getSession((err, session) => {
  if (err) {
  setSpanError(span, err);
@@ -2168,11 +2169,7 @@ class Database extends common.GrpcServiceObject {
  ? (optionsOrCallback as GetTransactionOptions)
  : {};
 
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- return startTrace('Database.getTransaction', traceConfig, span => {
+ return startTrace('Database.getTransaction', this._traceConfig, span => {
  this.pool_.getSession((err, session, transaction) => {
  if (options.requestOptions) {
  transaction!.requestOptions = Object.assign(
@@ -2800,8 +2797,7 @@ class Database extends common.GrpcServiceObject {
 
  const traceConfig = {
  sql: query,
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
+ ...this._traceConfig,
  };
  return startTrace('Database.run', traceConfig, span => {
  this.runStream(query, options)
@@ -3025,8 +3021,7 @@ class Database extends common.GrpcServiceObject {
  const proxyStream: Transform = through.obj();
  const traceConfig = {
  sql: query,
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
+ ...this._traceConfig,
  };
  return startTrace('Database.runStream', traceConfig, span => {
  this.pool_.getSession((err, session) => {
@@ -3205,11 +3200,7 @@ class Database extends common.GrpcServiceObject {
  ? (optionsOrRunFn as RunTransactionOptions)
  : {};
 
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- startTrace('Database.runTransaction', traceConfig, span => {
+ startTrace('Database.runTransaction', this._traceConfig, span => {
  this.pool_.getSession((err, session?, transaction?) => {
  if (err) {
  setSpanError(span, err);
@@ -3601,11 +3592,7 @@ class Database extends common.GrpcServiceObject {
  ? (optionsOrCallback as CallOptions)
  : {};
 
- const traceConfig = {
- opts: this._observabilityOptions,
- dbName: this.formattedName_,
- };
- return startTrace('Database.writeAtLeastOnce', traceConfig, span => {
+ return startTrace('Database.writeAtLeastOnce', this._traceConfig, span => {
  this.pool_.getSession((err, session?, transaction?) => {
  if (err && isSessionNotFoundError(err as grpc.ServiceError)) {
  span.addEvent('No session available', {