diff --git a/knip.json b/knip.json index 8e2d0f8a1..237330b99 100644 --- a/knip.json +++ b/knip.json @@ -37,7 +37,7 @@ "ignoreDependencies": ["tailwindcss", "tw-animate-css"] } }, - "ignoreDependencies": ["react-docgen-typescript"], + "ignoreDependencies": ["react-docgen-typescript", "esbuild", "ora", "zod"], "ignore": [ "**/*.generated.ts", "**/*.example.tsx", @@ -51,6 +51,7 @@ "packages/appkit/src/core/agent/load-agents.ts", "packages/appkit/src/connectors/mcp/index.ts", "packages/appkit/src/typedoc.entry.ts", + "packages/appkit/bin/appkit.js", "template/**", "tools/**", "docs/**", diff --git a/packages/appkit/src/plugins/lakebase/manifest.json b/packages/appkit/src/plugins/lakebase/manifest.json index 41c9366b8..fe93f032e 100644 --- a/packages/appkit/src/plugins/lakebase/manifest.json +++ b/packages/appkit/src/plugins/lakebase/manifest.json @@ -85,6 +85,10 @@ }, "scaffolding": { "rules": { + "must": [ + "Use withLakebaseTraceContext for direct Lakebase SQL to attach active request trace context at query execution", + "Keep SQL values in separate query parameters" + ], "should": [ "After init, run any database migrations for your chosen ORM before first request", "After init, verify Lakebase connectivity with 'psql $PGHOST -c \"select 1\"'" diff --git a/template/README.md.tmpl b/template/README.md.tmpl index c6564cb81..465c24cdb 100644 --- a/template/README.md.tmpl +++ b/template/README.md.tmpl @@ -48,6 +48,25 @@ DATABRICKS_APP_PORT=8000 #### Lakebase Configuration The Lakebase plugin requires additional environment variables for PostgreSQL connectivity. To learn how to configure the Lakebase plugin, see the [Lakebase plugin documentation](https://developers.databricks.com/docs/appkit/v0/plugins/lakebase). + +#### Correlate app requests with Lakebase query plans + +The sample routes execute Lakebase SQL through `withLakebaseTraceContext`. At query execution time, the wrapper reads the active AppKit request span and prepends its W3C `traceparent` plus a sanitized service label as a SQL comment. The label comes from `OTEL_SERVICE_NAME`, then `DATABRICKS_APP_NAME`, then `databricks-app`. If you configure a different AppKit telemetry service name, pass that name as the wrapper's second argument. The wrapper does not create a span, interpolate parameters, cache annotated SQL, or store request state on the connection pool. Without a valid active span, it sends the original SQL unchanged. + +When Databricks Apps provides `OTEL_EXPORTER_OTLP_ENDPOINT`, AppKit initializes request instrumentation and OTLP export. Configure an Apps OpenTelemetry destination to have the runtime inject that endpoint and make sampled application spans available in Unity Catalog: + +```yaml +resources: + apps: + app: + telemetry_export_destinations: + - unity_catalog: + logs_table: .._otel_logs + metrics_table: .._otel_metrics + traces_table: .._otel_spans +``` + +When Lakebase captures a query plan, it parses the SQL comment into `system.lakebase.plan_history`; it does not export the application span. For sampled requests, query the configured traces table for the app-side span and look for a captured plan with matching `trace_id` or `span_id`. Plan capture filters and retention still apply. Code that bypasses the wrapper (for example, a raw pool, checked-out client, ORM, or the opt-in Lakebase agent tool) is not covered by this template helper. {{- end}} {{- if .plugins.database}} diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index e32eb3ccb..bb4bad814 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -468,6 +468,10 @@ }, "scaffolding": { "rules": { + "must": [ + "Use withLakebaseTraceContext for direct Lakebase SQL to attach active request trace context at query execution", + "Keep SQL values in separate query parameters" + ], "should": [ "After init, run any database migrations for your chosen ORM before first request", "After init, verify Lakebase connectivity with 'psql $PGHOST -c \"select 1\"'" diff --git a/template/package-lock.json b/template/package-lock.json index 49239ced2..6ba2c9ec8 100644 --- a/template/package-lock.json +++ b/template/package-lock.json @@ -11,6 +11,7 @@ "dependencies": { "@databricks/appkit": "0.94.0", "@databricks/appkit-ui": "0.94.0", + "@opentelemetry/api": "1.9.0", "clsx": "2.1.1", "embla-carousel-react": "8.6.0", "lucide-react": "0.546.0", @@ -26,6 +27,8 @@ }, "devDependencies": { "@ast-grep/napi": "0.37.0", + "@opentelemetry/context-async-hooks": "2.12.0", + "@opentelemetry/sdk-trace-base": "2.12.0", "@tailwindcss/postcss": "4.1.17", "@tailwindcss/vite": "4.1.17", "@types/express": "5.0.5", diff --git a/template/package.json b/template/package.json index 827c945bf..215aaff59 100644 --- a/template/package.json +++ b/template/package.json @@ -34,6 +34,7 @@ "dependencies": { "@databricks/appkit": "0.94.0", "@databricks/appkit-ui": "0.94.0", + "@opentelemetry/api": "1.9.0", "clsx": "2.1.1", "embla-carousel-react": "8.6.0", "lucide-react": "0.546.0", @@ -49,6 +50,8 @@ }, "devDependencies": { "@ast-grep/napi": "0.37.0", + "@opentelemetry/context-async-hooks": "2.12.0", + "@opentelemetry/sdk-trace-base": "2.12.0", "@tailwindcss/postcss": "4.1.17", "@tailwindcss/vite": "4.1.17", "@types/express": "5.0.5", diff --git a/template/pnpm-lock.yaml b/template/pnpm-lock.yaml index 685d0e09a..c626c8032 100644 --- a/template/pnpm-lock.yaml +++ b/template/pnpm-lock.yaml @@ -23,6 +23,9 @@ importers: '@databricks/appkit-ui': specifier: 0.94.0 version: 0.94.0(@types/react-dom@19.2.2(@types/react@19.2.2))(@types/react@19.2.2)(react-dom@19.2.4(react@19.2.4))(react@19.2.4) + '@opentelemetry/api': + specifier: 1.9.0 + version: 1.9.0 clsx: specifier: 2.1.1 version: 2.1.1 @@ -63,6 +66,12 @@ importers: '@ast-grep/napi': specifier: 0.37.0 version: 0.37.0 + '@opentelemetry/context-async-hooks': + specifier: 2.12.0 + version: 2.12.0(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-trace-base': + specifier: 2.12.0 + version: 2.12.0(@opentelemetry/api@1.9.0) '@tailwindcss/postcss': specifier: 4.1.17 version: 4.1.17 diff --git a/template/server/routes/lakebase/todo-routes.ts b/template/server/routes/lakebase/todo-routes.ts index 8f04934f3..3a0684147 100644 --- a/template/server/routes/lakebase/todo-routes.ts +++ b/template/server/routes/lakebase/todo-routes.ts @@ -4,6 +4,7 @@ import { z } from 'zod'; import { Application } from 'express'; +import { withLakebaseTraceContext } from './trace-context'; interface AppKitWithLakebase { lakebase: { @@ -33,13 +34,15 @@ const CREATE_TABLE_SQL = ` const CreateTodoBody = z.object({ title: z.string().min(1) }); export async function setupSampleLakebaseRoutes(appkit: AppKitWithLakebase) { + const lakebase = withLakebaseTraceContext(appkit.lakebase); + try { - const { rows } = await appkit.lakebase.query(TABLE_EXISTS_SQL); + const { rows } = await lakebase.query(TABLE_EXISTS_SQL); if (rows.length > 0) { console.log('[lakebase] Table app.todos already exists, skipping setup'); } else { - await appkit.lakebase.query(SETUP_SCHEMA_SQL); - await appkit.lakebase.query(CREATE_TABLE_SQL); + await lakebase.query(SETUP_SCHEMA_SQL); + await lakebase.query(CREATE_TABLE_SQL); console.log('[lakebase] Created schema and table app.todos'); } } catch (err) { @@ -51,7 +54,7 @@ export async function setupSampleLakebaseRoutes(appkit: AppKitWithLakebase) { appkit.server.extend((app) => { app.get('/api/lakebase/todos', async (_req, res) => { try { - const result = await appkit.lakebase.query( + const result = await lakebase.query( 'SELECT id, title, completed, created_at FROM app.todos ORDER BY created_at DESC', ); res.json(result.rows); @@ -68,7 +71,7 @@ export async function setupSampleLakebaseRoutes(appkit: AppKitWithLakebase) { res.status(400).json({ error: 'title is required' }); return; } - const result = await appkit.lakebase.query( + const result = await lakebase.query( 'INSERT INTO app.todos (title) VALUES ($1) RETURNING id, title, completed, created_at', [parsed.data.title.trim()], ); @@ -86,7 +89,7 @@ export async function setupSampleLakebaseRoutes(appkit: AppKitWithLakebase) { res.status(400).json({ error: 'Invalid id' }); return; } - const result = await appkit.lakebase.query( + const result = await lakebase.query( 'UPDATE app.todos SET completed = NOT completed WHERE id = $1 RETURNING id, title, completed, created_at', [id], ); @@ -108,7 +111,7 @@ export async function setupSampleLakebaseRoutes(appkit: AppKitWithLakebase) { res.status(400).json({ error: 'Invalid id' }); return; } - const result = await appkit.lakebase.query( + const result = await lakebase.query( 'DELETE FROM app.todos WHERE id = $1 RETURNING id', [id], ); diff --git a/template/server/routes/lakebase/trace-context.test.ts b/template/server/routes/lakebase/trace-context.test.ts new file mode 100644 index 000000000..b939114c7 --- /dev/null +++ b/template/server/routes/lakebase/trace-context.test.ts @@ -0,0 +1,228 @@ +{{if .plugins.lakebase -}} +import { AsyncLocalStorageContextManager } from '@opentelemetry/context-async-hooks'; +import { context, ROOT_CONTEXT, trace, TraceFlags, type SpanContext } from '@opentelemetry/api'; +import { BasicTracerProvider, InMemorySpanExporter, SimpleSpanProcessor } from '@opentelemetry/sdk-trace-base'; +import { afterAll, beforeAll, describe, expect, test, vi } from 'vitest'; +import { withLakebaseTraceContext } from './trace-context'; + +const TRACE_A = '4bf92f3577b34da6a3ce929d0e0e4736'; +const TRACE_B = '70f5c2a87bdb4a84a8f5eab6329b86db'; +const SPAN_A = '00f067aa0ba902b7'; +const SPAN_B = '8f3b1c42d5e60718'; + +function runWithSpan(spanContext: SpanContext, fn: () => T): T { + const span = trace.wrapSpanContext(spanContext); + return context.with(trace.setSpan(context.active(), span), fn); +} + +describe('withLakebaseTraceContext', () => { + beforeAll(() => { + context.disable(); + context.setGlobalContextManager(new AsyncLocalStorageContextManager().enable()); + }); + + afterAll(() => { + context.disable(); + }); + + test('adds the valid active span and preserves its trace flags', async () => { + const params = ['todo-title']; + const query = vi.fn(() => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + + await runWithSpan( + { traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.NONE }, + () => lakebase.query('SELECT * FROM app.todos WHERE title = $1', params), + ); + + expect(query).toHaveBeenCalledWith( + `/*traceparent='00-${TRACE_A}-${SPAN_A}-00',service='orders-app'*/\nSELECT * FROM app.todos WHERE title = $1`, + params, + ); + }); + + test('keeps sampled flags instead of hard-coding them', async () => { + const query = vi.fn(() => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + + await runWithSpan( + { traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.SAMPLED }, + () => lakebase.query('SELECT 1'), + ); + + expect(query).toHaveBeenCalledWith( + `/*traceparent='00-${TRACE_A}-${SPAN_A}-01',service='orders-app'*/\nSELECT 1`, + undefined, + ); + }); + + test('formats uppercase valid span IDs as lowercase W3C hex', async () => { + const query = vi.fn(() => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + + await runWithSpan( + { traceId: TRACE_A.toUpperCase(), spanId: SPAN_A.toUpperCase(), traceFlags: TraceFlags.SAMPLED }, + () => lakebase.query('SELECT 1'), + ); + + expect(query).toHaveBeenCalledWith( + `/*traceparent='00-${TRACE_A}-${SPAN_A}-01',service='orders-app'*/\nSELECT 1`, + undefined, + ); + }); + + test('uses the app name when OTEL_SERVICE_NAME is empty', async () => { + vi.stubEnv('OTEL_SERVICE_NAME', ''); + vi.stubEnv('DATABRICKS_APP_NAME', 'inventory-app'); + try { + const query = vi.fn(() => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }); + await runWithSpan( + { traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.SAMPLED }, + () => lakebase.query('SELECT 1'), + ); + expect(query).toHaveBeenCalledWith( + `/*traceparent='00-${TRACE_A}-${SPAN_A}-01',service='inventory-app'*/\nSELECT 1`, + undefined, + ); + } finally { + vi.unstubAllEnvs(); + } + }); + + test.each([ + ['missing', undefined], + [ + 'invalid', + { + traceId: '00000000000000000000000000000000', + spanId: '0000000000000000', + traceFlags: TraceFlags.SAMPLED, + } satisfies SpanContext, + ], + ])('executes SQL unchanged when the active context is %s', async (_name, spanContext) => { + const params = [42]; + const query = vi.fn(() => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + const execute = () => lakebase.query('SELECT * FROM app.todos WHERE id = $1', params); + + if (spanContext) { + await runWithSpan(spanContext, execute); + } else { + await context.with(ROOT_CONTEXT, execute); + } + + expect(query).toHaveBeenCalledWith('SELECT * FROM app.todos WHERE id = $1', params); + }); + + test('isolates concurrent request contexts and keeps parameter binding', async () => { + const calls: Array<{ text: string; params?: unknown[] }> = []; + const query = vi.fn(async (text: string, params?: unknown[]) => { + await Promise.resolve(); + calls.push({ text, params }); + return { rows: [] }; + }); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + const paramsA = ['a']; + const paramsB = ['b']; + + await Promise.all([ + runWithSpan({ traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.SAMPLED }, async () => { + await Promise.resolve(); + await lakebase.query('SELECT $1', paramsA); + }), + runWithSpan({ traceId: TRACE_B, spanId: SPAN_B, traceFlags: TraceFlags.NONE }, async () => { + await Promise.resolve(); + await lakebase.query('SELECT $1', paramsB); + }), + ]); + + expect(calls).toEqual( + expect.arrayContaining([ + { + text: `/*traceparent='00-${TRACE_A}-${SPAN_A}-01',service='orders-app'*/\nSELECT $1`, + params: paramsA, + }, + { + text: `/*traceparent='00-${TRACE_B}-${SPAN_B}-00',service='orders-app'*/\nSELECT $1`, + params: paramsB, + }, + ]), + ); + }); + + test('uses a real exported span as the SQL correlation identity', async () => { + const exporter = new InMemorySpanExporter(); + const provider = new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] }); + trace.disable(); + trace.setGlobalTracerProvider(provider); + + try { + const query = vi.fn((_text: string, _params?: unknown[]) => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + const params = ['a']; + + await trace.getTracer('test').startActiveSpan('http.request', async (span) => { + try { + await lakebase.query('SELECT $1', params); + } finally { + span.end(); + } + }); + + const exportedSpan = exporter.getFinishedSpans()[0]; + expect(exportedSpan.name).toBe('http.request'); + expect(query).toHaveBeenCalledWith( + `/*traceparent='00-${exportedSpan.spanContext().traceId}-${exportedSpan.spanContext().spanId}-01',service='orders-app'*/\nSELECT $1`, + params, + ); + } finally { + trace.disable(); + await provider.shutdown(); + } + }); + + test('sanitizes the service label before putting it in a SQL comment', async () => { + const query = vi.fn((_text: string) => Promise.resolve({ rows: [] })); + const lakebase = withLakebaseTraceContext({ query }, "orders'*/\nDROP TABLE app.todos; --"); + + await runWithSpan( + { traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.SAMPLED }, + () => lakebase.query('SELECT 1'), + ); + + const sql = query.mock.calls[0][0]; + const match = sql.match(/^\/\*traceparent='00-[0-9a-f]{32}-[0-9a-f]{16}-01',service='([^']+)'\*\/\nSELECT 1$/); + expect(match).not.toBeNull(); + expect(match?.[1]).toMatch(/^[A-Za-z0-9._:/-]{1,128}$/); + expect(sql).not.toContain('DROP TABLE app.todos;'); + }); + + test('propagates query errors and reads context again on the next execution', async () => { + const failure = new Error('database unavailable'); + const query = vi.fn() + .mockRejectedValueOnce(failure) + .mockResolvedValueOnce({ rows: [] }); + const lakebase = withLakebaseTraceContext({ query }, 'orders-app'); + const params = [42]; + + await expect(runWithSpan( + { traceId: TRACE_A, spanId: SPAN_A, traceFlags: TraceFlags.SAMPLED }, + () => lakebase.query('SELECT $1', params), + )).rejects.toBe(failure); + await runWithSpan( + { traceId: TRACE_B, spanId: SPAN_B, traceFlags: TraceFlags.NONE }, + () => lakebase.query('SELECT $1', params), + ); + + expect(query.mock.calls[0]).toEqual([ + `/*traceparent='00-${TRACE_A}-${SPAN_A}-01',service='orders-app'*/\nSELECT $1`, + params, + ]); + expect(query.mock.calls[1]).toEqual([ + `/*traceparent='00-${TRACE_B}-${SPAN_B}-00',service='orders-app'*/\nSELECT $1`, + params, + ]); + }); +}); +{{- end}} diff --git a/template/server/routes/lakebase/trace-context.ts b/template/server/routes/lakebase/trace-context.ts new file mode 100644 index 000000000..ebe07a783 --- /dev/null +++ b/template/server/routes/lakebase/trace-context.ts @@ -0,0 +1,51 @@ +{{if .plugins.lakebase -}} +import { context, isSpanContextValid, trace } from '@opentelemetry/api'; + +export interface LakebaseQueryClient { + query(text: string, params?: unknown[]): Promise<{ rows: Record[] }>; +} + +const DEFAULT_SERVICE_NAME = 'databricks-app'; +const MAX_SERVICE_NAME_LENGTH = 128; + +function serviceNameFromEnvironment(): string { + return process.env.OTEL_SERVICE_NAME || process.env.DATABRICKS_APP_NAME || DEFAULT_SERVICE_NAME; +} + +function safeServiceName(serviceName: string): string { + const sanitized = serviceName + .trim() + .replace(/[^A-Za-z0-9._:/-]/g, '_') + .slice(0, MAX_SERVICE_NAME_LENGTH); + return sanitized || DEFAULT_SERVICE_NAME; +} + +/** + * Adds the active OpenTelemetry span context immediately before a query is + * executed. No span is created here, and no request state is stored on the + * shared Lakebase pool. + */ +export function withLakebaseTraceContext( + lakebase: LakebaseQueryClient, + serviceName = serviceNameFromEnvironment(), +): LakebaseQueryClient { + const service = safeServiceName(serviceName); + + return { + query(text, params) { + const span = trace.getSpan(context.active()); + const spanContext = span?.spanContext(); + + if (!spanContext || !isSpanContextValid(spanContext)) { + return lakebase.query(text, params); + } + + const flags = (spanContext.traceFlags & 0xff).toString(16).padStart(2, '0'); + const traceparent = `00-${spanContext.traceId.toLowerCase()}-${spanContext.spanId.toLowerCase()}-${flags}`; + const annotatedSql = `/*traceparent='${traceparent}',service='${service}'*/\n${text}`; + + return lakebase.query(annotatedSql, params); + }, + }; +} +{{- end}}