Skip to content

Commit 701e226

Browse files
odeke-emodeke-em
authored
feat: (observability): trace Database.runPartitionedUpdate (#2176)
This change traces Database.runPartitionedUpdate along with the appropriate tests for it with and without errors. Updates #2079
1 parent cbc86fa commit 701e226

File tree

3 files changed

+287
-7
lines changed

3 files changed

+287
-7
lines changed

‎observability-test/database.ts

+220-1
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,10 @@ const {
3838
InMemorySpanExporter,
3939
} = require('@opentelemetry/sdk-trace-node');
4040
// eslint-disable-next-line n/no-extraneous-require
41-
const {SimpleSpanProcessor} = require('@opentelemetry/sdk-trace-base');
41+
const {
42+
ReadableSpan,
43+
SimpleSpanProcessor,
44+
} = require('@opentelemetry/sdk-trace-base');
4245
import * as db from '../src/database';
4346
import {Instance, MutationGroup, Spanner} from '../src';
4447
import * as pfy from '@google-cloud/promisify';
@@ -1954,4 +1957,220 @@ describe('Database', () => {
19541957
fakeStream2.push(null);
19551958
});
19561959
});
1960+
1961+
describe('runPartitionedUpdate', () => {
1962+
const QUERY = {
1963+
sql: 'INSERT INTO `MyTable` (Key, Thing) VALUES(@key, @thing)',
1964+
params: {
1965+
key: 'k999',
1966+
thing: 'abc',
1967+
},
1968+
};
1969+
1970+
let fakePool: FakeSessionPool;
1971+
let fakeSession: FakeSession;
1972+
let fakePartitionedDml = new FakeTransaction(
1973+
{} as google.spanner.v1.TransactionOptions.PartitionedDml
1974+
);
1975+
1976+
let getSessionStub;
1977+
let beginStub;
1978+
let runUpdateStub;
1979+
1980+
beforeEach(() => {
1981+
fakePool = database.pool_;
1982+
fakeSession = new FakeSession();
1983+
fakePartitionedDml = new FakeTransaction(
1984+
{} as google.spanner.v1.TransactionOptions.PartitionedDml
1985+
);
1986+
1987+
getSessionStub = (
1988+
sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
1989+
).callsFake(callback => {
1990+
callback(null, fakeSession);
1991+
});
1992+
1993+
sandbox.stub(fakeSession, 'partitionedDml').returns(fakePartitionedDml);
1994+
1995+
beginStub = (
1996+
sandbox.stub(fakePartitionedDml, 'begin') as sinon.SinonStub
1997+
).callsFake(callback => callback(null));
1998+
1999+
runUpdateStub = (
2000+
sandbox.stub(fakePartitionedDml, 'runUpdate') as sinon.SinonStub
2001+
).callsFake((_, callback) => callback(null));
2002+
});
2003+
2004+
interface traceExportResults {
2005+
spanNames: string[];
2006+
spans: (typeof ReadableSpan)[];
2007+
eventNames: string[];
2008+
}
2009+
2010+
async function getTraceExportResults(): Promise<traceExportResults> {
2011+
await provider.forceFlush();
2012+
await traceExporter.forceFlush();
2013+
const spans = traceExporter.getFinishedSpans();
2014+
withAllSpansHaveDBName(spans);
2015+
2016+
const actualSpanNames: string[] = [];
2017+
const actualEventNames: string[] = [];
2018+
spans.forEach(span => {
2019+
actualSpanNames.push(span.name);
2020+
span.events.forEach(event => {
2021+
actualEventNames.push(event.name);
2022+
});
2023+
});
2024+
2025+
return Promise.resolve({
2026+
spanNames: actualSpanNames,
2027+
spans: spans,
2028+
eventNames: actualEventNames,
2029+
});
2030+
}
2031+
2032+
it('with pool errors', done => {
2033+
const fakeError = new Error('err');
2034+
const fakeCallback = sandbox.spy();
2035+
2036+
getSessionStub.callsFake(callback => callback(fakeError));
2037+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2038+
assert.strictEqual(err, fakeError);
2039+
assert.strictEqual(rowCount, 0);
2040+
2041+
const exportResults = await getTraceExportResults();
2042+
const actualSpanNames = exportResults.spanNames;
2043+
const spans = exportResults.spans;
2044+
const actualEventNames = exportResults.eventNames;
2045+
2046+
const expectedSpanNames = [
2047+
'CloudSpanner.Database.runPartitionedUpdate',
2048+
];
2049+
assert.deepStrictEqual(
2050+
actualSpanNames,
2051+
expectedSpanNames,
2052+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2053+
);
2054+
2055+
// Ensure that the first span actually produced an error that was recorded.
2056+
const parentSpan = spans[0];
2057+
assert.deepStrictEqual(
2058+
SpanStatusCode.ERROR,
2059+
parentSpan.status.code,
2060+
'Expected an ERROR span status'
2061+
);
2062+
assert.deepStrictEqual(
2063+
fakeError.message,
2064+
parentSpan.status.message.toString(),
2065+
'Mismatched span status message'
2066+
);
2067+
2068+
const expectedEventNames = [];
2069+
assert.deepStrictEqual(
2070+
actualEventNames,
2071+
expectedEventNames,
2072+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2073+
);
2074+
2075+
done();
2076+
});
2077+
});
2078+
2079+
it('with begin errors', done => {
2080+
const fakeError = new Error('err');
2081+
2082+
beginStub.callsFake(callback => callback(fakeError));
2083+
2084+
const releaseStub = (
2085+
sandbox.stub(fakePool, 'release') as sinon.SinonStub
2086+
).withArgs(fakeSession);
2087+
2088+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2089+
assert.strictEqual(err, fakeError);
2090+
assert.strictEqual(rowCount, 0);
2091+
assert.strictEqual(releaseStub.callCount, 1);
2092+
2093+
const exportResults = await getTraceExportResults();
2094+
const actualSpanNames = exportResults.spanNames;
2095+
const spans = exportResults.spans;
2096+
const actualEventNames = exportResults.eventNames;
2097+
2098+
const expectedSpanNames = [
2099+
'CloudSpanner.Database.runPartitionedUpdate',
2100+
];
2101+
assert.deepStrictEqual(
2102+
actualSpanNames,
2103+
expectedSpanNames,
2104+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2105+
);
2106+
2107+
// Ensure that the first span actually produced an error that was recorded.
2108+
const parentSpan = spans[0];
2109+
assert.deepStrictEqual(
2110+
SpanStatusCode.ERROR,
2111+
parentSpan.status.code,
2112+
'Expected an ERROR span status'
2113+
);
2114+
assert.deepStrictEqual(
2115+
fakeError.message,
2116+
parentSpan.status.message.toString(),
2117+
'Mismatched span status message'
2118+
);
2119+
2120+
const expectedEventNames = [];
2121+
assert.deepStrictEqual(
2122+
actualEventNames,
2123+
expectedEventNames,
2124+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2125+
);
2126+
done();
2127+
});
2128+
});
2129+
2130+
it('session released on transaction end', done => {
2131+
const releaseStub = (
2132+
sandbox.stub(fakePool, 'release') as sinon.SinonStub
2133+
).withArgs(fakeSession);
2134+
2135+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2136+
const exportResults = await getTraceExportResults();
2137+
const actualSpanNames = exportResults.spanNames;
2138+
const spans = exportResults.spans;
2139+
const actualEventNames = exportResults.eventNames;
2140+
2141+
const expectedSpanNames = [
2142+
'CloudSpanner.Database.runPartitionedUpdate',
2143+
];
2144+
assert.deepStrictEqual(
2145+
actualSpanNames,
2146+
expectedSpanNames,
2147+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2148+
);
2149+
2150+
// Ensure that the first span actually produced an error that was recorded.
2151+
const parentSpan = spans[0];
2152+
assert.deepStrictEqual(
2153+
SpanStatusCode.UNSET,
2154+
parentSpan.status.code,
2155+
'Unexpected span status'
2156+
);
2157+
assert.deepStrictEqual(
2158+
undefined,
2159+
parentSpan.status.message,
2160+
'Mismatched span status message'
2161+
);
2162+
2163+
const expectedEventNames = [];
2164+
assert.deepStrictEqual(
2165+
actualEventNames,
2166+
expectedEventNames,
2167+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2168+
);
2169+
done();
2170+
});
2171+
2172+
fakePartitionedDml.emit('end');
2173+
assert.strictEqual(releaseStub.callCount, 1);
2174+
});
2175+
});
19572176
});

