Wrap node-postgres

Author: timokoessler

instrumentation-wasm/src/js_transformer/transformer.rs

@@ -1,5 +1,5 @@
-use oxc_allocator::Allocator;
-use oxc_ast::ast::MethodDefinition;
+use oxc_allocator::{Allocator, Box};
+use oxc_ast::ast::{FunctionBody, MethodDefinition};
 use oxc_codegen::{Codegen, CodegenOptions};
 use oxc_parser::{ParseOptions, Parser};
 use oxc_semantic::SemanticBuilder;
@@ -48,9 +48,6 @@ pub fn transform_code_str(code: &str, instructions_json: &str, src_type: i32) ->
     let t = &mut Transformer {
         allocator: &allocator,
         file_instructions: &file_instructions,
-        //current_function_identifier: None,
-        //modify_return_value: false,
-        //sub_function_counter: 0,
     };
 
     traverse_mut(t, &allocator, program, symbols, scopes);
@@ -77,18 +74,12 @@ pub fn transform_code_str(code: &str, instructions_json: &str, src_type: i32) ->
         })
         .build(&program);
 
-    // Debug helper
-    //format!("{:?}", parser_result.program);
-
     js.code
 }
 
 struct Transformer<'a> {
     allocator: &'a Allocator,
     file_instructions: &'a FileInstructions,
-    //current_function_identifier: Option<String>, // Only set if we want to instrument the current function
-    //sub_function_counter: i32, // Counter to keep track of how many sub functions we are in
-    //modify_return_value: bool,
 }
 
 impl<'a> Traverse<'a> for Transformer<'a> {
@@ -106,8 +97,6 @@ impl<'a> Traverse<'a> for Transformer<'a> {
             return;
         }
 
-        // Todo implement submethod counting for nested functions for supporting modifications of return value
-
         let method_name = node.key.name().unwrap().to_string();
 
         let found_instruction = self
@@ -117,17 +106,11 @@ impl<'a> Traverse<'a> for Transformer<'a> {
             .find(|f| f.node_type == "MethodDefinition" && f.name == method_name);
 
         if found_instruction.is_none() {
-            //self.current_function_identifier = None;
-            //self.modify_return_value = false;
             return;
         }
 
         let instruction = found_instruction.unwrap();
 
-        /*let function_identifier = format!("{}.{}", self.module_name, method_name);
-        self.current_function_identifier = Some(function_identifier.clone());
-        self.modify_return_value = instruction.modify_return_value;*/
-
         // We need to collect the arg names before we make the body mutable
         let arg_names = if instruction.modify_args {
             // Todo check whats appens if rest parameter is used
@@ -149,13 +132,7 @@ impl<'a> Traverse<'a> for Transformer<'a> {
                 arg_names_str, instruction.identifier, arg_names_str
             ));
 
-            body.statements.insert(
-                0,
-                parse_js_code_to_statements(self.allocator, &source_text, SourceType::mjs())
-                    .into_iter()
-                    .next()
-                    .unwrap(),
-            );
+            insert_single_statement_into_func(self.allocator, body, 0, &source_text);
         }
 
         if instruction.inspect_args {
@@ -165,52 +142,24 @@ impl<'a> Traverse<'a> for Transformer<'a> {
                 instruction.identifier
             ));
 
-            body.statements.insert(
-                0,
-                parse_js_code_to_statements(self.allocator, &source_text, SourceType::mjs())
-                    .into_iter()
-                    .next()
-                    .unwrap(),
-            );
-        }
-    }
-
-    /*fn exit_method_definition(
-        &mut self,
-        _node: &mut MethodDefinition<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
-        if self.current_function_identifier.is_none() {
-            return;
-        }
-
-        // If we are in a sub function, we need to decrease the counter
-        // If we leave the last sub function, we need to modify the return value again
-        if self.sub_function_counter > 0 {
-            self.sub_function_counter -= 1;
-
-            if self.sub_function_counter == 0 {
-                self.modify_return_value = true;
-            }
-
-            return;
+            insert_single_statement_into_func(self.allocator, body, 0, source_text);
         }
 
-        // We leave the current function we want to instrument
-        self.modify_return_value = false;
-        self.current_function_identifier = None;
+        // Todo support return value modification
     }
+}
 
-    fn enter_return_statement(
-        &mut self,
-        node: &mut oxc_ast::ast::ReturnStatement<'a>,
-        ctx: &mut TraverseCtx<'a>,
-    ) {
-        if !self.modify_return_value {
-            return;
-        }
-
-        // Todo support modifying return value
-        //AstBuilder::new(self.allocator).return_statement(node.span, argument)
-    }*/
+fn insert_single_statement_into_func<'a>(
+    allocator: &'a Allocator,
+    body: &mut Box<'a, FunctionBody<'a>>,
+    pos: usize,
+    source_text: &'a str,
+) {
+    body.statements.insert(
+        pos,
+        parse_js_code_to_statements(&allocator, &source_text, SourceType::mjs())
+            .into_iter()
+            .next()
+            .unwrap(),
+    );
 }

