Skip to content

Commit

Permalink
feat: (observability): trace Database.runPartitionedUpdate (googleapi…
Browse files Browse the repository at this point in the history 
…s#2176)

This change traces Database.runPartitionedUpdate along with the appropriate tests for it with and without errors.

Updates googleapis#2079
  • Loading branch information
@odeke-em @surbhigarg92
odeke-em authored and surbhigarg92 committed Oct 29, 2024
1 parent 4e85d45 commit 20866a6
Showing 1 changed file with 44 additions and 127 deletions.
171 changes: 44 additions & 127 deletions observability-test/spanner.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,8 +53,6 @@ const {
} = require('@opentelemetry/context-async-hooks');

const {ObservabilityOptions} = require('../src/instrument');
import {SessionPool} from '../src/session-pool';

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

Expand Down Expand Up @@ -267,70 +265,18 @@ describe('EndToEnd', async () => {
});

it('run', async () => {
const withAllSpansHaveDBName = generateWithAllSpansHaveDBName(
database.formattedName_
);
const [rows] = await database.run('SELECT 1');

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
withAllSpansHaveDBName(spans);

// Sort the spans by duration.
spans.sort((spanA, spanB) => {
spanA.duration < spanB.duration;
});

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

const expectedSpanNames = [
'CloudSpanner.Snapshot.runStream',
'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.ok(
spanRun.spanContext().spanId,
'Expected that createQueryPartitions has a defined spanId'
);
assert.ok(
spanRunStream.spanContext().spanId,
'Expected that RunStream has a defined spanId'
);

const expectedEventNames = [
'Starting stream',
...cacheSessionEvents,
'Using Session',
];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
);
await verifySpansAndEvents(traceExporter, expectedSpanNames, expectedEventNames);
});

it('runTransaction', done => {
Expand Down Expand Up @@ -420,8 +366,6 @@ describe('EndToEnd', async () => {
expectedEventNames
);
done();
});
});

it('runPartitionedUpdate', async () => {
const [rowCount] = await database.runPartitionedUpdate({
Expand Down Expand Up @@ -469,82 +413,55 @@ describe('EndToEnd', async () => {
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
);
});
});
});

describe('SessionPool', async () => {
const traceExporter = new InMemorySpanExporter();
const sampler = new AlwaysOnSampler();
const provider = new NodeTracerProvider({
sampler: sampler,
exporter: traceExporter,
});
provider.addSpanProcessor(new SimpleSpanProcessor(traceExporter));

const setupResult = await setup({
tracerProvider: provider,
enableExtendedTracing: false,
});

const spanner = setupResult.spanner;
const server = setupResult.server;
const spannerMock = setupResult.spannerMock;
const instance = spanner.instance('instance');

after(async () => {
traceExporter.reset();
await provider.shutdown();
spannerMock.resetRequests();
spanner.close();
server.tryShutdown(() => {});
});

it('_createSessions', async () => {
// The first invocation of new SessionPool shall implicitly happen in here.
const database = instance.database('database');
await database.run('SELECT 1');

await provider.forceFlush();
traceExporter.reset();

// Explicitly invoking new SessionPool.
const sessionPool = new SessionPool(database);
});
});

const OPTIONS = 3;
await sessionPool._createSessions(OPTIONS);
it('runPartitionedUpdate', async () => {
const [rowCount] = await database.runPartitionedUpdate({
sql: updateSql,
});

traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();
await tracerProvider.forceFlush();
await traceExporter.forceFlush();
const spans = traceExporter.getFinishedSpans();

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

const expectedSpanNames = [
'CloudSpanner.Database.batchCreateSessions',
'CloudSpanner.SessionPool.createSessions',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);
const expectedSpanNames = [
'CloudSpanner.Snapshot.begin',
'CloudSpanner.Snapshot.runStream',
'CloudSpanner.Snapshot.run',
'CloudSpanner.Dml.runUpdate',
'CloudSpanner.PartitionedDml.runUpdate',
'CloudSpanner.Database.runPartitionedUpdate',
];
assert.deepStrictEqual(
actualSpanNames,
expectedSpanNames,
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
);

const expectedEventNames = [
'Requesting 3 sessions',
'Creating 3 sessions',
'Requested for 3 sessions returned 3',
];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
);
const expectedEventNames = [
'Begin Transaction',
'Transaction Creation Done',
'Starting stream',
'Acquiring session',
'Cache hit: has usable session',
'Acquired session',
];
assert.deepStrictEqual(
actualEventNames,
expectedEventNames,
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
);
});
});
});

Expand Down

0 comments on commit 20866a6

Please sign in to comment.