Skip to content

Commit

Permalink
Initialize baseline _traceConfig then extend it as needed
Browse files Browse the repository at this point in the history
  • Loading branch information
@odeke-em
odeke-em committed Oct 10, 2024
1 parent 8f5ea2b commit 0c902ad
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 66 deletions.
8 changes: 4 additions & 4 deletions observability-test/database.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -243,13 +243,13 @@ describe('Database', () => {
beforeEach(() => {
fakeCodec.encode = util.noop;
extend(Database, DatabaseCached);
database = new Database(INSTANCE, NAME, POOL_OPTIONS);
database.parent = INSTANCE;
database.databaseRole = 'parent_role';
database._observabilityOptions = {
INSTANCE._observabilityOptions = {
tracerProvider: provider,
enableExtendedTracing: false,
};
database = new Database(INSTANCE, NAME, POOL_OPTIONS);
database.parent = INSTANCE;
database.databaseRole = 'parent_role';
const gaxOpts = {};
const options: {
a: string;
Expand Down
111 changes: 49 additions & 62 deletions src/database.ts
View file
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,7 @@ import {
startTrace,
setSpanError,
setSpanErrorAndException,
traceConfig,
} from './instrument';

export type GetDatabaseRolesCallback = RequestCallback<
Expand Down Expand Up @@ -344,7 +345,8 @@ class Database extends common.GrpcServiceObject {
databaseDialect?: EnumKey<
typeof databaseAdmin.spanner.admin.database.v1.DatabaseDialect
> | null;
_observabilityOptions?: ObservabilityOptions;
_observabilityOptions?: ObservabilityOptions; // TODO: exmaine if we can remove it
private _traceConfig: traceConfig;
constructor(
instance: Instance,
name: string,
Expand Down Expand Up @@ -455,6 +457,12 @@ class Database extends common.GrpcServiceObject {
}
this.formattedName_ = formattedName_;
this.instance = instance;
this._observabilityOptions = instance._observabilityOptions;
this._traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};

this.resourceHeader_ = {
[CLOUD_RESOURCE_HEADER]: this.formattedName_,
};
Expand All @@ -467,7 +475,6 @@ class Database extends common.GrpcServiceObject {
Object.assign({}, queryOptions),
Database.getEnvironmentQueryOptions()
);
this._observabilityOptions = instance._observabilityOptions;
}
/**
* @typedef {array} SetDatabaseMetadataResponse
Expand Down Expand Up @@ -693,7 +700,7 @@ class Database extends common.GrpcServiceObject {

const sessions = (resp!.session || []).map(metadata => {
const session = this.session(metadata.name!);
session._observabilityOptions = this._observabilityOptions;
session._observabilityOptions = this._traceConfig!.opts;
session.metadata = metadata;
return session;
});
Expand Down Expand Up @@ -738,7 +745,7 @@ class Database extends common.GrpcServiceObject {
const id = identifier.transaction;
const transaction = new BatchTransaction(session, options);
transaction.id = id;
transaction._observabilityOptions = this._observabilityOptions;
transaction._observabilityOptions = this._traceConfig!.opts;
transaction.readTimestamp = identifier.timestamp as PreciseDate;
return transaction;
}
Expand Down Expand Up @@ -827,39 +834,41 @@ class Database extends common.GrpcServiceObject {
typeof optionsOrCallback === 'object'
? (optionsOrCallback as TimestampBounds)
: {};

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
return startTrace('Database.createBatchTransaction', traceConfig, span => {
this.pool_.getSession((err, session) => {
if (err) {
setSpanError(span, err);
span.end();
callback!(err as ServiceError, null, undefined);
return;
}
const transaction = this.batchTransaction({session: session!}, options);
this._releaseOnEnd(session!, transaction, span);
transaction.begin((err, resp) => {
return startTrace(
'Database.createBatchTransaction',
this._traceConfig,
span => {
this.pool_.getSession((err, session) => {
if (err) {
setSpanError(span, err);
if (isSessionNotFoundError(err)) {
span.addEvent('No session available', {
'session.id': session?.id,
});
}
span.end();
callback!(err, null, resp!);
callback!(err as ServiceError, null, undefined);
return;
}
span.addEvent('Using Session', {'session.id': session?.id});
span.end();
callback!(null, transaction, resp!);
const transaction = this.batchTransaction(
{session: session!},
options
);
this._releaseOnEnd(session!, transaction, span);
transaction.begin((err, resp) => {
if (err) {
setSpanError(span, err);
if (isSessionNotFoundError(err)) {
span.addEvent('No session available', {
'session.id': session?.id,
});
}
span.end();
callback!(err, null, resp!);
return;
}
span.addEvent('Using Session', {'session.id': session?.id});
span.end();
callback!(null, transaction, resp!);
});
});
});
});
}
);
}
/**
* Create a new session.
Expand Down Expand Up @@ -1089,7 +1098,7 @@ class Database extends common.GrpcServiceObject {
/CREATE TABLE `*([^\s`(]+)/
)![1];
const table = this.table(tableName!);
table._observabilityOptions = this._observabilityOptions;
table._observabilityOptions = this._traceConfig!.opts;
callback!(null, table, operation!, resp!);
});
}
Expand Down Expand Up @@ -1878,11 +1887,7 @@ class Database extends common.GrpcServiceObject {
delete (gaxOpts as GetSessionsOptions).pageToken;
}

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
return startTrace('Database.getSessions', traceConfig, span => {
return startTrace('Database.getSessions', this._traceConfig, span => {
this.request<
google.spanner.v1.ISession,
google.spanner.v1.IListSessionsResponse
Expand All @@ -1903,7 +1908,7 @@ class Database extends common.GrpcServiceObject {
sessionInstances = sessions.map(metadata => {
const session = self.session(metadata.name!);
session.metadata = metadata;
session._observabilityOptions = this._observabilityOptions;
session._observabilityOptions = this._traceConfig!.opts;
return session;
});
}
Expand Down Expand Up @@ -2064,11 +2069,7 @@ class Database extends common.GrpcServiceObject {
? (optionsOrCallback as TimestampBounds)
: {};

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
return startTrace('Database.getSnapshot', traceConfig, span => {
return startTrace('Database.getSnapshot', this._traceConfig, span => {
this.pool_.getSession((err, session) => {
if (err) {
setSpanError(span, err);
Expand Down Expand Up @@ -2168,11 +2169,7 @@ class Database extends common.GrpcServiceObject {
? (optionsOrCallback as GetTransactionOptions)
: {};

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
return startTrace('Database.getTransaction', traceConfig, span => {
return startTrace('Database.getTransaction', this._traceConfig, span => {
this.pool_.getSession((err, session, transaction) => {
if (options.requestOptions) {
transaction!.requestOptions = Object.assign(
Expand Down Expand Up @@ -2800,8 +2797,7 @@ class Database extends common.GrpcServiceObject {

const traceConfig = {
sql: query,
opts: this._observabilityOptions,
dbName: this.formattedName_,
...this._traceConfig,
};
return startTrace('Database.run', traceConfig, span => {
this.runStream(query, options)
Expand Down Expand Up @@ -3025,8 +3021,7 @@ class Database extends common.GrpcServiceObject {
const proxyStream: Transform = through.obj();
const traceConfig = {
sql: query,
opts: this._observabilityOptions,
dbName: this.formattedName_,
...this._traceConfig,
};
return startTrace('Database.runStream', traceConfig, span => {
this.pool_.getSession((err, session) => {
Expand Down Expand Up @@ -3205,11 +3200,7 @@ class Database extends common.GrpcServiceObject {
? (optionsOrRunFn as RunTransactionOptions)
: {};

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
startTrace('Database.runTransaction', traceConfig, span => {
startTrace('Database.runTransaction', this._traceConfig, span => {
this.pool_.getSession((err, session?, transaction?) => {
if (err) {
setSpanError(span, err);
Expand Down Expand Up @@ -3601,11 +3592,7 @@ class Database extends common.GrpcServiceObject {
? (optionsOrCallback as CallOptions)
: {};

const traceConfig = {
opts: this._observabilityOptions,
dbName: this.formattedName_,
};
return startTrace('Database.writeAtLeastOnce', traceConfig, span => {
return startTrace('Database.writeAtLeastOnce', this._traceConfig, span => {
this.pool_.getSession((err, session?, transaction?) => {
if (err && isSessionNotFoundError(err as grpc.ServiceError)) {
span.addEvent('No session available', {
Expand Down

0 comments on commit 0c902ad

Please sign in to comment.