library/agent/hooks/instrumentation/types.ts

@@ -59,15 +59,17 @@ export type IntereptorFunctionsObj = {
   modifyReturnValue?: ModifyReturnValueInterceptor;
 };
 
+export type PackageFunctionInstrumentationInstruction = {
+  nodeType: "MethodDefinition";
+  name: string;
+  inspectArgs?: InspectArgsInterceptor;
+  modifyArgs?: ModifyArgsInterceptor;
+  modifyReturnValue?: ModifyReturnValueInterceptor;
+};
+
 export type PackageFileInstrumentationInstruction = {
   path: string; // Relative path to required file inside the package folder
-  functions: {
-    nodeType: "MethodDefinition";
-    name: string;
-    inspectArgs?: InspectArgsInterceptor;
-    modifyArgs?: ModifyArgsInterceptor;
-    modifyReturnValue?: ModifyReturnValueInterceptor;
-  }[];
+  functions: PackageFunctionInstrumentationInstruction[];
 };
 
 export type PackageFileInstrumentationInstructionJSON = {

library/sinks/Postgres.ts

@@ -56,6 +56,16 @@ export class Postgres implements Wrapper {
         wrapExport(exports.Client.prototype, "query", pkgInfo, {
           inspectArgs: (args) => this.inspectQuery(args),
         });
+      })
+      .addFileInstrumentation({
+        path: "lib/client.js",
+        functions: [
+          {
+            nodeType: "MethodDefinition",
+            name: "query",
+            inspectArgs: (args) => this.inspectQuery(args),
+          },
+        ],
       });
   }
 }

library/sources/Hono.ts

@@ -5,6 +5,7 @@ import { Wrapper } from "../agent/Wrapper";
 import { wrapRequestHandler } from "./hono/wrapRequestHandler";
 import { wrapExport } from "../agent/hooks/wrapExport";
 import { wrapNewInstance } from "../agent/hooks/wrapNewInstance";
+import type { PackageFunctionInstrumentationInstruction } from "../agent/hooks/instrumentation/types";
 
 const METHODS = [
   "get",
@@ -38,6 +39,17 @@ export class Hono implements Wrapper {
   }
 
   wrap(hooks: Hooks) {
+    // Instrumentation instructions for the new system
+    const functionsToInstrument: PackageFunctionInstrumentationInstruction[] = [
+      {
+        nodeType: "MethodDefinition",
+        name: "addRoute",
+        modifyArgs: (args) => {
+          return this.wrapArgs(args);
+        },
+      },
+    ];
+
     hooks
       .addPackage("hono")
       .withVersion("^4.0.0")
@@ -56,15 +68,11 @@ export class Hono implements Wrapper {
       })
       .addFileInstrumentation({
         path: "dist/hono-base.js", //CJS has a different path
-        functions: [
-          {
-            nodeType: "MethodDefinition",
-            name: "addRoute",
-            modifyArgs: (args) => {
-              return this.wrapArgs(args);
-            },
-          },
-        ],
+        functions: functionsToInstrument,
+      })
+      .addFileInstrumentation({
+        path: "dist/cjs/hono-base.js",
+        functions: functionsToInstrument,
       });
   }
 }

sample-apps/hono-esm/app.js

@@ -0,0 +1,47 @@
+import { Hono } from "hono";
+import { serve } from "@hono/node-server";
+import { readFile } from "fs/promises";
+import { join } from "path";
+import { createConnection } from "./db.js";
+
+const app = new Hono();
+const db = await createConnection();
+
+app.get("/", async (c) => {
+  return c.text("Hello, World!");
+});
+
+app.get("/file", async (c) => {
+  const path = c.req.query("path");
+  const data = await readFile(
+    join(import.meta.dirname, "files", path || "test.txt"),
+    "utf8"
+  );
+  return c.text(data);
+});
+
+app.get("/clear", async (c) => {
+  await db.query("DELETE FROM cats_3;");
+  return c.text("Table cleared");
+});
+
+function getPort() {
+  const port = parseInt(process.argv[2], 10) || 4000;
+
+  if (isNaN(port)) {
+    console.error("Invalid port");
+    process.exit(1);
+  }
+
+  return port;
+}
+
+serve(
+  {
+    fetch: app.fetch,
+    port: getPort(),
+  },
+  (info) => {
+    console.log(`Server is running on http://${info.address}:${info.port}`);
+  }
+);

sample-apps/hono-esm/package-lock.json

@@ -0,0 +1,22 @@
+import pg from "pg";
+const { Client } = pg;
+
+export async function createConnection() {
+  const client = new Client({
+    user: "root",
+    host: "127.0.0.1",
+    database: "main_db",
+    password: "password",
+    port: 27016,
+  });
+
+  await client.connect();
+  await client.query(`
+    CREATE TABLE IF NOT EXISTS cats_3 (
+        petname varchar(255),
+        comment varchar(255)
+    );
+  `);
+
+  return client;
+}

sample-apps/hono-pg-esm/app.js

@@ -0,0 +1 @@
+Hi