feat(server-utils): Add the orchestrion config for @neondatabase/serverless

packages/server-utils/package.json

@@ -102,6 +102,7 @@
   },
   "devDependencies": {
     "@apm-js-collab/code-transformer-bundler-plugins": "^0.7.4",
+    "@neondatabase/serverless": "1.2.0",
     "@types/node": "^18.19.1",
     "meriyah": "^6.1.4",
     "vite": "^8.3.1"

packages/server-utils/src/orchestrion/config/neon.ts

@@ -0,0 +1,51 @@
+import type { InstrumentationConfig } from '../apmTypes';
+
+import { getModuleNames } from './module-names';
+
+// `@neondatabase/serverless` ships as two minified esbuild bundles (`index.js` for CJS,
+// `index.mjs` for ESM) that inline their own copy of `pg`. Class and function identifiers are
+// mangled and differ between the two builds, so the name-based `functionQuery` matchers can't
+// find anything. Every selector below keys on property names and string literals, which survive
+// minification and are identical in both bundles.
+const module = { name: '@neondatabase/serverless', versionRange: '>=1.0.0 <2', filePath: /^index\.m?js$/ };
+
+export const neonConfig = [
+  // The bundled pg `Client.prototype.query` (WebSocket driver). `_pulseQueryQueue` is a pg
+  // internal that only `Client` defines, which keeps the pool and wire-protocol `query` methods
+  // out. `methodName` is ignored for matching (`astQuery` wins) but makes the transformer emit a
+  // method wrapper that exposes `self`, which carries `connectionParameters`.
+  {
+    channelName: 'query',
+    module,
+    astQuery:
+      'ClassBody:has(MethodDefinition[key.name="_pulseQueryQueue"]) > MethodDefinition[key.name="query"] > FunctionExpression',
+    functionQuery: { methodName: 'query', kind: 'Auto' },
+  },
+  // The HTTP executor nested inside `neon()`. Every HTTP query funnels through it: lazy
+  // `NeonQueryPromise`s call it from `then`/`catch`/`finally`, and `transaction()` calls it once
+  // with the whole batch, so wrapping it gives one span per request with no double counting.
+  {
+    channelName: 'http-query',
+    module,
+    astQuery: 'FunctionDeclaration[async=true]:has(Literal[value="Neon-Connection-String"])',
+    functionQuery: { kind: 'Async' },
+  },
+  // `resolveConnectionParams`, awaited by the HTTP executor. The connection string lives in the
+  // `neon()` closure and never reaches the executor's channel context, so the subscriber reads
+  // `server.address`/`db.namespace` from this call's result instead, inside the query span.
+  {
+    channelName: 'resolve-connection',
+    module,
+    astQuery:
+      'FunctionDeclaration[async=true]:has(ReturnStatement > ObjectExpression > Property[key.name="resolvedURL"])',
+    functionQuery: { kind: 'Async' },
+  },
+] satisfies InstrumentationConfig[];
+
+export const neonModuleNames = getModuleNames(neonConfig);
+
+export const neonChannels = {
+  NEON_QUERY: 'orchestrion:@neondatabase/serverless:query',
+  NEON_HTTP_QUERY: 'orchestrion:@neondatabase/serverless:http-query',
+  NEON_RESOLVE_CONNECTION: 'orchestrion:@neondatabase/serverless:resolve-connection',
+} as const;

packages/server-utils/test/orchestrion/neon-bundle.test.ts

@@ -0,0 +1,121 @@
+import { createCodeTransformer } from '@apm-js-collab/code-transformer-bundler-plugins/core';
+import { mkdtempSync, readFileSync, rmSync, writeFileSync } from 'node:fs';
+import { createRequire } from 'node:module';
+import { tmpdir } from 'node:os';
+import { dirname, join } from 'node:path';
+import { pathToFileURL } from 'node:url';
+import { tracingChannel } from 'node:diagnostics_channel';
+import { afterAll, beforeAll, describe, expect, it } from 'vitest';
+import { neonChannels, neonConfig } from '../../src/orchestrion/config/neon';
+
+// The selectors key on internals of Neon's minified bundles, so this runs the bundler plugins'
+// code transformer over the real published files and executes the result: a release that
+// reshapes the bundle fails here rather than silently dropping spans.
+const nodeRequire = createRequire(import.meta.url);
+const packageDir = dirname(nodeRequire.resolve('@neondatabase/serverless'));
+
+interface NeonModule {
+  neon: (connectionString: string) => {
+    (strings: TemplateStringsArray, ...values: unknown[]): Promise<unknown>;
+    query: (text: string, params?: unknown[]) => Promise<unknown>;
+    transaction: (queries: Promise<unknown>[]) => Promise<unknown>;
+  };
+  neonConfig: { fetchFunction: unknown };
+  Client: new (connectionString: string) => {
+    query: (text: string) => Promise<unknown>;
+    end: () => Promise<void>;
+  };
+}
+
+function neonResponse(body: unknown): () => Promise<Response> {
+  return () => Promise.resolve(new Response(JSON.stringify(body), { status: 200 }));
+}
+
+const EMPTY_RESULT = { rows: [], fields: [], command: 'SELECT', rowCount: 0 };
+
+// Every wrapper opens with a `hasSubscribers` guard on its channel, one per instrumented site.
+function countWrappedSites(code: string, channelVariable: string): number {
+  return code.split(`tr_ch_apm_hasSubscribers(${channelVariable})`).length - 1;
+}
+
+describe('neon orchestrion config on the published bundles', () => {
+  let outDir: string;
+  const events: string[] = [];
+
+  beforeAll(() => {
+    outDir = mkdtempSync(join(tmpdir(), 'neon-bundle-'));
+    for (const name of ['query', 'http-query', 'resolve-connection']) {
+      tracingChannel(`orchestrion:@neondatabase/serverless:${name}`).subscribe({
+        start: data => {
+          const ctx = data as { arguments: unknown[]; self?: { connectionParameters?: { database?: string } } };
+          events.push(`${name}:start:${ctx.self?.connectionParameters?.database ?? JSON.stringify(ctx.arguments[0])}`);
+        },
+        end: () => {},
+        asyncStart: () => {},
+        asyncEnd: data => {
+          const ctx = data as { result?: { resolvedURL?: URL } };
+          events.push(`${name}:asyncEnd:${ctx.result?.resolvedURL?.hostname ?? ''}`);
+        },
+        error: () => {},
+      });
+    }
+  });
+
+  afterAll(() => rmSync(outDir, { recursive: true, force: true }));
+
+  it.each([
+    ['index.js', 'cjs'],
+    ['index.mjs', 'esm'],
+  ])('instruments both drivers in %s', async (file, format) => {
+    const transformer = createCodeTransformer({ instrumentations: neonConfig });
+    const source = join(packageDir, file);
+    const result = transformer.transform(readFileSync(source, 'utf8'), source);
+    expect(result).not.toBeNull();
+
+    for (const channel of Object.values(neonChannels)) {
+      expect(result!.code).toContain(channel);
+    }
+    // Only pg's `Client.prototype.query` gets the channel, not the pool or wire-protocol `query`.
+    expect(countWrappedSites(result!.code, 'tr_ch_apm$query')).toBe(1);
+    expect(countWrappedSites(result!.code, 'tr_ch_apm$http_query')).toBe(1);
+    expect(countWrappedSites(result!.code, 'tr_ch_apm$resolve_connection')).toBe(1);
+    // The query wrapper exposes `self`, which carries pg's `connectionParameters`.
+    const querySite = result!.code.indexOf('Client was passed a null or undefined query');
+    expect(result!.code.slice(result!.code.lastIndexOf('__apm$ctx = {', querySite), querySite)).toContain('self: this');
+
+    const outFile = join(outDir, file);
+    writeFileSync(outFile, result!.code);
+    const mod: NeonModule =
+      format === 'cjs' ? nodeRequire(outFile) : ((await import(pathToFileURL(outFile).href)) as NeonModule);
+
+    events.length = 0;
+    mod.neonConfig.fetchFunction = neonResponse(EMPTY_RESULT);
+    const sql = mod.neon('postgres://user:pass@ep-x.aws.neon.tech/neondb');
+    await sql.query('SELECT $1', [1]);
+    await sql`SELECT 2`;
+    mod.neonConfig.fetchFunction = neonResponse({ results: [EMPTY_RESULT, EMPTY_RESULT] });
+    await sql.transaction([sql`SELECT 3`, sql`SELECT 4`]);
+
+    const client = new mod.Client('postgres://user:pass@ep-x.aws.neon.tech/neondb');
+    // Never connects: the channel fires at the call, before the socket is needed.
+    client.query('SELECT 5').catch(() => {});
+    await client.end().catch(() => {});
+
+    const resolve = [
+      'resolve-connection:start:"postgres://user:pass@ep-x.aws.neon.tech/neondb"',
+      'resolve-connection:asyncEnd:ep-x.aws.neon.tech',
+    ];
+    expect(events).toEqual([
+      'http-query:start:{"query":"SELECT $1","params":[1]}',
+      ...resolve,
+      'http-query:asyncEnd:',
+      'http-query:start:{"strings":["SELECT 2"],"values":[]}',
+      ...resolve,
+      'http-query:asyncEnd:',
+      'http-query:start:[{"strings":["SELECT 3"],"values":[]},{"strings":["SELECT 4"],"values":[]}]',
+      ...resolve,
+      'http-query:asyncEnd:',
+      'query:start:neondb',
+    ]);
+  });
+});

yarn.lock

@@ -5664,6 +5664,11 @@
     "@emnapi/runtime" "^1.1.0"
     "@tybys/wasm-util" "^0.9.0"
 
+"@neondatabase/serverless@1.2.0":
+  version "1.2.0"
+  resolved "https://registry.yarnpkg.com/@neondatabase/serverless/-/serverless-1.2.0.tgz#d58553f0804eb9ddacb09662d607756b7dc74b77"
+  integrity sha512-tnaqgixUengZQNd9ZrS3ytQfamJzLKY8QxgqWQTNHHtB+KrFQWk77JTbbpg4RD+J3qE4eO2lVytvGDX5Kj6FdQ==
+
 "@nestjs/common@^10.0.0":
   version "10.4.15"
   resolved "https://registry.yarnpkg.com/@nestjs/common/-/common-10.4.15.tgz#27c291466d9100eb86fdbe6f7bbb4d1a6ad55f70"