Skip to content

Commit 2e4dd94

Browse files
pcncTrigger.dev RepoOps
pcnc
authored and
Trigger.dev RepoOps
committed
feat(webapp): label database spans with the Prisma client they use
Database query spans from the webapp now carry a db.client_type attribute naming the database client that ran the query, the same name used on the webapp's database metrics. Query latency can now be broken down by query type (INSERT, UPDATE, SELECT, ...) for each client. Mono-RevId: 6dead141fe60243a6f6b7bb9ea4847920dd288fb
1 parent 1ca7cfd commit 2e4dd94

4 files changed

Lines changed: 260 additions & 72 deletions

File tree

‎apps/webapp/app/db.server.ts‎

Lines changed: 21 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -47,9 +47,9 @@ import {
4747
runOpsTransactionResilience,
4848
} from "./v3/transactionResilience.server";
4949
import { assertControlPlaneCoresidencyAdvisory } from "./v3/runOpsMigration/controlPlaneCoresidencySentinel.server";
50-
import { DATASOURCE_CONTEXT_KEY, startActiveSpan } from "./v3/tracer.server";
50+
import { startActiveSpan } from "./v3/tracer.server";
51+
import { tagDatasource, tagDatasourceRunOps } from "./utils/datasourceTagging.server";
5152
import type { Span } from "@opentelemetry/api";
52-
import { context, trace } from "@opentelemetry/api";
5353
import { queryPerformanceMonitor } from "./utils/queryPerformanceMonitor.server";
5454

