feat: make astQuery a first-class feature

README.md

@@ -161,6 +161,52 @@ substituted return value only applies when the function returns (or
 resolves) normally. If no subscriber reassigns `message.result`, the
 original return value is preserved unchanged.
 
+### AST Query
+
+The name-based `FunctionQuery` variants cover the common cases
+(named functions, class/object methods, expressions). When you
+need to target a node they can't express, such as an **anonymous
+function returned by a factory** (a decorator factory, a
+per-request handler), you can set `astQuery` to a raw
+[esquery](https://github.com/estools/esquery) selector instead.
+
+When present, `astQuery` chooses the nodes to instrument and
+takes precedence over `functionQuery`'s matching fields;
+`functionQuery` then only supplies behaviour (`kind`, `index`,
+`callbackIndex`) and may be omitted (it defaults to `kind:
+"Sync"`).
+
+For example, to instrument the decorator returned by a factory:
+
+```js
+function Injectable(options) {
+    return (target) => { /* applied to the decorated class */ };
+}
+```
+
+```js
+const instrumentation = {
+    channelName: "injectable-apply",
+    module: { name: "@nestjs/common", versionRange: ">=8.0.0", filePath: "decorators/core/injectable.decorator.js" },
+    // Match the arrow returned from `Injectable`. There is no name to target!
+    astQuery: 'FunctionDeclaration[id.name="Injectable"] ReturnStatement > ArrowFunctionExpression',
+    functionQuery: { kind: "Sync" },
+};
+```
+
+The channel then fires each time the decorator is applied, with
+the decorated target available as `message.arguments[0]`, which a
+subscriber can mutate, for example to wrap prototype methods.
+
+An `astQuery` is used verbatim, so it can match any node,
+including ones the name-based variants don't expose, such as
+anonymous or deeply nested functions. (Both name-based and
+`astQuery` matching work on synchronous and async functions
+alike.)
+
+If an `astQuery` matches no nodes, the "failed to find injection
+points" error includes the selector so it can be debugged.
+
 ### API Reference
 
 ```ts
@@ -212,11 +258,28 @@ type ModuleMatcher = {
 #### **`InstrumentationConfig`**
 
 ```ts
-type InstrumentationConfig = {
-    channelName: string; // Name of the diagnostics channel
-    module: ModuleMatcher;
-    functionQuery: FunctionQuery;
+// Behaviour-only fields, used when `astQuery` does the matching.
+type FunctionBehavior = {
+    kind?: FunctionKind;
+    index?: number | null;
+    callbackIndex?: number;
 };
+
+type InstrumentationConfig =
+    | {
+        channelName: string; // Name of the diagnostics channel
+        module: ModuleMatcher;
+        functionQuery: FunctionQuery; // Name-based matching
+        astQuery?: string; // Raw esquery selector; takes precedence over functionQuery matching
+        transform?: string; // Name of a custom transform registered via addTransform
+    }
+    | {
+        channelName: string;
+        module: ModuleMatcher;
+        astQuery: string; // Raw esquery selector chooses the node(s)
+        functionQuery?: FunctionBehavior; // Behaviour only; matching fields ignored
+        transform?: string;
+    };
 ```
 
 ### Functions

index.d.ts

@@ -36,28 +36,58 @@ export type FunctionQuery = { className: string; methodName: string; kind: Funct
 export type CustomTransform = (state: unknown, node: Node, parent: Node, ancestry: Node[]) => void;
 
 /**
- * Configuration for injecting instrumentation code
+ * The behaviour-only fields of a `FunctionQuery`. Used together with `astQuery`,
+ * where the raw selector chooses the node and these fields drive how it is
+ * wrapped (the name-based matching fields are ignored).
  */
-export interface InstrumentationConfig {
-    /**
-     * The name of the diagnostics channel to publish to
-     */
-    channelName: string;
-    /**
-     * The module matcher to identify the module and file to instrument
-     */
-    module: ModuleMatcher;
-    /**
-     * The function query to identify the function to instrument
-     */
-    functionQuery: FunctionQuery;
-    /**
-     * The name of a custom transform registered via `addTransform`.
-     * When set, takes precedence over `functionQuery.kind`.
-     */
-    transform?: string;
+export interface FunctionBehavior {
+    kind?: FunctionKind;
+    index?: number | null;
+    callbackIndex?: number;
+    mutableResult?: boolean;
 }
 
+/**
+ * Configuration for injecting instrumentation code.
+ *
+ * Either `functionQuery` (name-based matching) or `astQuery` (a raw esquery
+ * selector) must identify the node(s) to instrument. When `astQuery` is set it
+ * takes precedence over `functionQuery`'s matching fields, and `functionQuery`
+ * becomes an optional bag of behaviour options ({@link FunctionBehavior}).
+ */
+export type InstrumentationConfig =
+    | {
+        /** The name of the diagnostics channel to publish to */
+        channelName: string;
+        /** The module matcher to identify the module and file to instrument */
+        module: ModuleMatcher;
+        /** The function query to identify the function to instrument */
+        functionQuery: FunctionQuery;
+        /**
+         * A raw esquery selector that chooses the node(s) to instrument. When
+         * set, it takes precedence over `functionQuery`'s matching fields.
+         */
+        astQuery?: string;
+        /**
+         * The name of a custom transform registered via `addTransform`.
+         * When set, takes precedence over `functionQuery.kind`.
+         */
+        transform?: string;
+    }
+    | {
+        channelName: string;
+        module: ModuleMatcher;
+        /**
+         * A raw esquery selector that chooses the node(s) to instrument. This is
+         * the escape hatch for shapes the name-based `functionQuery` can't
+         * express, e.g. an anonymous arrow returned by a factory function.
+         */
+        astQuery: string;
+        /** Behaviour options for the matched node(s); matching fields are ignored. */
+        functionQuery?: FunctionBehavior;
+        transform?: string;
+    };
+
 /**
  * Describes the module and file path you would like to match
  */

lib/transformer.js

@@ -101,6 +101,11 @@ class Transformer {
       }
 
       const resolvedFunctionQuery = this.#resolveExportAlias(functionQuery, aliases)
+      // A raw `astQuery` selector takes precedence over the name-based
+      // `functionQuery`: it chooses which node(s) to instrument, while
+      // `functionQuery` still supplies behaviour (kind/index/callbackIndex).
+      // This is the escape hatch for shapes the name-based queries can't
+      // express (e.g. an anonymous arrow returned by a factory).
       const query = astQuery || this.#fromFunctionQuery(resolvedFunctionQuery)
       const state = {
         ...config,
@@ -119,7 +124,8 @@ class Transformer {
     }
 
     if (injectionCount === 0 && this.#configs.length > 0) {
-      const names = this.#configs.map(({ functionQuery = {} }) => {
+      const names = this.#configs.map(({ astQuery, functionQuery = {} }) => {
+        if (astQuery) return astQuery
         const resolvedQuery = this.#resolveExportAlias(functionQuery, aliases)
         const queryName = (q) => q.methodName || q.privateMethodName || q.functionName || q.expressionName || 'constructor'
         const originalName = queryName(functionQuery)

tests/ast_query_returned_arrow_cjs/mod.js

@@ -0,0 +1,12 @@
+// Mirrors a NestJS-style decorator factory: `Decorator(options)` returns the
+// actual decorator `(target) => {...}` that is applied to a class. The returned
+// arrow is anonymous and synchronous, so the name-based FunctionQuery variants
+// cannot target it. Only a raw `astQuery` can do that.
+function Decorator (options) {
+  return (target) => {
+    target.decorated = options
+    return target
+  }
+}
+
+module.exports = { Decorator }

tests/ast_query_returned_arrow_cjs/test.js

@@ -0,0 +1,26 @@
+const { Decorator } = require('./instrumented.js')
+const assert = require('node:assert')
+const { tracingChannel } = require('node:diagnostics_channel')
+
+// The channel fires when the (anonymous, sync) returned decorator is applied.
+// A subscriber can mutate the decorated target (arg 0): building block for
+// instrumenting decorated classes.
+let started = false
+tracingChannel('orchestrion:undici:decorator_apply').subscribe({
+  start (message) {
+    started = true
+    const target = message.arguments[0]
+    target.instrumented = true
+  }
+})
+
+class MyService {}
+const decorate = Decorator({ scope: 'request' })
+const result = decorate(MyService)
+
+assert.strictEqual(started, true)
+// Original decorator still ran.
+assert.strictEqual(result, MyService)
+assert.deepStrictEqual(MyService.decorated, { scope: 'request' })
+// Subscriber's mutation of the target took effect.
+assert.strictEqual(MyService.instrumented, true)

tests/tests.test.mjs

@@ -446,6 +446,35 @@ describe('mutable_result_async_cjs', () => {
   })
 })
 
+describe('ast_query_returned_arrow_cjs', () => {
+  test('instruments an anonymous arrow returned by a factory via astQuery', () => {
+    runTest('ast_query_returned_arrow_cjs', [
+      {
+        channelName: 'decorator_apply',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        astQuery: 'FunctionDeclaration[id.name="Decorator"] ReturnStatement > ArrowFunctionExpression',
+        functionQuery: { kind: 'Sync' },
+      },
+    ])
+  })
+
+  test('reports the astQuery selector when no injection points are found', () => {
+    const instrumentor = create([
+      {
+        channelName: 'no_match',
+        module: { name: TEST_MODULE_NAME, versionRange: '>=0.0.1', filePath: TEST_MODULE_PATH },
+        astQuery: 'FunctionDeclaration[id.name="doesNotExist"]',
+        functionQuery: { kind: 'Sync' },
+      },
+    ])
+    const transformer = instrumentor.getTransformer(TEST_MODULE_NAME, TEST_MODULE_VERSION, TEST_MODULE_PATH)
+    assert.throws(
+      () => transformer.transform('function fetch () { return 42 }', 'cjs'),
+      /doesNotExist/
+    )
+  })
+})
+
 describe('polyfill_cjs', () => {
   test('instruments with a custom dc module (cjs)', () => {
     runTest('polyfill_cjs', [