Skip to content

Commit ec48f2c

Browse files
odeke-emodeke-em
committed
feat(observability): add tracing spans to Session
This change adds tracing spans to Session for methods: * create * getMetadata * keepAlive and corresponding tests. Updates googleapis#2079 Built from PR googleapis#2087 Updates googleapis#2114
1 parent d51aae9 commit ec48f2c

File tree

2 files changed

+318
-33
lines changed

2 files changed

+318
-33
lines changed

‎observability-test/session.ts

+252
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,252 @@
1+
/*!
2+
* Copyright 2024 Google LLC. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
import * as assert from 'assert';
18+
import {grpc} from 'google-gax';
19+
import {google} from '../protos/protos';
20+
import {Database, Session, Spanner} from '../src';
21+
import protobuf = google.spanner.v1;
22+
import * as mock from '../test/mockserver/mockspanner';
23+
import * as mockInstanceAdmin from '../test/mockserver/mockinstanceadmin';
24+
import * as mockDatabaseAdmin from '../test/mockserver/mockdatabaseadmin';
25+
const {
26+
AlwaysOnSampler,
27+
NodeTracerProvider,
28+
InMemorySpanExporter,
29+
} = require('@opentelemetry/sdk-trace-node');
30+
// eslint-disable-next-line n/no-extraneous-require
31+
const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base');
32+
33+
/** A simple result set for SELECT 1. */
34+
function createSelect1ResultSet(): protobuf.ResultSet {
35+
const fields = [
36+
protobuf.StructType.Field.create({
37+
name: 'NUM',
38+
type: protobuf.Type.create({code: protobuf.TypeCode.INT64}),
39+
}),
40+
];
41+
const metadata = new protobuf.ResultSetMetadata({
42+
rowType: new protobuf.StructType({
43+
fields,
44+
}),
45+
});
46+
return protobuf.ResultSet.create({
47+
metadata,
48+
rows: [{values: [{stringValue: '1'}]}],
49+
});
50+
}
51+
52+
interface setupResults {
53+
server: grpc.Server;
54+
spanner: Spanner;
55+
spannerMock: mock.MockSpanner;
56+
}
57+
58+
async function setup(): Promise<setupResults> {
59+
const server = new grpc.Server();
60+
61+
const spannerMock = mock.createMockSpanner(server);
62+
mockInstanceAdmin.createMockInstanceAdmin(server);
63+
mockDatabaseAdmin.createMockDatabaseAdmin(server);
64+
65+
const port: number = await new Promise((resolve, reject) => {
66+
server.bindAsync(
67+
'0.0.0.0:0',
68+
grpc.ServerCredentials.createInsecure(),
69+
(err, assignedPort) => {
70+
if (err) {
71+
reject(err);
72+
} else {
73+
resolve(assignedPort);
74+
}
75+
}
76+
);
77+
});
78+
79+
const selectSql = 'SELECT 1';
80+
const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2';
81+
spannerMock.putStatementResult(
82+
selectSql,
83+
mock.StatementResult.resultSet(createSelect1ResultSet())
84+
);
85+
spannerMock.putStatementResult(
86+
updateSql,
87+
mock.StatementResult.updateCount(1)
88+
);
89+
90+
const spanner = new Spanner({
91+
projectId: 'observability-project-id',
92+
servicePath: 'localhost',
93+
port,
94+
sslCreds: grpc.credentials.createInsecure(),
95+
});
96+
97+
return Promise.resolve({
98+
spanner: spanner,
99+
server: server,
100+
spannerMock: spannerMock,
101+
});
102+
}
103+
104+
describe('Session', () => {
105+
let server: grpc.Server;
106+
let spanner: Spanner;
107+
let database: Database;
108+
let spannerMock: mock.MockSpanner;
109+
let traceExporter: typeof InMemorySpanExporter;
110+
111+
after(() => {
112+
spanner.close();
113+
server.tryShutdown(() => {});
114+
});
115+
116+
before(async () => {
117+
const setupResult = await setup();
118+
spanner = setupResult.spanner;
119+
server = setupResult.server;
120+
spannerMock = setupResult.spannerMock;
121+
122+
const selectSql = 'SELECT 1';
123+
const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2';
124+
spannerMock.putStatementResult(
125+
selectSql,
126+
mock.StatementResult.resultSet(createSelect1ResultSet())
127+
);
128+
spannerMock.putStatementResult(
129+
updateSql,
130+
mock.StatementResult.updateCount(1)
131+
);
132+
133+
traceExporter = new InMemorySpanExporter();
134+
const sampler = new AlwaysOnSampler();
135+
136+
const provider = new NodeTracerProvider({
137+
sampler: sampler,
138+
exporter: traceExporter,
139+
});
140+
provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter));
141+
142+
const instance = spanner.instance('instance');
143+
database = instance.database('database');
144+
database.observabilityConfig = {
145+
tracerProvider: provider,
146+
enableExtendedTracing: false,
147+
};
148+
});
149+
150+
beforeEach(() => {
151+
spannerMock.resetRequests();
152+
});
153+
154+
afterEach(() => {
155+
traceExporter.reset();
156+
});
157+
158+
it('create with constructor', done => {
159+
const session = new Session(database);
160+
session.create(err => {
161+
traceExporter.forceFlush();
162+
const spans = traceExporter.getFinishedSpans();
163+
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
164+
165+
const actualSpanNames: string[] = [];
166+
spans.forEach(span => {
167+
actualSpanNames.push(span.name);
168+
});
169+
170+
const expectedSpanNames = ['CloudSpanner.Session.create'];
171+
assert.deepStrictEqual(
172+
actualSpanNames,
173+
expectedSpanNames,
174+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
175+
);
176+
177+
done();
178+
});
179+
});
180+
181+
it('create with database.session()', done => {
182+
const session = database.session();
183+
session.create(err => {
184+
assert.ifError(err);
185+
traceExporter.forceFlush();
186+
const spans = traceExporter.getFinishedSpans();
187+
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
188+
189+
const actualSpanNames: string[] = [];
190+
spans.forEach(span => {
191+
actualSpanNames.push(span.name);
192+
});
193+
194+
const expectedSpanNames = ['CloudSpanner.Session.create'];
195+
assert.deepStrictEqual(
196+
actualSpanNames,
197+
expectedSpanNames,
198+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
199+
);
200+
201+
done();
202+
});
203+
});
204+
205+
it('getMetadata', done => {
206+
const session = database.session();
207+
session.create(err => {
208+
assert.ifError(err);
209+
traceExporter.forceFlush();
210+
traceExporter.reset();
211+
212+
session.getMetadata((err, metadata) => {
213+
const spans = traceExporter.getFinishedSpans();
214+
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
215+
const span = spans[0];
216+
217+
const expectedSpanName = 'CloudSpanner.Session.getMetadata';
218+
assert.deepStrictEqual(
219+
span.name,
220+
expectedSpanName,
221+
`span names mismatch:\n\tGot: ${span.name}\n\tWant: ${expectedSpanName}`
222+
);
223+
224+
done();
225+
});
226+
});
227+
});
228+
229+
it('keepAlive', done => {
230+
const session = database.session();
231+
session.create(err => {
232+
assert.ifError(err);
233+
traceExporter.forceFlush();
234+
traceExporter.reset();
235+
236+
session.keepAlive(err => {
237+
const spans = traceExporter.getFinishedSpans();
238+
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
239+
const span = spans[0];
240+
241+
const expectedSpanName = 'CloudSpanner.Session.keepAlive';
242+
assert.deepStrictEqual(
243+
span.name,
244+
expectedSpanName,
245+
`span names mismatch:\n\tGot: ${span.name}\n\tWant: ${expectedSpanName}`
246+
);
247+
248+
done();
249+
});
250+
});
251+
});
252+
});