‎observability-test/spanner.ts

+47
Original file line numberDiff line numberDiff line change
@@ -613,6 +613,53 @@ describe('EndToEnd', async () => {
613613
done();
614614
});
615615
});
616+
617+
it('runPartitionedUpdate', async () => {
618+
const [rowCount] = await database.runPartitionedUpdate({
619+
sql: updateSql,
620+
});
621+
622+
await tracerProvider.forceFlush();
623+
await traceExporter.forceFlush();
624+
const spans = traceExporter.getFinishedSpans();
625+
626+
const actualEventNames: string[] = [];
627+
const actualSpanNames: string[] = [];
628+
spans.forEach(span => {
629+
actualSpanNames.push(span.name);
630+
span.events.forEach(event => {
631+
actualEventNames.push(event.name);
632+
});
633+
});
634+
635+
const expectedSpanNames = [
636+
'CloudSpanner.Snapshot.begin',
637+
'CloudSpanner.Snapshot.runStream',
638+
'CloudSpanner.Snapshot.run',
639+
'CloudSpanner.Dml.runUpdate',
640+
'CloudSpanner.PartitionedDml.runUpdate',
641+
'CloudSpanner.Database.runPartitionedUpdate',
642+
];
643+
assert.deepStrictEqual(
644+
actualSpanNames,
645+
expectedSpanNames,
646+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
647+
);
648+
649+
const expectedEventNames = [
650+
'Begin Transaction',
651+
'Transaction Creation Done',
652+
'Starting stream',
653+
'Acquiring session',
654+
'Cache hit: has usable session',
655+
'Acquired session',
656+
];
657+
assert.deepStrictEqual(
658+
actualEventNames,
659+
expectedEventNames,
660+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
661+
);
662+
});
616663
});
617664
});
618665

