fix(plpgsql-parse): comment reinjection fixes — pad stitched body with newlines, skip DECLARE sections during statement matching

Author: pyramation

packages/plpgsql-parse/__tests__/__snapshots__/roundtrip.test.ts.snap

@@ -17,7 +17,7 @@ EXCEPTION
     -- Log the error and return null
     RAISE NOTICE 'Division by zero: % / %', a, b;
     RETURN NULL;
-END;
+END
 $$;"
 `;
 
@@ -31,15 +31,18 @@ DECLARE
   r RECORD;
 BEGIN
   -- Process items in batches
-  FOR r IN SELECT id, data FROM pending_items LIMIT p_batch_size LOOP
-    -- Process each item
-    PERFORM process_item(r.id, r.data);
-    v_processed := v_processed + 1;
+  FOR r IN SELECT
+    id,
+    data
+  FROM pending_items
+  LIMIT p_batch_size LOOP
+      -- Process each item
+      PERFORM process_item(r.id, r.data);
+      v_processed := v_processed + 1;
   END LOOP;
-
   -- Return the count of processed items
   RETURN v_processed;
-END;
+END
 $$;"
 `;
 
@@ -54,7 +57,7 @@ CREATE FUNCTION add_numbers(
 BEGIN
   -- Simple addition
   RETURN a + b;
-END;
+END
 $$;
 
 -- Second function with its own body comments
@@ -65,7 +68,7 @@ CREATE FUNCTION multiply_numbers(
 BEGIN
   -- Multiply the inputs
   RETURN a * b;
-END;
+END
 $$;"
 `;
 
@@ -79,19 +82,20 @@ DECLARE
   v_status text;
 BEGIN
   -- First, calculate the order total
-  SELECT sum(amount) INTO v_total FROM order_items WHERE order_id = p_order_id;
-
+  SELECT sum(amount) INTO v_total
+      FROM order_items
+  WHERE
+    order_id = p_order_id;
   -- Then update the order status
   -- based on the total amount
   IF v_total > 1000 THEN
-    v_status := 'premium';
+      v_status := 'premium';
   ELSE
-    v_status := 'standard';
+      v_status := 'standard';
   END IF;
-
   -- Finally, record the result
-  UPDATE orders SET status = v_status, total = v_total WHERE id = p_order_id;
-END;
+  UPDATE orders SET status = v_status,total = v_total WHERE id = p_order_id;
+END
 $$;"
 `;
 
@@ -105,19 +109,20 @@ DECLARE
 BEGIN
   -- Initialize result
   v_result := 'unknown';
-
   -- Try the main logic
   BEGIN
-    -- Fetch and process
-    SELECT status INTO v_result FROM items WHERE id = p_id;
+      -- Fetch and process
+      SELECT status INTO v_result
+          FROM items
+      WHERE
+        id = p_id;
   EXCEPTION
-    WHEN no_data_found THEN
-      -- Handle missing item
-      v_result := 'not_found';
+      WHEN no_data_found THEN
+        -- Handle missing item
+        v_result := 'not_found';
   END;
-
   RETURN v_result;
-END;
+END
 $$;"
 `;
 
@@ -136,9 +141,12 @@ DECLARE
   v_count integer;
 BEGIN
   -- Count all active users
-  SELECT count(*) INTO v_count FROM users WHERE is_active = true;
+  SELECT count(*) INTO v_count
+      FROM users
+  WHERE
+    is_active = true;
   RETURN v_count;
-END;
+END
 $$;"
 `;
 
@@ -147,15 +155,18 @@ exports[`fixture round-trip tests trigger-function.sql deparsed output matches s
 CREATE FUNCTION audit_trigger() RETURNS trigger LANGUAGE plpgsql AS $$
 BEGIN
   -- Set the updated_at timestamp
-  NEW.updated_at := now();
-
+  new.updated_at := now();
   -- Record the change in audit log
-  IF TG_OP = 'UPDATE' THEN
-    INSERT INTO audit_log (table_name, operation, old_data, new_data)
-    VALUES (TG_TABLE_NAME, TG_OP, row_to_json(OLD), row_to_json(NEW));
+  IF tg_op = 'UPDATE' THEN
+      INSERT INTO audit_log (
+        table_name,
+        operation,
+        old_data,
+        new_data
+      ) VALUES
+        (tg_table_name, tg_op, row_to_json(old), row_to_json(new));
   END IF;
-
-  RETURN NEW;
-END;
+  RETURN new;
+END
 $$;"
 `;

packages/plpgsql-parse/src/deparse.ts

@@ -194,11 +194,23 @@ function reinjectBodyComments(
   let stmtKeyIdx = 0;
   const usedAnchors = new Set<number>();
 
+  // Track DECLARE sections: lines between DECLARE and the following BEGIN
+  // are variable declarations, not statements, and must not consume
+  // statement keyword matches (e.g. a declaration `v_x int;` would
+  // otherwise match an assignment to `v_x`).
+  let inDeclare = false;
+
   for (const line of depLines) {
     const trimmed = line.trim().toUpperCase();
 
+    if (trimmed === 'DECLARE') {
+      inDeclare = true;
+    } else if (inDeclare && trimmed.startsWith('BEGIN')) {
+      inDeclare = false;
+    }
+
     // Try to match this line to the next expected statement
-    if (stmtKeyIdx < stmtKeywords.length) {
+    if (!inDeclare && stmtKeyIdx < stmtKeywords.length) {
       const { lineno, keywords } = stmtKeywords[stmtKeyIdx];
 
       if (lineMatchesKeywords(trimmed, keywords)) {
@@ -574,19 +586,26 @@ function findLastEndLine(lines: string[]): number {
 
 /**
  * Replace the function body in a CREATE FUNCTION AST node.
+ * The body is padded with newlines so it renders as
+ * `AS $$\n<body>\n$$` rather than gluing onto the dollar quotes.
  */
 function stitchBodyIntoAst(createFunctionStmt: any, newBody: string): void {
   if (!createFunctionStmt?.options) return;
 
+  const padded =
+    (newBody.startsWith('\n') ? '' : '\n') +
+    newBody +
+    (newBody.endsWith('\n') ? '' : '\n');
+
   for (const opt of createFunctionStmt.options) {
     if (opt?.DefElem?.defname === 'as') {
       const arg = opt.DefElem.arg;
       if (arg?.List?.items?.[0]?.String) {
-        arg.List.items[0].String.sval = newBody;
+        arg.List.items[0].String.sval = padded;
         return;
       }
       if (arg?.String) {
-        arg.String.sval = newBody;
+        arg.String.sval = padded;
         return;
       }
     }