5555
export type {
@@ -164,53 +164,6 @@ async function $transactionInner<R>(
164164

165165
export { Prisma, boundedIn };
166166

167-
type DatasourceLabel =
168-
| "control-plane-writer"
169-
| "control-plane-replica"
170-
| "legacy-run-ops-writer"
171-
| "legacy-run-ops-replica"
172-
| "run-ops-writer"
173-
| "run-ops-replica"
174-
| "webhook-writer"
175-
| "webhook-partitions"
176-
| "webhook-replica"
177-
| "alert-partitions";
178-
179-
function tagDatasource<T extends PrismaClient>(datasource: DatasourceLabel, client: T): T {
180-
return client.$extends({
181-
name: "datasource-tagger",
182-
query: {
183-
$allOperations: ({ query, args }) => {
184-
trace.getActiveSpan()?.setAttribute("db.datasource", datasource);
185-
return context.with(
186-
context.active().setValue(DATASOURCE_CONTEXT_KEY, datasource),
187-
async () => await query(args)
188-
);
189-
},
190-
},
191-
}) as unknown as T;
192-
}
193-
194-
// Same extension as tagDatasource but typed for RunOpsPrismaClient (different
195-
// generated package — does not extend @trigger.dev/database.PrismaClient).
196-
function tagDatasourceRunOps(
197-
datasource: DatasourceLabel,
198-
client: RunOpsPrismaClient
199-
): RunOpsPrismaClient {
200-
return client.$extends({
201-
name: "datasource-tagger",
202-
query: {
203-
$allOperations: ({ query, args }) => {
204-
trace.getActiveSpan()?.setAttribute("db.datasource", datasource);
205-
return context.with(
206-
context.active().setValue(DATASOURCE_CONTEXT_KEY, datasource),
207-
async () => await query(args)
208-
);
209-
},
210-
},
211-
}) as unknown as RunOpsPrismaClient;
212-
}
213-
214167
// Same wrapper as captureInfrastructureErrors, bridged via double cast because
215168
// that helper is constrained to T extends @trigger.dev/database.PrismaClient.
216169
function captureInfraErrorsRunOps(client: RunOpsPrismaClient): RunOpsPrismaClient {
@@ -221,7 +174,9 @@ function captureInfraErrorsRunOps(client: RunOpsPrismaClient): RunOpsPrismaClien
221174

222175
export const prisma = singleton("prisma", () =>
223176
registerTransactionResilience(
224-
captureInfrastructureErrors(tagDatasource("control-plane-writer", getClient())),
177+
captureInfrastructureErrors(
178+
tagDatasource("control-plane-writer", "control-plane-writer", getClient())
179+
),
225180
controlPlaneTransactionResilience
226181
)
227182
);
@@ -232,7 +187,9 @@ export const $replica: PrismaReplicaClient = singleton("replica", () => {
232187
// No replica configured → fall back to the writer `prisma`, which must stay UNBRANDED.
233188
return replica
234189
? markReadReplicaClient(
235-
captureInfrastructureErrors(tagDatasource("control-plane-replica", replica))
190+
captureInfrastructureErrors(
191+
tagDatasource("control-plane-replica", "control-plane-replica", replica)
192+
)
236193
)
237194
: prisma;
238195
});
@@ -248,6 +205,7 @@ export const webhookPrisma: WebhookDatabase = singleton("webhookPrisma", () => {
248205
}
249206
return captureInfrastructureErrors(
250207
tagDatasource(
208+
"webhook-writer",
251209
"webhook-writer",
252210
buildWriterClient({
253211
url: env.WEBHOOK_DATABASE_URL,
@@ -269,6 +227,7 @@ export const webhookPartitionPrisma: WebhookDatabase = singleton("webhookPartiti
269227
}
270228
return captureInfrastructureErrors(
271229
tagDatasource(
230+
"webhook-partitions",
272231
"webhook-partitions",
273232
buildWriterClient({
274233
url: env.WEBHOOK_DATABASE_DIRECT_URL,
@@ -283,6 +242,7 @@ export const webhookPartitionPrisma: WebhookDatabase = singleton("webhookPartiti
283242
export function createAlertPartitionPrisma(): PrismaClient {
284243
return captureInfrastructureErrors(
285244
tagDatasource(
245+
"alert-partitions",
286246
"alert-partitions",
287247
buildWriterClient({
288248
url: env.DATABASE_URL,
@@ -304,6 +264,7 @@ export const webhookReplica: WebhookReplicaDatabase = singleton("webhookReplica"
304264
captureInfrastructureErrors(
305265
tagDatasource(
306266
"webhook-replica",
267+
"webhook-reader",
307268
buildReplicaClient({
308269
url: env.WEBHOOK_DATABASE_READ_REPLICA_URL,
309270
clientType: "webhook-reader",
@@ -429,7 +390,7 @@ export function selectRunOpsTopology(
429390
// assertRunOpsSplitSentinel(). Because the builder callbacks only run when
430391
// splitEnabled is true, single-DB reuses prisma/$replica by reference and opens
431392
// nothing new. The builders apply the SAME wrapper pair the control-plane
432-
// singletons use (captureInfrastructureErrors(tagDatasource(role, raw))).
393+
// singletons use (captureInfrastructureErrors(tagDatasource(datasource, clientType, raw))).
433394
const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
434395
const newUrl = env.RUN_OPS_DATABASE_URL;
435396
// Gate on the opt-in flag too: the distinct-DB sentinel only runs when the flag is on.
@@ -486,6 +447,7 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
486447
captureInfraErrorsRunOps(
487448
tagDatasourceRunOps(
488449
"run-ops-writer",
450+
clientType,
489451
buildRunOpsClient({
490452
url,
491453
clientType,
@@ -507,6 +469,7 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
507469
captureInfraErrorsRunOps(
508470
tagDatasourceRunOps(
509471
"run-ops-replica",
472+
clientType,
510473
buildRunOpsClient({
511474
url,
512475
clientType,
@@ -526,6 +489,7 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
526489
captureInfrastructureErrors(
527490
tagDatasource(
528491
"legacy-run-ops-writer",
492+
clientType,
529493
buildWriterClient({
530494
url,
531495
clientType,
@@ -542,6 +506,7 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
542506
captureInfrastructureErrors(
543507
tagDatasource(
544508
"legacy-run-ops-replica",
509+
clientType,
545510
buildReplicaClient({
546511
url,
547512
clientType,
@@ -554,14 +519,16 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
554519
),
555520
// A gen-2 shard is a dedicated run-ops DB, so it mirrors buildNewWriter/buildNewReplica: same
556521
// client class, same wrapper stack, its OWN resilience budget, and the "new"-role pool knobs
557-
// merged with the descriptor's per-shard overrides. Shards share the run-ops datasource tag.
522+
// merged with the descriptor's per-shard overrides. Shards share the run-ops datasource tag;
523+
// the db.client_type span attribute (the clientType) tells them apart.
558524
buildShardWriter: (shard) => {
559525
const descriptor = shardDescriptorsByKey.get(shard.key);
560526
const knobs = resolveRunOpsPoolKnobs("new", descriptor?.knobs);
561527
return registerTransactionResilience(
562528
captureInfraErrorsRunOps(
563529
tagDatasourceRunOps(
564530
"run-ops-writer",
531+
`run-ops-shard-${shard.key}-writer`,
565532
buildRunOpsClient({
566533
url: shard.url!,
567534
clientType: `run-ops-shard-${shard.key}-writer`,
@@ -583,6 +550,7 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => {
583550
captureInfraErrorsRunOps(
584551
tagDatasourceRunOps(
585552
"run-ops-replica",
553+
`run-ops-shard-${shard.key}-replica`,
586554
buildRunOpsClient({
587555
url: shard.replicaUrl!,
588556
clientType: `run-ops-shard-${shard.key}-replica`,

‎apps/webapp/app/utils/datasourceTagging.server.ts‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import { type Context, type Span, context, createContextKey, trace } from "@opentelemetry/api";
2+
import type { SpanProcessor } from "@opentelemetry/sdk-trace-base";
3+
import type { RunOpsPrismaClient } from "@internal/run-ops-database";
4+
import type { PrismaClient } from "@trigger.dev/database";
5+
6+
const DATASOURCE_CONTEXT_KEY = createContextKey("trigger.db.datasource");
7+
const CLIENT_TYPE_CONTEXT_KEY = createContextKey("trigger.db.client_type");
8+
9+
export class DatasourceAttributeSpanProcessor implements SpanProcessor {
10+
onStart(span: Span, parentContext: Context): void {
11+
const ds = parentContext.getValue(DATASOURCE_CONTEXT_KEY);
12+
if (typeof ds === "string") {
13+
span.setAttribute("db.datasource", ds);
14+
}
15+
const clientType = parentContext.getValue(CLIENT_TYPE_CONTEXT_KEY);
16+
if (typeof clientType === "string") {
17+
span.setAttribute("db.client_type", clientType);
18+
}
19+
}
20+
onEnd(): void {}
21+
shutdown(): Promise<void> {
22+
return Promise.resolve();
23+
}
24+
forceFlush(): Promise<void> {
25+
return Promise.resolve();
26+
}
27+
}
28+
29+
export type DatasourceLabel =
30+
| "control-plane-writer"
31+
| "control-plane-replica"
32+
| "legacy-run-ops-writer"
33+
| "legacy-run-ops-replica"
34+
| "run-ops-writer"
35+
| "run-ops-replica"
36+
| "webhook-writer"
37+
| "webhook-partitions"
38+
| "webhook-replica"
39+
| "alert-partitions";
40+
41+
export function tagDatasource<T extends PrismaClient>(
42+
datasource: DatasourceLabel,
43+
clientType: string,
44+
client: T
45+
): T {
46+
return client.$extends({
47+
name: "datasource-tagger",
48+
query: {
49+
$allOperations: ({ query, args }) => {
50+
trace
51+
.getActiveSpan()
52+
?.setAttributes({ "db.datasource": datasource, "db.client_type": clientType });
53+
return context.with(
54+
context
55+
.active()
56+
.setValue(DATASOURCE_CONTEXT_KEY, datasource)
57+
.setValue(CLIENT_TYPE_CONTEXT_KEY, clientType),
58+
async () => await query(args)
59+
);
60+
},
61+
},
62+
}) as unknown as T;
63+
}
64+
65+
// Same extension as tagDatasource but typed for RunOpsPrismaClient (different
66+
// generated package — does not extend @trigger.dev/database.PrismaClient).
67+
export function tagDatasourceRunOps(
68+
datasource: DatasourceLabel,
69+
clientType: string,
70+
client: RunOpsPrismaClient
71+
): RunOpsPrismaClient {
72+
return client.$extends({
73+
name: "datasource-tagger",
74+
query: {
75+
$allOperations: ({ query, args }) => {
76+
trace
77+
.getActiveSpan()
78+
?.setAttributes({ "db.datasource": datasource, "db.client_type": clientType });
79+
return context.with(
80+
context
81+
.active()
82+
.setValue(DATASOURCE_CONTEXT_KEY, datasource)
83+
.setValue(CLIENT_TYPE_CONTEXT_KEY, clientType),
84+
async () => await query(args)
85+
);
86+
},
87+
},
88+
}) as unknown as RunOpsPrismaClient;
89+
}

‎apps/webapp/app/v3/tracer.server.ts‎

Lines changed: 1 addition & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@ import {
22
type Attributes,
33
type Context,
44
context as otelContext,
5-
createContextKey,
65
DiagConsoleLogger,
76
DiagLogLevel,
87
type Link,
@@ -70,27 +69,10 @@ import { logger } from "~/services/logger.server";
7069
import { metricsRegister } from "~/metrics.server";
7170
import { collectDatabaseClientMetrics } from "~/utils/databaseMetrics.server";
7271
import { performance } from "node:perf_hooks";
72+
import { DatasourceAttributeSpanProcessor } from "~/utils/datasourceTagging.server";
7373

7474
export const SEMINTATTRS_FORCE_RECORDING = "forceRecording";
7575

76-
export const DATASOURCE_CONTEXT_KEY = createContextKey("trigger.db.datasource");
77-
78-
class DatasourceAttributeSpanProcessor implements SpanProcessor {
79-
onStart(span: Span, parentContext: Context): void {
80-
const ds = parentContext.getValue(DATASOURCE_CONTEXT_KEY);
81-
if (typeof ds === "string") {
82-
span.setAttribute("db.datasource", ds);
83-
}
84-
}
85-
onEnd(): void {}
86-
shutdown(): Promise<void> {
87-
return Promise.resolve();
88-
}
89-
forceFlush(): Promise<void> {
90-
return Promise.resolve();
91-
}
92-
}
93-
9476
// Mirrors name-prefixed spans into a second exporter; they still flow to the main one
9577
class SpanNamePrefixMirrorProcessor implements SpanProcessor {
9678
constructor(

0 commit comments

Comments
 (0)