‎src/database.ts

+20-6
Original file line numberDiff line numberDiff line change
@@ -2858,13 +2858,27 @@ class Database extends common.GrpcServiceObject {
28582858
query: string | RunPartitionedUpdateOptions,
28592859
callback?: RunUpdateCallback
28602860
): void | Promise<[number]> {
2861-
this.pool_.getSession((err, session) => {
2862-
if (err) {
2863-
callback!(err as ServiceError, 0);
2864-
return;
2865-
}
2861+
const traceConfig = {
2862+
sql: query,
2863+
...this._traceConfig,
2864+
};
2865+
return startTrace('Database.runPartitionedUpdate', traceConfig, span => {
2866+
this.pool_.getSession((err, session) => {
2867+
if (err) {
2868+
setSpanError(span, err);
2869+
span.end();
2870+
callback!(err as ServiceError, 0);
2871+
return;
2872+
}
28662873

2867-
this._runPartitionedUpdate(session!, query, callback);
2874+
this._runPartitionedUpdate(session!, query, (err, count) => {
2875+
if (err) {
2876+
setSpanError(span, err);
2877+
}
2878+
span.end();
2879+
callback!(err, count);
2880+
});
2881+
});
28682882
});
28692883
}
28702884

0 commit comments

Comments
 (0)