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 googleapis#2079
Built from PR googleapis#2087
Updates googleapis#2114
  • Loading branch information
odeke-em committed Sep 22, 2024
1 parent d51aae9 commit b8d0789
Show file tree
Hide file tree
Showing 2 changed files with 763 additions and 481 deletions.
196 changes: 196 additions & 0 deletions observability-test/transaction.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,196 @@
/*!
* 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 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.observabilityConfig = {
tracerProvider: provider,
enableExtendedTracing: false,
};
});

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

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

it('run', done => {
database.runTransaction((err, tx) => {
assert.ifError(err);
assert.strictEqual(!tx, false, 'Expecting a non-null transaction');
console.log('db.runTransaction');
// Reset the trace exporter to clear any prior spans
// so that we only have spans for Transaction.
// traceExporter.reset();

console.log('ts.run', tx);
console.log('fini');
tx!.run('SELECT 1', (err, rows) => {
console.log('tx!.run');
traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
assert.strictEqual(spans.length, 2, 'Exactly 2 spans expected');

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

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

done();
});
});
});
});
Loading

0 comments on commit b8d0789

Please sign in to comment.