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 30, 2024
1 parent 5b70b70 commit 4fc8fe7
Show file tree
Hide file tree
Showing 5 changed files with 602 additions and 291 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
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) => {
// Reset the trace exporter to clear any prior spans
// so that we only have spans for Transaction.
traceExporter.forceFlush();
traceExporter.reset();

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.Transaction.upsert'];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);

const expectedEventNames = [];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
`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
6 changes: 5 additions & 1 deletion src/database.ts
Original file line number Diff line number Diff line change
Expand Up @@ -102,7 +102,7 @@ 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} from './instrument';

export type GetDatabaseRolesCallback = RequestCallback<
IDatabaseRole,
Expand Down Expand Up @@ -337,6 +337,7 @@ class Database extends common.GrpcServiceObject {
databaseDialect?: EnumKey<
typeof databaseAdmin.spanner.admin.database.v1.DatabaseDialect
> | null;
observabilityOptions_?: ObservabilityOptions;
constructor(
instance: Instance,
name: string,
Expand Down Expand Up @@ -729,6 +730,7 @@ class Database extends common.GrpcServiceObject {
const transaction = new BatchTransaction(session, options);
transaction.id = id;
transaction.readTimestamp = identifier.timestamp as PreciseDate;
transaction.observabilityOptions_ = this.observabilityOptions_;
return transaction;
}
/**
Expand Down Expand Up @@ -822,6 +824,7 @@ class Database extends common.GrpcServiceObject {
callback!(err as ServiceError, null, undefined);
return;
}
session!.observabilityOptions_ = this.observabilityOptions_;
const transaction = this.batchTransaction({session: session!}, options);
this._releaseOnEnd(session!, transaction);
transaction.begin((err, resp) => {
Expand Down Expand Up @@ -3129,6 +3132,7 @@ class Database extends common.GrpcServiceObject {
runFn!(err as grpc.ServiceError);
return;
}
transaction!.observabilityOptions_ = this.observabilityOptions_;
if (options.optimisticLock) {
transaction!.useOptimisticLock();
}
Expand Down
Loading

0 comments on commit 4fc8fe7

Please sign in to comment.