Skip to content

Commit

Permalink
feat(observability): trace Transaction
Browse files Browse the repository at this point in the history
This change adds observability tracing for
Transaction along with tests.

Updates #2079
Built from PR #2087
Updates #2114
  • Loading branch information
odeke-em committed Oct 1, 2024
1 parent 1f06871 commit 838225e
Show file tree
Hide file tree
Showing 8 changed files with 654 additions and 316 deletions.
2 changes: 1 addition & 1 deletion observability-test/batch-transaction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ describe('BatchTransaction', () => {
batchTransaction = new BatchTransaction(SESSION as {} as Session);
batchTransaction.session = SESSION as {} as Session;
batchTransaction.id = ID;
batchTransaction.observabilityOptions = {tracerProvider: provider};
batchTransaction.observabilityOptions_ = {tracerProvider: provider};
REQUEST.callsFake((_, callback) => callback(null, RESPONSE));
});

Expand Down
2 changes: 1 addition & 1 deletion observability-test/database.ts
Original file line number Diff line number Diff line change
Expand Up @@ -241,7 +241,7 @@ describe('Database', () => {
database = new Database(INSTANCE, NAME, POOL_OPTIONS);
database.parent = INSTANCE;
database.databaseRole = 'parent_role';
database.observabilityConfig = {
database.observabilityOptions_ = {
tracerProvider: provider,
enableExtendedTracing: false,
};
Expand Down
35 changes: 25 additions & 10 deletions observability-test/spanner.ts
Original file line number Diff line number Diff line change
Expand Up @@ -149,7 +149,7 @@ describe('EndToEnd', () => {

const instance = spanner.instance('instance');
database = instance.database('database');
database.observabilityConfig = {
database.observabilityOptions_ = {
tracerProvider: provider,
enableExtendedTracing: false,
};
Expand Down Expand Up @@ -202,7 +202,7 @@ describe('EndToEnd', () => {

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
assert.strictEqual(spans.length, 3, 'Exactly 1 span expected');

const actualSpanNames: string[] = [];
const actualEventNames: string[] = [];
Expand All @@ -213,14 +213,18 @@ describe('EndToEnd', () => {
});
});

const expectedSpanNames = ['CloudSpanner.Database.getSnapshot'];
const expectedSpanNames = [
'CloudSpanner.Snapshot.begin',
'CloudSpanner.Database.getSnapshot',
'CloudSpanner.Snapshot.run',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);

const expectedEventNames = [];
const expectedEventNames = ['Begin Transaction'];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
Expand Down Expand Up @@ -310,7 +314,7 @@ describe('EndToEnd', () => {

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');
assert.strictEqual(spans.length, 3, 'Exactly 2 spans expected');

// Sort the spans by duration.
spans.sort((spanA, spanB) => {
Expand All @@ -329,6 +333,7 @@ describe('EndToEnd', () => {
const expectedSpanNames = [
'CloudSpanner.Database.runStream',
'CloudSpanner.Database.run',
'CloudSpanner.Snapshot.runStream',
];
assert.deepStrictEqual(
actualSpanNames,
Expand Down Expand Up @@ -372,7 +377,7 @@ describe('EndToEnd', () => {

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
assert.strictEqual(spans.length, 2, 'Exactly 1 span expected');

const actualEventNames: string[] = [];
const actualSpanNames: string[] = [];
Expand All @@ -383,7 +388,10 @@ describe('EndToEnd', () => {
});
});

const expectedSpanNames = ['CloudSpanner.Database.runTransaction'];
const expectedSpanNames = [
'CloudSpanner.Database.runTransaction',
'CloudSpanner.Snapshot.run',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
Expand All @@ -410,7 +418,7 @@ describe('EndToEnd', () => {

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 1, 'Exactly 1 span expected');
assert.strictEqual(spans.length, 2, 'Exactly 1 span expected');

const actualEventNames: string[] = [];
const actualSpanNames: string[] = [];
Expand All @@ -421,14 +429,21 @@ describe('EndToEnd', () => {
});
});

const expectedSpanNames = ['CloudSpanner.Database.writeAtLeastOnce'];
const expectedSpanNames = [
'CloudSpanner.Transaction.commit',
'CloudSpanner.Database.writeAtLeastOnce',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);

const expectedEventNames = ['Using Session'];
const expectedEventNames = [
'Starting Commit',
'Commit Done',
'Using Session',
];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
Expand Down
203 changes: 203 additions & 0 deletions observability-test/transaction.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,203 @@
/*!
* 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, Spanner, Transaction} from '../src';
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;
}

const selectSql = 'SELECT 1';
const updateSql = 'UPDATE FOO SET BAR=1 WHERE BAZ=2';

async function setup(): Promise<setupResults> {
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);
}
}
);
});
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('Transaction', () => {
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';
const upsertSql = 'INSERTORUPDATE INTO FOO(BAR, BAZ) VALUES(@bar, @baz)';
spannerMock.putStatementResult(
selectSql,
mock.StatementResult.resultSet(createSelect1ResultSet())
);
spannerMock.putStatementResult(
updateSql,
mock.StatementResult.updateCount(1)
);
spannerMock.putStatementResult(
upsertSql,
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.observabilityOptions_ = {
tracerProvider: provider,
enableExtendedTracing: false,
};
});

beforeEach(() => {
spannerMock.resetRequests();
});

afterEach(() => {
traceExporter.reset();
});

it('run', done => {
database.getTransaction((err, tx) => {
assert.ifError(err);

tx!.run('SELECT 1', (err, rows) => {
traceExporter.forceFlush();

const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');

const actualSpanNames: string[] = [];
const actualEventNames: string[] = [];
spans.forEach(span => {
actualSpanNames.push(span.name);
span.events.forEach(event => {
actualEventNames.push(event.name);
});
});

const expectedSpanNames = [
'CloudSpanner.Database.getTransaction',
'CloudSpanner.Snapshot.run',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);

const expectedEventNames = ['Using Session'];
assert.strictEqual(
actualEventNames.every(value => expectedEventNames.includes(value)),
true,
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
);

done();
});
});
});
});
6 changes: 3 additions & 3 deletions src/batch-transaction.ts
Original file line number Diff line number Diff line change
Expand Up @@ -139,7 +139,7 @@ class BatchTransaction extends Snapshot {

const traceConfig: traceConfig = {
sql: query,
opts: this.observabilityOptions,
opts: this.observabilityOptions_,
};
return startTrace(
'BatchTransaction.createQueryPartitions',
Expand Down Expand Up @@ -182,7 +182,7 @@ class BatchTransaction extends Snapshot {
*/
createPartitions_(config, callback) {
const traceConfig: traceConfig = {
opts: this.observabilityOptions,
opts: this.observabilityOptions_,
};

return startTrace(
Expand Down Expand Up @@ -259,7 +259,7 @@ class BatchTransaction extends Snapshot {
*/
createReadPartitions(options, callback) {
const traceConfig: traceConfig = {
opts: this.observabilityOptions,
opts: this.observabilityOptions_,
};

return startTrace(
Expand Down
Loading

0 comments on commit 838225e

Please sign in to comment.