diff --git a/observability-test/database.ts b/observability-test/database.ts index 5af9c18e9..d955817a7 100644 --- 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 index a7f877acf..0c5a0ff21 100644 --- 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', {