ISSUE-2103 Handle trailing and trailing statement comments

cli/api/BUILD

@@ -1,12 +1,16 @@
 package(default_visibility = ["//visibility:public"])
 
 load("//tools:ts_library.bzl", "ts_library")
+load("//testing:index.bzl", "ts_test_suite")
 
 ts_library(
     name = "api",
     srcs = glob(
         ["**/*.ts"],
-        exclude = ["utils/**/*.*"],
+        exclude = [
+            "utils/**/*.*",
+            "**/*_test.ts",
+        ],
     ),
     deps = [
         "//cli/api/utils",
@@ -42,3 +46,15 @@ ts_library(
         "@npm//tmp",
     ],
 )
+
+ts_test_suite(
+    name = "tests",
+    srcs = ["tasks_test.ts"],
+    deps = [
+        "//cli/api",
+        "//protos:ts",
+        "//testing",
+        "@npm//@types/chai",
+        "@npm//chai",
+    ],
+)

cli/api/dbadapters/tasks.ts

@@ -1,16 +1,44 @@
 import { dataform } from "df/protos/ts";
 
 export function concatenateQueries(statements: string[], modifier?: (mod: string) => string) {
-  return statements
+  const processed = statements
     .filter(statement => !!statement)
     .map(statement => statement.trim())
-    .map(statement =>
-      statement.length > 0 && statement.charAt(statement.length - 1) === ";"
-        ? statement.substring(0, statement.length - 1)
-        : statement
-    )
-    .map(statement => (!!modifier ? modifier(statement) : statement))
-    .join("\n;\n");
+    .filter(statement => statement.length > 0);
+
+  return processed
+    .map((statement, index) => {
+      const isLast = index === processed.length - 1;
+      const commentIndex = statement.indexOf("--");
+
+      let code: string;
+      let comment: string;
+
+      if (commentIndex === 0) {
+        code = "";
+        comment = statement;
+      } else if (commentIndex > 0) {
+        code = statement.substring(0, commentIndex).trimEnd();
+        comment = statement.substring(commentIndex);
+      } else {
+        code = statement;
+        comment = "";
+      }
+
+      if (modifier && code.length > 0) {
+        code = modifier(code);
+      }
+
+      if (!isLast && code.length > 0 && !code.endsWith(";")) {
+        code = code + ";";
+      }
+
+      if (code.length > 0 && comment.length > 0) {
+        return code + " " + comment;
+      }
+      return code || comment;
+    })
+    .join("\n");
 }
 
 export class Tasks {

cli/api/tasks_test.ts

@@ -0,0 +1,77 @@
+import { expect } from "chai";
+
+import { concatenateQueries } from "df/cli/api/dbadapters/tasks";
+import { suite, test } from "df/testing";
+
+suite("concatenateQueries", () => {
+  test("adds semicolons to non-last statements and joins with newlines", () => {
+    expect(concatenateQueries(["SELECT 1", "SELECT 2"])).deep.equals(
+      "SELECT 1;\nSELECT 2"
+    );
+  });
+
+  test("does not add semicolon to the last statement", () => {
+    expect(concatenateQueries(["SELECT 1", "SELECT 2", "SELECT 3"])).deep.equals(
+      "SELECT 1;\nSELECT 2;\nSELECT 3"
+    );
+  });
+
+  test("preserves existing semicolons on non-last statements", () => {
+    expect(concatenateQueries(["SELECT 1;", "SELECT 2"])).deep.equals(
+      "SELECT 1;\nSELECT 2"
+    );
+  });
+
+  test("preserves existing semicolons on last statement", () => {
+    expect(concatenateQueries(["SELECT 1;", "SELECT 2;"])).deep.equals(
+      "SELECT 1;\nSELECT 2;"
+    );
+  });
+
+  test("adds semicolon before trailing comment if necessary, verifies fix for https://github.com/dataform-co/dataform/issues/1231", () => {
+    expect(concatenateQueries(["SELECT 1 -- blah", "drop table foo; -- baz", "SELECT 2"])).deep.equals(
+      "SELECT 1; -- blah\ndrop table foo; -- baz\nSELECT 2"
+    );
+  });
+
+  test("preserves comment-only lines, verifies fix for https://github.com/dataform-co/dataform/issues/2103", () => {
+    expect(concatenateQueries(["SELECT 1;", "--drop table foo", "SELECT 2", "--baz"])).deep.equals(
+      "SELECT 1;\n--drop table foo\nSELECT 2;\n--baz"
+    );
+  });
+
+  test("trims whitespace from statements", () => {
+    expect(concatenateQueries(["  SELECT 1  ", "  SELECT 2  ", "  SELECT 3 -- blah"])).deep.equals(
+      "SELECT 1;\nSELECT 2;\nSELECT 3 -- blah"
+    );
+  });
+
+  test("filters out empty and falsy statements", () => {
+    expect(concatenateQueries(["SELECT 1", "", "SELECT 2"])).deep.equals(
+      "SELECT 1;\nSELECT 2"
+    );
+    expect(concatenateQueries(["SELECT 1", undefined, "SELECT 2"])).deep.equals(
+      "SELECT 1;\nSELECT 2"
+    );
+  });
+
+  test("returns empty string for empty input", () => {
+    expect(concatenateQueries([])).deep.equals("");
+  });
+
+  test("returns single statement without semicolon", () => {
+    expect(concatenateQueries(["SELECT 1"])).deep.equals("SELECT 1");
+  });
+
+  test("applies modifier to code portion of each statement", () => {
+    expect(
+      concatenateQueries(["SELECT 1", "SELECT 2"], stmt => `(${stmt})`)
+    ).deep.equals("(SELECT 1);\n(SELECT 2)");
+  });
+
+  test("applies modifier only to code portion, not comment", () => {
+    expect(
+      concatenateQueries(["SELECT 1 --comment", "SELECT 2"], stmt => `(${stmt})`)
+    ).deep.equals("(SELECT 1); --comment\n(SELECT 2)");
+  });
+});

tests/integration/bigquery.spec.ts

@@ -384,7 +384,7 @@ suite("@dataform/integration/bigquery", { parallel: true }, ({ before, after })
       const bqadapter = new ExecutionSql({ warehouse: "bigquery" }, "1.4.8");
 
       const refresh = bqadapter.publishTasks(table, { fullRefresh: true }, { fields: [] }).build();
-      const splitRefresh = refresh[0].statement.split("\n;\n");
+      const splitRefresh = refresh[0].statement.split("\n");
       expect([...splitRefresh.slice(0, 2), ...splitRefresh.slice(-2)]).to.eql([
         ...table.preOps,
         ...table.postOps
@@ -394,7 +394,7 @@ suite("@dataform/integration/bigquery", { parallel: true }, ({ before, after })
         .publishTasks(table, { fullRefresh: false }, { fields: [] })
         .build();
 
-      const splitIncrement = increment[0].statement.split("\n;\n");
+      const splitIncrement = increment[0].statement.split("\n");
       expect([...splitIncrement.slice(0, 2), ...splitIncrement.slice(-2)]).to.eql([
         ...table.preOps,
         ...table.postOps

tests/integration/bigquery_project/definitions/example_trailing_comment_in_operations.sqlx

@@ -0,0 +1,13 @@
+config {
+  "type": "table",
+  "description": "Verifies fix for https://github.com/dataform-co/dataform/issues/2103"
+}
+
+SELECT 1 AS foo
+
+pre_operations {
+
+  DECLARE bar DATE DEFAULT DATE_TRUNC(CURRENT_DATE(), WEEK(MONDAY));
+
+  -- DECLARE bar DATE DEFAULT DATE_TRUNC(DATE('2026-01-01'), WEEK(MONDAY))
+}

tests/integration/bigquery_project/definitions/example_trailing_statement_comment_with_operations.sqlx

@@ -0,0 +1,12 @@
+config { 
+    type: "table", 
+    schema: "test", 
+    name: "testtab",
+    description: "Verifies fix for https://github.com/dataform-co/dataform/issues/1231" 
+} 
+
+post_operations { 
+    drop table ${self()} 
+}
+
+select 1 as col1 -- test comment