Skip to content

Commit a8c1982

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

File tree

3 files changed

+300
-7
lines changed

3 files changed

+300
-7
lines changed

‎observability-test/database.ts

+233-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,233 @@ 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+
const fakeDirectedReadOptions = {
1981+
includeReplicas: {
1982+
replicaSelections: [
1983+
{
1984+
location: 'us-west1',
1985+
type: google.spanner.v1.DirectedReadOptions.ReplicaSelection.Type
1986+
.READ_WRITE,
1987+
},
1988+
],
1989+
autoFailoverDisabled: true,
1990+
},
1991+
};
1992+
1993+
beforeEach(() => {
1994+
fakePool = database.pool_;
1995+
fakeSession = new FakeSession();
1996+
fakePartitionedDml = new FakeTransaction(
1997+
{} as google.spanner.v1.TransactionOptions.PartitionedDml
1998+
);
1999+
2000+
getSessionStub = (
2001+
sandbox.stub(fakePool, 'getSession') as sinon.SinonStub
2002+
).callsFake(callback => {
2003+
callback(null, fakeSession);
2004+
});
2005+
2006+
sandbox.stub(fakeSession, 'partitionedDml').returns(fakePartitionedDml);
2007+
2008+
beginStub = (
2009+
sandbox.stub(fakePartitionedDml, 'begin') as sinon.SinonStub
2010+
).callsFake(callback => callback(null));
2011+
2012+
runUpdateStub = (
2013+
sandbox.stub(fakePartitionedDml, 'runUpdate') as sinon.SinonStub
2014+
).callsFake((_, callback) => callback(null));
2015+
});
2016+
2017+
interface traceExportResults {
2018+
spanNames: string[];
2019+
spans: (typeof ReadableSpan)[];
2020+
eventNames: string[];
2021+
}
2022+
2023+
async function getTraceExportResults(): Promise<traceExportResults> {
2024+
await provider.forceFlush();
2025+
await traceExporter.forceFlush();
2026+
const spans = traceExporter.getFinishedSpans();
2027+
withAllSpansHaveDBName(spans);
2028+
2029+
const actualSpanNames: string[] = [];
2030+
const actualEventNames: string[] = [];
2031+
spans.forEach(span => {
2032+
actualSpanNames.push(span.name);
2033+
span.events.forEach(event => {
2034+
actualEventNames.push(event.name);
2035+
});
2036+
});
2037+
2038+
return Promise.resolve({
2039+
spanNames: actualSpanNames,
2040+
spans: spans,
2041+
eventNames: actualEventNames,
2042+
});
2043+
}
2044+
2045+
it('with pool errors', done => {
2046+
const fakeError = new Error('err');
2047+
const fakeCallback = sandbox.spy();
2048+
2049+
getSessionStub.callsFake(callback => callback(fakeError));
2050+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2051+
assert.strictEqual(err, fakeError);
2052+
assert.strictEqual(rowCount, 0);
2053+
2054+
const exportResults = await getTraceExportResults();
2055+
const actualSpanNames = exportResults.spanNames;
2056+
const spans = exportResults.spans;
2057+
const actualEventNames = exportResults.eventNames;
2058+
2059+
const expectedSpanNames = [
2060+
'CloudSpanner.Database.runPartitionedUpdate',
2061+
];
2062+
assert.deepStrictEqual(
2063+
actualSpanNames,
2064+
expectedSpanNames,
2065+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2066+
);
2067+
2068+
// Ensure that the first span actually produced an error that was recorded.
2069+
const parentSpan = spans[0];
2070+
assert.deepStrictEqual(
2071+
SpanStatusCode.ERROR,
2072+
parentSpan.status.code,
2073+
'Expected an ERROR span status'
2074+
);
2075+
assert.deepStrictEqual(
2076+
fakeError.message,
2077+
parentSpan.status.message.toString(),
2078+
'Mismatched span status message'
2079+
);
2080+
2081+
const expectedEventNames = [];
2082+
assert.deepStrictEqual(
2083+
actualEventNames,
2084+
expectedEventNames,
2085+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2086+
);
2087+
2088+
done();
2089+
});
2090+
});
2091+
2092+
it('with begin errors', done => {
2093+
const fakeError = new Error('err');
2094+
2095+
beginStub.callsFake(callback => callback(fakeError));
2096+
2097+
const releaseStub = (
2098+
sandbox.stub(fakePool, 'release') as sinon.SinonStub
2099+
).withArgs(fakeSession);
2100+
2101+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2102+
assert.strictEqual(err, fakeError);
2103+
assert.strictEqual(rowCount, 0);
2104+
assert.strictEqual(releaseStub.callCount, 1);
2105+
2106+
const exportResults = await getTraceExportResults();
2107+
const actualSpanNames = exportResults.spanNames;
2108+
const spans = exportResults.spans;
2109+
const actualEventNames = exportResults.eventNames;
2110+
2111+
const expectedSpanNames = [
2112+
'CloudSpanner.Database.runPartitionedUpdate',
2113+
];
2114+
assert.deepStrictEqual(
2115+
actualSpanNames,
2116+
expectedSpanNames,
2117+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2118+
);
2119+
2120+
// Ensure that the first span actually produced an error that was recorded.
2121+
const parentSpan = spans[0];
2122+
assert.deepStrictEqual(
2123+
SpanStatusCode.ERROR,
2124+
parentSpan.status.code,
2125+
'Expected an ERROR span status'
2126+
);
2127+
assert.deepStrictEqual(
2128+
fakeError.message,
2129+
parentSpan.status.message.toString(),
2130+
'Mismatched span status message'
2131+
);
2132+
2133+
const expectedEventNames = [];
2134+
assert.deepStrictEqual(
2135+
actualEventNames,
2136+
expectedEventNames,
2137+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2138+
);
2139+
done();
2140+
});
2141+
});
2142+
2143+
it('session released on transaction end', done => {
2144+
const releaseStub = (
2145+
sandbox.stub(fakePool, 'release') as sinon.SinonStub
2146+
).withArgs(fakeSession);
2147+
2148+
database.runPartitionedUpdate(QUERY, async (err, rowCount) => {
2149+
const exportResults = await getTraceExportResults();
2150+
const actualSpanNames = exportResults.spanNames;
2151+
const spans = exportResults.spans;
2152+
const actualEventNames = exportResults.eventNames;
2153+
2154+
const expectedSpanNames = [
2155+
'CloudSpanner.Database.runPartitionedUpdate',
2156+
];
2157+
assert.deepStrictEqual(
2158+
actualSpanNames,
2159+
expectedSpanNames,
2160+
`span names mismatch:\n\tGot: ${actualSpanNames}\n\tWant: ${expectedSpanNames}`
2161+
);
2162+
2163+
// Ensure that the first span actually produced an error that was recorded.
2164+
const parentSpan = spans[0];
2165+
assert.deepStrictEqual(
2166+
SpanStatusCode.UNSET,
2167+
parentSpan.status.code,
2168+
'Unexpected span status'
2169+
);
2170+
assert.deepStrictEqual(
2171+
undefined,
2172+
parentSpan.status.message,
2173+
'Mismatched span status message'
2174+
);
2175+
2176+
const expectedEventNames = [];
2177+
assert.deepStrictEqual(
2178+
actualEventNames,
2179+
expectedEventNames,
2180+
`Unexpected events:\n\tGot: ${actualEventNames}\n\tWant: ${expectedEventNames}`
2181+
);
2182+
done();
2183+
});
2184+
2185+
fakePartitionedDml.emit('end');
2186+
assert.strictEqual(releaseStub.callCount, 1);
2187+
});
2188+
});
19572189
});

‎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)