Add bindContext support

Author: timokoessler

.nvmrc

@@ -1 +1 @@
-22.12
+24.3.0

library/agent/hooks/instrumentation/injectedFunctions.ts

@@ -1,5 +1,5 @@
 import { getInstance } from "../../AgentSingleton";
-import { getContext } from "../../Context";
+import { bindContext, getContext } from "../../Context";
 import { inspectArgs } from "../wrapExport";
 import { getPackageCallbackInfo } from "./instructions";
 
@@ -67,7 +67,17 @@ export function __instrumentModifyArgs(
     const newArgs = cbInfo.funcs.modifyArgs(args, agent, subject);
     // Only return the new arguments if they are an array
     if (Array.isArray(newArgs)) {
-      return newArgs;
+      if (!cbInfo.funcs.bindContext) {
+        return newArgs;
+      }
+
+      // Ensure that all functions in the new arguments are bound to the current execution context (only if bindContext is true)
+      return newArgs.map((arg) => {
+        if (typeof arg === "function") {
+          return bindContext(arg as () => unknown);
+        }
+        return arg;
+      });
     }
   } catch (error) {
     if (error instanceof Error) {

library/agent/hooks/instrumentation/instructions.ts

@@ -1,3 +1,4 @@
+/* eslint-disable max-lines-per-function */
 import { Package } from "../Package";
 import { BuiltinModule } from "../BuiltinModule";
 import type {
@@ -45,6 +46,12 @@ export function setPackagesToInstrument(_packages: Package[]) {
               functions: file.functions.map((func) => {
                 const identifier = `${pkg.getName()}.${file.path}.${func.name}.${func.nodeType}.${versionedPackage.getRange()}`;
 
+                // If bindContext is set to true, but no modifyArgs is defined, modifyArgs will be set to a stub function
+                // The reason for this is that the bindContext logic needs to modify the arguments
+                if (func.bindContext && !func.modifyArgs) {
+                  func.modifyArgs = (args) => args;
+                }
+
                 packageCallbackInfo.set(identifier, {
                   pkgName: pkg.getName(),
                   methodName: func.name,
@@ -53,6 +60,7 @@ export function setPackagesToInstrument(_packages: Package[]) {
                     inspectArgs: func.inspectArgs,
                     modifyArgs: func.modifyArgs,
                     modifyReturnValue: func.modifyReturnValue,
+                    bindContext: func.bindContext ?? false,
                   },
                 });
 

library/agent/hooks/instrumentation/types.ts

@@ -58,6 +58,7 @@ export type IntereptorFunctionsObj = {
   inspectArgs?: InspectArgsInterceptor;
   modifyArgs?: ModifyArgsInterceptor;
   modifyReturnValue?: ModifyReturnValueInterceptor;
+  bindContext: boolean;
 };
 
 export type IntereptorCallbackInfoObj = {
@@ -92,6 +93,11 @@ export type PackageFunctionInstrumentationInstruction = {
    * In strict mode, the arguments object and the named arguments are not synced, so changing the arguments object will not change the named arguments and vice versa.
    */
   modifyArgumentsObject?: boolean;
+  /**
+   * Binds the given function to the current execution context to prevent that the Zen context using the async local storage is lost.
+   * If enabled, the bindContext function will be called for all callbacks that are passed to the function.
+   */
+  bindContext?: boolean;
 };
 
 export type PackageFileInstrumentationInstruction = {

library/sinks/MySQL.ts

@@ -65,6 +65,7 @@ export class MySQL implements Wrapper {
             name: "Connection.prototype.query",
             nodeType: "FunctionAssignment",
             operationKind: "sql_op",
+            bindContext: true,
             inspectArgs: (args) => this.inspectQuery(args),
           },
         ],

library/sinks/Postgres.ts

@@ -65,6 +65,7 @@ export class Postgres implements Wrapper {
             nodeType: "MethodDefinition",
             name: "query",
             operationKind: "sql_op",
+            bindContext: true,
             inspectArgs: (args) => this.inspectQuery(args),
           },
         ],

scripts/run-tap.js

@@ -30,8 +30,6 @@ if (process.argv.includes("--test-new-instrumentation")) {
     "**/sinks/SQLite3.test.ts",
     "**/sinks/Prisma.test.ts",
     "**/sinks/AwsSDK*",
-    "**/sinks/Postgres.test.ts", // Bind context
-    "**/sinks/MySQL.test.ts", // Bind context
     "**/sinks/MySQL2.v3.10.test.ts",
 
     "**/sources/PubSub.test.ts",