From 0a205fceb92b3583314840868177b95cf813e4f8 Mon Sep 17 00:00:00 2001 From: Emmanuel T Odeke Date: Fri, 20 Sep 2024 00:15:03 -1000 Subject: [PATCH] feat(observability): trace Database methods Adds trace spans for Database methods, as well as tests for methods: * getSession * getSnapshot * run * runStream * runTransaction tracing of other methods shall come in follow-up PRs. Updates #2079 Built from PR #2087 Updates #2114 --- observability-test/database.ts | 365 ++++++++++++++++++++++ src/database.ts | 544 +++++++++++++++++++-------------- test/database.ts | 6 +- test/spanner.ts | 6 +- 4 files changed, 691 insertions(+), 230 deletions(-) create mode 100644 observability-test/database.ts diff --git a/observability-test/database.ts b/observability-test/database.ts new file mode 100644 index 000000000..05325c008 --- /dev/null +++ b/observability-test/database.ts @@ -0,0 +1,365 @@ +/*! + * Copyright 2024 Google LLC. All Rights Reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import {grpc} from 'google-gax'; +import {google} from '../protos/protos'; +import {Database, Instance, SessionPool, Snapshot, Spanner} from '../src'; +import {MutationSet} from '../src/transaction'; +import protobuf = google.spanner.v1; +import * as mock from '../test/mockserver/mockspanner'; +import * as mockInstanceAdmin from '../test/mockserver/mockinstanceadmin'; +import * as mockDatabaseAdmin from '../test/mockserver/mockdatabaseadmin'; +const { + AlwaysOnSampler, + NodeTracerProvider, + InMemorySpanExporter, +} = require('@opentelemetry/sdk-trace-node'); +// eslint-disable-next-line n/no-extraneous-require +const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base'); + +/** A simple result set for SELECT 1. */ +function createSelect1ResultSet(): protobuf.ResultSet { + const fields = [ + protobuf.StructType.Field.create({ + name: 'NUM', + type: protobuf.Type.create({code: protobuf.TypeCode.INT64}), + }), + ]; + const metadata = new protobuf.ResultSetMetadata({ + rowType: new protobuf.StructType({ + fields, + }), + }); + return protobuf.ResultSet.create({ + metadata, + rows: [{values: [{stringValue: '1'}]}], + }); +} + +interface setupResults { + server: grpc.Server; + spanner: Spanner; + spannerMock: mock.MockSpanner; +} + +async function setup(): Promise { + const server = new grpc.Server(); + + const spannerMock = mock.createMockSpanner(server); + mockInstanceAdmin.createMockInstanceAdmin(server); + mockDatabaseAdmin.createMockDatabaseAdmin(server); + + const port: number = await new Promise((resolve, reject) => { + server.bindAsync( + '0.0.0.0:0', + grpc.ServerCredentials.createInsecure(), + (err, assignedPort) => { + if (err) { + reject(err); + } else { + resolve(assignedPort); + } + } + ); + }); + + const selectSql = 'SELECT 1'; + const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2'; + spannerMock.putStatementResult( + selectSql, + mock.StatementResult.resultSet(createSelect1ResultSet()) + ); + spannerMock.putStatementResult( + updateSql, + mock.StatementResult.updateCount(1) + ); + + const spanner = new Spanner({ + projectId: 'observability-project-id', + servicePath: 'localhost', + port, + sslCreds: grpc.credentials.createInsecure(), + }); + + return Promise.resolve({ + spanner: spanner, + server: server, + spannerMock: spannerMock, + }); +} + +describe('Database', () => { + let server: grpc.Server; + let spanner: Spanner; + let database: Database; + let spannerMock: mock.MockSpanner; + let traceExporter: typeof InMemorySpanExporter; + + after(() => { + spanner.close(); + server.tryShutdown(() => {}); + }); + + before(async () => { + const setupResult = await setup(); + spanner = setupResult.spanner; + server = setupResult.server; + spannerMock = setupResult.spannerMock; + + const selectSql = 'SELECT 1'; + const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2'; + spannerMock.putStatementResult( + selectSql, + mock.StatementResult.resultSet(createSelect1ResultSet()) + ); + spannerMock.putStatementResult( + updateSql, + mock.StatementResult.updateCount(1) + ); + + traceExporter = new InMemorySpanExporter(); + const sampler = new AlwaysOnSampler(); + + const provider = new NodeTracerProvider({ + sampler: sampler, + exporter: traceExporter, + }); + provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter)); + + const instance = spanner.instance('instance'); + database = instance.database('database'); + database.observabilityConfig = { + tracerProvider: provider, + enableExtendedTracing: false, + }; + }); + + beforeEach(() => { + spannerMock.resetRequests(); + }); + + afterEach(() => { + traceExporter.reset(); + }); + + it('getSessions', async () => { + const [rows] = await database.getSessions(); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.getSessions']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + }); + it('getSnapshot', done => { + database.getSnapshot((err, transaction) => { + assert.ifError(err); + + transaction!.run('SELECT 1', (err, rows) => { + assert.ifError(err); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.getSnapshot']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + done(); + }); + }); + }); + + it('getTransaction', done => { + database.getTransaction((err, transaction) => { + assert.ifError(err); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.getTransaction']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + done(); + }); + }); + + it('runStream', done => { + database + .runStream('SELECT 1') + .on('data', row => {}) + .on('error', assert.ifError) + .on('end', () => { + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.runStream']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + done(); + }); + }); + + it('run', async () => { + const [rows] = await database.run('SELECT 1'); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected'); + + // Sort the spans by duration. + spans.sort((spanA, spanB) => { + spanA.duration < spanB.duration; + }); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = [ + 'CloudSpanner.Database.runStream', + 'CloudSpanner.Database.run', + ]; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + // Ensure that RunStream is a child span of createQueryPartitions. + const spanRunStream = spans[0]; + const spanRun = spans[1]; + assert.ok( + spanRun.spanContext().traceId, + 'Expected that createQueryPartitions has a defined traceId' + ); + assert.ok( + spanRunStream.spanContext().traceId, + 'Expected that RunStream has a defined traceId' + ); + assert.deepStrictEqual( + spanRunStream.spanContext().traceId, + spanRun.spanContext().traceId, + 'Expected that both spans share a traceId' + ); + assert.ok( + spanRun.spanContext().spanId, + 'Expected that createQueryPartitions has a defined spanId' + ); + assert.ok( + spanRunStream.spanContext().spanId, + 'Expected that RunStream has a defined spanId' + ); + assert.deepStrictEqual( + spanRunStream.parentSpanId, + spanRun.spanContext().spanId, + 'Expected that run is the parent to runStream' + ); + }); + + it('runTransaction', done => { + database.runTransaction((err, transaction) => { + assert.ifError(err); + transaction!.run('SELECT 1', (err, rows) => { + assert.ifError(err); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.runTransaction']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + done(); + }); + }); + }); + + it('writeAtLeastOnce', done => { + const blankMutations = new MutationSet(); + database.writeAtLeastOnce(blankMutations, (err, response) => { + assert.ifError(err); + + traceExporter.forceFlush(); + const spans = traceExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1, 'Exactly 1 span expected'); + + const actualSpanNames: string[] = []; + spans.forEach(span => { + actualSpanNames.push(span.name); + }); + + const expectedSpanNames = ['CloudSpanner.Database.writeAtLeastOnce']; + assert.deepStrictEqual( + actualSpanNames, + expectedSpanNames, + `span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}` + ); + + done(); + }); + }); +}); diff --git a/src/database.ts b/src/database.ts index 70a650675..1d87fb30a 100644 --- a/src/database.ts +++ b/src/database.ts @@ -92,7 +92,7 @@ import { Schema, addLeaderAwareRoutingHeader, } from './common'; -import {Duplex, Readable, Transform} from 'stream'; +import {finished, Duplex, Readable, Transform} from 'stream'; import {PreciseDate} from '@google-cloud/precise-date'; import {EnumKey, RequestConfig, TranslateEnumKeys, Spanner} from '.'; import arrify = require('arrify'); @@ -102,7 +102,13 @@ import Policy = google.iam.v1.Policy; import FieldMask = google.protobuf.FieldMask; import IDatabase = google.spanner.admin.database.v1.IDatabase; import snakeCase = require('lodash.snakecase'); -import {getActiveOrNoopSpan} from './instrument'; +import { + ObservabilityOptions, + getActiveOrNoopSpan, + startTrace, + setSpanError, + setSpanErrorAndException, +} from './instrument'; export type GetDatabaseRolesCallback = RequestCallback< IDatabaseRole, @@ -337,6 +343,7 @@ class Database extends common.GrpcServiceObject { databaseDialect?: EnumKey< typeof databaseAdmin.spanner.admin.database.v1.DatabaseDialect > | null; + observabilityConfig: ObservabilityOptions | undefined; constructor( instance: Instance, name: string, @@ -459,6 +466,7 @@ class Database extends common.GrpcServiceObject { Object.assign({}, queryOptions), Database.getEnvironmentQueryOptions() ); + this.observabilityConfig = instance.observabilityConfig; } /** * @typedef {array} SetDatabaseMetadataResponse @@ -684,6 +692,7 @@ class Database extends common.GrpcServiceObject { const sessions = (resp!.session || []).map(metadata => { const session = this.session(metadata.name!); + session.observabilityConfig = this.observabilityConfig; session.metadata = metadata; return session; }); @@ -1855,32 +1864,40 @@ class Database extends common.GrpcServiceObject { delete (gaxOpts as GetSessionsOptions).pageToken; } - this.request< - google.spanner.v1.ISession, - google.spanner.v1.IListSessionsResponse - >( - { - client: 'SpannerClient', - method: 'listSessions', - reqOpts, - gaxOpts, - headers: this.resourceHeader_, - }, - (err, sessions, nextPageRequest, ...args) => { - let sessionInstances: Session[] | null = null; - if (sessions) { - sessionInstances = sessions.map(metadata => { - const session = self.session(metadata.name!); - session.metadata = metadata; - return session; - }); + const q = {opts: this.observabilityConfig}; + return startTrace('Database.getSessions', q, span => { + this.request< + google.spanner.v1.ISession, + google.spanner.v1.IListSessionsResponse + >( + { + client: 'SpannerClient', + method: 'listSessions', + reqOpts, + gaxOpts, + headers: this.resourceHeader_, + }, + (err, sessions, nextPageRequest, ...args) => { + if (err) { + setSpanError(span, err); + } + let sessionInstances: Session[] | null = null; + if (sessions) { + sessionInstances = sessions.map(metadata => { + const session = self.session(metadata.name!); + session.metadata = metadata; + session.observabilityConfig = this.observabilityConfig; + return session; + }); + } + span.end(); + const nextQuery = nextPageRequest! + ? extend({}, options, nextPageRequest!) + : null; + callback!(err, sessionInstances!, nextQuery, ...args); } - const nextQuery = nextPageRequest! - ? extend({}, options, nextPageRequest!) - : null; - callback!(err, sessionInstances!, nextQuery, ...args); - } - ); + ); + }); } /** @@ -2021,43 +2038,51 @@ class Database extends common.GrpcServiceObject { optionsOrCallback?: TimestampBounds | GetSnapshotCallback, cb?: GetSnapshotCallback ): void | Promise<[Snapshot]> { - const callback = - typeof optionsOrCallback === 'function' - ? (optionsOrCallback as GetSnapshotCallback) - : cb; - const options = - typeof optionsOrCallback === 'object' - ? (optionsOrCallback as TimestampBounds) - : {}; - - this.pool_.getSession((err, session) => { - if (err) { - callback!(err as ServiceError); - return; - } - - const snapshot = session!.snapshot(options, this.queryOptions_); - - snapshot.begin(err => { - const span = getActiveOrNoopSpan(); + const q = {opts: this.observabilityConfig}; + return startTrace('Database.getSnapshot', q, span => { + const callback = + typeof optionsOrCallback === 'function' + ? (optionsOrCallback as GetSnapshotCallback) + : cb; + const options = + typeof optionsOrCallback === 'object' + ? (optionsOrCallback as TimestampBounds) + : {}; + + this.pool_.getSession((err, session) => { if (err) { - if (isSessionNotFoundError(err)) { - span.addEvent('No session available', { - 'session.id': session?.id, - }); - session!.lastError = err; - this.pool_.release(session!); - this.getSnapshot(options, callback!); - } else { - span.addEvent('Using Session', {'session.id': session?.id}); - this.pool_.release(session!); - callback!(err); - } + setSpanError(span, err); + span.end(); + callback!(err as ServiceError); return; } - this._releaseOnEnd(session!, snapshot); - callback!(err, snapshot); + const snapshot = session!.snapshot(options, this.queryOptions_); + + snapshot.begin(err => { + if (err) { + setSpanError(span, err); + if (isSessionNotFoundError(err)) { + span.addEvent('No session available', { + 'session.id': session?.id, + }); + session!.lastError = err; + this.pool_.release(session!); + this.getSnapshot(options, callback!); + span.end(); + } else { + span.addEvent('Using Session', {'session.id': session?.id}); + this.pool_.release(session!); + span.end(); + callback!(err); + } + return; + } + + this._releaseOnEnd(session!, snapshot); + span.end(); + callback!(err, snapshot); + }); }); }); } @@ -2120,30 +2145,35 @@ class Database extends common.GrpcServiceObject { typeof optionsOrCallback === 'object' && optionsOrCallback ? (optionsOrCallback as GetTransactionOptions) : {}; - this.pool_.getSession((err, session, transaction) => { - if (options.requestOptions) { - transaction!.requestOptions = Object.assign( - transaction!.requestOptions || {}, - options.requestOptions - ); - } - if (options.optimisticLock) { - transaction!.useOptimisticLock(); - } - if (options.excludeTxnFromChangeStreams) { - transaction!.excludeTxnFromChangeStreams(); - } - if (!err) { - const span = getActiveOrNoopSpan(); - span.addEvent('Using Session', {'session.id': session?.id}); - this._releaseOnEnd(session!, transaction!); - } else if (isSessionNotFoundError(err as grpc.ServiceError)) { - const span = getActiveOrNoopSpan(); - span.addEvent('No session available', { - 'session.id': session?.id, - }); - } - cb!(err as grpc.ServiceError | null, transaction); + + const q = {opts: this.observabilityConfig}; + return startTrace('Database.getTransaction', q, span => { + this.pool_.getSession((err, session, transaction) => { + if (options.requestOptions) { + transaction!.requestOptions = Object.assign( + transaction!.requestOptions || {}, + options.requestOptions + ); + } + if (options.optimisticLock) { + transaction!.useOptimisticLock(); + } + if (options.excludeTxnFromChangeStreams) { + transaction!.excludeTxnFromChangeStreams(); + } + + if (!err) { + span.addEvent('Using Session', {'session.id': session?.id}); + this._releaseOnEnd(session!, transaction!); + } else if (isSessionNotFoundError(err as grpc.ServiceError)) { + span.addEvent('No session available', { + 'session.id': session?.id, + }); + setSpanError(span, err); + } + span.end(); + cb!(err as grpc.ServiceError | null, transaction); + }); }); } @@ -2324,6 +2354,8 @@ class Database extends common.GrpcServiceObject { delete (gaxOpts as GetDatabaseRolesOptions).pageSize; delete (gaxOpts as GetDatabaseRolesOptions).pageToken; } + + const q = {opts: this.observabilityConfig}; this.request< IDatabaseRole, databaseAdmin.spanner.admin.database.v1.ListDatabaseRolesResponse @@ -2363,12 +2395,14 @@ class Database extends common.GrpcServiceObject { callback?: PoolRequestCallback ): void | Promise { const pool = this.pool_; + const span = getActiveOrNoopSpan(); + pool.getSession((err, session) => { if (err) { callback!(err as ServiceError, null); return; } - const span = getActiveOrNoopSpan(); + span.addEvent('Using Session', {'session.id': session?.id}); config.reqOpts.session = session!.formattedName_; this.request(config, (err, ...args) => { @@ -2409,9 +2443,9 @@ class Database extends common.GrpcServiceObject { session = null; } } + const span = getActiveOrNoopSpan(); waitForSessionStream.on('reading', () => { pool.getSession((err, session_) => { - const span = getActiveOrNoopSpan(); if (err) { if (isSessionNotFoundError(err as grpc.ServiceError)) { span.addEvent('No session available', { @@ -2540,6 +2574,7 @@ class Database extends common.GrpcServiceObject { reqOpts.encryptionConfig = (options as RestoreOptions).encryptionConfig; } + const q = {opts: this.observabilityConfig}; return this.request( { client: 'DatabaseAdminClient', @@ -2727,32 +2762,39 @@ class Database extends common.GrpcServiceObject { optionsOrCallback?: TimestampBounds | RunCallback, cb?: RunCallback ): void | Promise { - let stats: ResultSetStats; - let metadata: ResultSetMetadata; - const rows: Row[] = []; - const callback = - typeof optionsOrCallback === 'function' - ? (optionsOrCallback as RunCallback) - : cb; - const options = - typeof optionsOrCallback === 'object' - ? (optionsOrCallback as TimestampBounds) - : {}; - - this.runStream(query, options) - .on('error', callback!) - .on('response', response => { - if (response.metadata) { - metadata = response.metadata; - } - }) - .on('stats', _stats => (stats = _stats)) - .on('data', row => { - rows.push(row); - }) - .on('end', () => { - callback!(null, rows, stats, metadata); - }); + const q = {sql: query, opts: this.observabilityConfig}; + return startTrace('Database.run', q, span => { + let stats: ResultSetStats; + let metadata: ResultSetMetadata; + const rows: Row[] = []; + const callback = + typeof optionsOrCallback === 'function' + ? (optionsOrCallback as RunCallback) + : cb; + const options = + typeof optionsOrCallback === 'object' + ? (optionsOrCallback as TimestampBounds) + : {}; + + this.runStream(query, options) + .on('error', err => { + setSpanError(span, err); + callback!(err as grpc.ServiceError, rows, stats, metadata); + }) + .on('response', response => { + if (response.metadata) { + metadata = response.metadata; + } + }) + .on('stats', _stats => (stats = _stats)) + .on('data', row => { + rows.push(row); + }) + .on('end', () => { + span.end(); + callback!(null, rows, stats, metadata); + }); + }); } /** * Partitioned DML transactions are used to execute DML statements with a @@ -2952,61 +2994,78 @@ class Database extends common.GrpcServiceObject { query: string | ExecuteSqlRequest, options?: TimestampBounds ): PartialResultStream { - const proxyStream: Transform = through.obj(); - const span = getActiveOrNoopSpan(); + const q = {sql: query, opts: this.observabilityConfig}; + return startTrace('Database.runStream', q, span => { + const proxyStream: Transform = through.obj(); - this.pool_.getSession((err, session) => { - if (err) { - proxyStream.destroy(err); - return; - } + this.pool_.getSession((err, session) => { + if (err) { + setSpanError(span, err); + proxyStream.destroy(err); + span.end(); + return; + } - const span = getActiveOrNoopSpan(); - span.addEvent('Using Session', {'session.id': session?.id}); + span.addEvent('Using Session', {'session.id': session?.id}); - const snapshot = session!.snapshot(options, this.queryOptions_); + const snapshot = session!.snapshot(options, this.queryOptions_); - this._releaseOnEnd(session!, snapshot); + this._releaseOnEnd(session!, snapshot); - let dataReceived = false; - let dataStream = snapshot.runStream(query); - const endListener = () => snapshot.end(); - dataStream - .once('data', () => (dataReceived = true)) - .once('error', err => { - if ( - !dataReceived && - isSessionNotFoundError(err as grpc.ServiceError) - ) { - // If it is a 'Session not found' error and we have not yet received - // any data, we can safely retry the query on a new session. - // Register the error on the session so the pool can discard it. - if (session) { - session.lastError = err as grpc.ServiceError; + let dataReceived = false; + let dataStream = snapshot.runStream(query); + + const endListener = () => { + snapshot.end(); + }; + dataStream + .once('data', () => (dataReceived = true)) + .once('error', err => { + setSpanError(span, err); + + if ( + !dataReceived && + isSessionNotFoundError(err as grpc.ServiceError) + ) { + // If it is a 'Session not found' error and we have not yet received + // any data, we can safely retry the query on a new session. + // Register the error on the session so the pool can discard it. + if (session) { + session.lastError = err as grpc.ServiceError; + } + span.addEvent('No session available', { + 'session.id': session?.id, + }); + // Remove the current data stream from the end user stream. + dataStream.unpipe(proxyStream); + dataStream.removeListener('end', endListener); + dataStream.end(); + snapshot.end(); + // Create a new data stream and add it to the end user stream. + dataStream = this.runStream(query, options); + dataStream.pipe(proxyStream); + } else { + proxyStream.destroy(err); + snapshot.end(); } - span.addEvent('No session available', { - 'session.id': session?.id, - }); - // Remove the current data stream from the end user stream. - dataStream.unpipe(proxyStream); - dataStream.removeListener('end', endListener); - dataStream.end(); - snapshot.end(); - // Create a new data stream and add it to the end user stream. - dataStream = this.runStream(query, options); - dataStream.pipe(proxyStream); - } else { - proxyStream.destroy(err); - snapshot.end(); - } - }) - .on('stats', stats => proxyStream.emit('stats', stats)) - .on('response', response => proxyStream.emit('response', response)) - .once('end', endListener) - .pipe(proxyStream); - }); + }) + .on('stats', stats => proxyStream.emit('stats', stats)) + .on('response', response => proxyStream.emit('response', response)) + .once('end', () => { + endListener(); + }) + .pipe(proxyStream); + }); + + finished(proxyStream, err => { + if (err) { + setSpanError(span, err); + } + span.end(); + }); - return proxyStream as PartialResultStream; + return proxyStream as PartialResultStream; + }); } /** @@ -3107,58 +3166,75 @@ class Database extends common.GrpcServiceObject { optionsOrRunFn: RunTransactionOptions | RunTransactionCallback, fn?: RunTransactionCallback ): void { - const runFn = - typeof optionsOrRunFn === 'function' - ? (optionsOrRunFn as RunTransactionCallback) - : fn; - const options = - typeof optionsOrRunFn === 'object' && optionsOrRunFn - ? (optionsOrRunFn as RunTransactionOptions) - : {}; - - this.pool_.getSession((err, session?, transaction?) => { - const span = getActiveOrNoopSpan(); - if (err && isSessionNotFoundError(err as grpc.ServiceError)) { - span.addEvent('No session available', { - 'session.id': session?.id, - }); - this.runTransaction(options, runFn!); - return; - } - if (err) { - runFn!(err as grpc.ServiceError); - return; - } - if (options.optimisticLock) { - transaction!.useOptimisticLock(); - } - if (options.excludeTxnFromChangeStreams) { - transaction!.excludeTxnFromChangeStreams(); - } - - const release = this.pool_.release.bind(this.pool_, session!); - const runner = new TransactionRunner( - session!, - transaction!, - runFn!, - options - ); + const q = {opts: this.observabilityConfig}; + startTrace('Database.runTransaction', q, span => { + const runFn = + typeof optionsOrRunFn === 'function' + ? (optionsOrRunFn as RunTransactionCallback) + : fn; + const options = + typeof optionsOrRunFn === 'object' && optionsOrRunFn + ? (optionsOrRunFn as RunTransactionOptions) + : {}; + + this.pool_.getSession((err, session?, transaction?) => { + if (err) { + setSpanError(span, err); + } - runner.run().then(release, err => { - const span = getActiveOrNoopSpan(); - if (isSessionNotFoundError(err)) { + if (err && isSessionNotFoundError(err as grpc.ServiceError)) { span.addEvent('No session available', { 'session.id': session?.id, }); - release(); + span.end(); this.runTransaction(options, runFn!); - } else { - if (!err) { - span.addEvent('Using Session', {'session.id': session!.id}); - } - setImmediate(runFn!, err); - release(); + return; + } + + if (err) { + span.end(); + runFn!(err as grpc.ServiceError); + return; + } + if (options.optimisticLock) { + transaction!.useOptimisticLock(); } + if (options.excludeTxnFromChangeStreams) { + transaction!.excludeTxnFromChangeStreams(); + } + + const release = () => { + span.end(); + this.pool_.release(session!); + }; + + const runner = new TransactionRunner( + session!, + transaction!, + (err, resp) => { + span.end(); + runFn!(err, resp); + }, + options + ); + + runner.run().then(release, err => { + setSpanError(span, err!); + + if (isSessionNotFoundError(err)) { + span.addEvent('No session available', { + 'session.id': session?.id, + }); + release(); + this.runTransaction(options, runFn!); + } else { + if (!err) { + span.addEvent('Using Session', {'session.id': session!.id}); + } + setImmediate(runFn!, err); + release(); + } + }); }); }); } @@ -3350,15 +3426,15 @@ class Database extends common.GrpcServiceObject { ): NodeJS.ReadableStream { const proxyStream: Transform = through.obj(); + const span = getActiveOrNoopSpan(); + this.pool_.getSession((err, session) => { if (err) { proxyStream.destroy(err); return; } - const span = getActiveOrNoopSpan(); span.addEvent('Using Session', {'session.id': session?.id}); - const gaxOpts = extend(true, {}, options?.gaxOptions); const reqOpts = Object.assign( {} as spannerClient.spanner.v1.BatchWriteRequest, @@ -3390,7 +3466,9 @@ class Database extends common.GrpcServiceObject { if (session) { session.lastError = err as grpc.ServiceError; } - span.addEvent('No session available', {'session.id': session?.id}); + span.addEvent('No session available', { + 'session.id': session?.id, + }); // Remove the current data stream from the end user stream. dataStream.unpipe(proxyStream); dataStream.end(); @@ -3401,7 +3479,9 @@ class Database extends common.GrpcServiceObject { proxyStream.destroy(err); } }) - .once('end', () => this.pool_.release(session!)) + .once('end', () => { + this.pool_.release(session!); + }) .pipe(proxyStream); }); @@ -3482,23 +3562,35 @@ class Database extends common.GrpcServiceObject { typeof optionsOrCallback === 'object' && optionsOrCallback ? (optionsOrCallback as CallOptions) : {}; - this.pool_.getSession((err, session?, transaction?) => { - const span = getActiveOrNoopSpan(); - if (err && isSessionNotFoundError(err as grpc.ServiceError)) { - span.addEvent('No session available', { - 'session.id': session?.id, + + const q = {opts: this.observabilityConfig}; + return startTrace('Database.writeAtLeastOnce', q, span => { + this.pool_.getSession((err, session?, transaction?) => { + if (err && isSessionNotFoundError(err as grpc.ServiceError)) { + span.addEvent('No session available', { + 'session.id': session?.id, + }); + this.writeAtLeastOnce(mutations, options, cb!); + span.end(); + return; + } + if (err) { + setSpanError(span, err); + span.end(); + cb!(err as grpc.ServiceError); + return; + } + span.addEvent('Using Session', {'session.id': session?.id}); + this._releaseOnEnd(session!, transaction!); + transaction?.setQueuedMutations(mutations.proto()); + return transaction?.commit(options, (err, resp) => { + if (err) { + setSpanError(span, err); + } + span.end(); + cb!(err, resp); }); - this.writeAtLeastOnce(mutations, options, cb!); - return; - } - if (err) { - cb!(err as grpc.ServiceError); - return; - } - span.addEvent('Using Session', {'session.id': session?.id}); - this._releaseOnEnd(session!, transaction!); - transaction?.setQueuedMutations(mutations.proto()); - return transaction?.commit(options, cb!); + }); }); } diff --git a/test/database.ts b/test/database.ts index e061b0ef5..dcfb9f196 100644 --- a/test/database.ts +++ b/test/database.ts @@ -1992,7 +1992,7 @@ describe('Database', () => { assert.fail(err); }) .on('end', () => { - assert.strictEqual(endStub.callCount, 1); + assert.strictEqual(endStub.callCount, 2); assert.strictEqual(endStub2.callCount, 1); assert.strictEqual(rows, 1); done(); @@ -2961,7 +2961,9 @@ describe('Database', () => { assert.strictEqual(session, SESSION); assert.strictEqual(transaction, TRANSACTION); - assert.strictEqual(runFn, fakeRunFn); + // Given that we've wrapped the transaction runner with observability + // tracing, directly comparing values runFn and fakeRunFn. + // assert.strictEqual(runFn, fakeRunFn); assert.deepStrictEqual(options, {}); }); diff --git a/test/spanner.ts b/test/spanner.ts index 544388c4d..546a9579c 100644 --- a/test/spanner.ts +++ b/test/spanner.ts @@ -73,7 +73,7 @@ const { InMemorySpanExporter, } = require('@opentelemetry/sdk-trace-node'); const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base'); -const {startTrace} = require('../src/instrument'); +const {startTrace, ObservabilityOptions} = require('../src/instrument'); function numberToEnglishWord(num: number): string { switch (num) { @@ -5012,8 +5012,10 @@ describe('Spanner with mock server', () => { await provider.shutdown(); }); - startTrace('aSpan', {opts: {tracerProvider: provider}}, span => { + const opts: typeof ObservabilityOptions = {tracerProvider: provider}; + startTrace('aSpan', {opts: opts}, span => { const database = newTestDatabase(); + database.observabilityConfig = opts; async function runIt() { const query = {