‎src/session.ts

+66-33
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import {
4444
import {grpc, CallOptions} from 'google-gax';
4545
import IRequestOptions = google.spanner.v1.IRequestOptions;
4646
import {Spanner} from '.';
47+
import {ObservabilityOptions, startTrace, setSpanError} from './instrument';
4748

4849
export type GetSessionResponse = [Session, r.Response];
4950

@@ -118,6 +119,7 @@ export class Session extends common.GrpcServiceObject {
118119
lastUsed?: number;
119120
lastError?: grpc.ServiceError;
120121
resourceHeader_: {[k: string]: string};
122+
observabilityConfig: ObservabilityOptions | undefined;
121123
constructor(database: Database, name?: string) {
122124
const methods = {
123125
/**
@@ -247,14 +249,24 @@ export class Session extends common.GrpcServiceObject {
247249
this.databaseRole =
248250
options.databaseRole || database.databaseRole || null;
249251

250-
return database.createSession(options, (err, session, apiResponse) => {
251-
if (err) {
252-
callback(err, null, apiResponse);
253-
return;
254-
}
252+
const q = {opts: database.observabilityConfig};
253+
return startTrace('Session.create', q, span => {
254+
return database.createSession(
255+
options,
256+
(err, session, apiResponse) => {
257+
if (err) {
258+
setSpanError(span, err);
259+
span.end();
260+
callback(err, null, apiResponse);
261+
return;
262+
}
255263

256-
extend(this, session);
257-
callback(null, this, apiResponse);
264+
this.observabilityConfig = database.ovservabilityConfig;
265+
span.end();
266+
extend(this, session);
267+
callback(null, this, apiResponse);
268+
}
269+
);
258270
});
259271
},
260272
} as {} as ServiceObjectConfig);
@@ -268,6 +280,8 @@ export class Session extends common.GrpcServiceObject {
268280
if (name) {
269281
this.formattedName_ = Session.formatName_(database.formattedName_, name);
270282
}
283+
284+
this.observabilityConfig = database.observabilityConfig;
271285
}
272286
/**
273287
* Delete a session.
@@ -388,23 +402,32 @@ export class Session extends common.GrpcServiceObject {
388402
if (this._getSpanner().routeToLeaderEnabled) {
389403
addLeaderAwareRoutingHeader(headers);
390404
}
391-
return this.request(
392-
{
393-
client: 'SpannerClient',
394-
method: 'getSession',
395-
reqOpts,
396-
gaxOpts,
397-
headers: headers,
398-
},
399-
(err, resp) => {
400-
if (resp) {
401-
resp.databaseRole = resp.creatorRole;
402-
delete resp.creatorRole;
403-
this.metadata = resp;
405+
406+
const q = {opts: this.observabilityConfig};
407+
return startTrace('Session.getMetadata', q, span => {
408+
return this.request(
409+
{
410+
client: 'SpannerClient',
411+
method: 'getSession',
412+
reqOpts,
413+
gaxOpts,
414+
headers: headers,
415+
},
416+
(err, resp) => {
417+
if (err) {
418+
setSpanError(span, err);
419+
}
420+
421+
if (resp) {
422+
resp.databaseRole = resp.creatorRole;
423+
delete resp.creatorRole;
424+
this.metadata = resp;
425+
}
426+
span.end();
427+
callback!(err, resp);
404428
}
405-
callback!(err, resp);
406-
}
407-
);
429+
);
430+
});
408431
}
409432
/**
410433
* Ping the session with `SELECT 1` to prevent it from expiring.
@@ -440,16 +463,26 @@ export class Session extends common.GrpcServiceObject {
440463
session: this.formattedName_,
441464
sql: 'SELECT 1',
442465
};
443-
return this.request(
444-
{
445-
client: 'SpannerClient',
446-
method: 'executeSql',
447-
reqOpts,
448-
gaxOpts,
449-
headers: this.resourceHeader_,
450-
},
451-
callback!
452-
);
466+
467+
const q = {opts: this.observabilityConfig, sql: reqOpts.sql};
468+
return startTrace('Session.keepAlive', q, span => {
469+
return this.request(
470+
{
471+
client: 'SpannerClient',
472+
method: 'executeSql',
473+
reqOpts,
474+
gaxOpts,
475+
headers: this.resourceHeader_,
476+
},
477+
err => {
478+
if (err) {
479+
setSpanError(span, err);
480+
}
481+
span.end();
482+
callback!(err);
483+
}
484+
);
485+
});
453486
}
454487
/**
455488
* Create a PartitionedDml transaction.

0 commit comments

Comments
 (0)