fix(test-runner): split $$ heredocs and escaped backticks in splitQueries()

tests/clickhouse-test-runner/__tests__/split-queries.test.ts

@@ -77,4 +77,68 @@ describe("splitQueries", () => {
       "SELECT 2",
     ]);
   });
+
+  it("does not split on quotes or semicolons inside a $$ heredoc", () => {
+    // Condensed 03032_save_bad_json_escape_sequences: the heredoc body holds
+    // double quotes and backslashes that must not toggle the quote state.
+    const sql = [
+      "select * from format(JSONEachRow, $$",
+      '{"key" : "\\u"}',
+      '{"key" : "\\ud800\\"}',
+      "$$);",
+      "SELECT 2;",
+    ].join("\n");
+    expect(splitQueries(sql)).toEqual([
+      [
+        "select * from format(JSONEachRow, $$",
+        '{"key" : "\\u"}',
+        '{"key" : "\\ud800\\"}',
+        "$$)",
+      ].join("\n"),
+      "SELECT 2",
+    ]);
+  });
+
+  it("does not split on quotes or semicolons inside a $tag$ heredoc", () => {
+    expect(
+      splitQueries("SELECT $doc$it's; a trap$doc$ AS s; SELECT 2"),
+    ).toEqual(["SELECT $doc$it's; a trap$doc$ AS s", "SELECT 2"]);
+  });
+
+  it("keeps the boundaries of statements that follow a heredoc", () => {
+    // 02685_bson2 and friends: the heredoc body itself contains an apostrophe.
+    expect(splitQueries("SELECT $$it's$$; SELECT 2; SELECT 3")).toEqual([
+      "SELECT $$it's$$",
+      "SELECT 2",
+      "SELECT 3",
+    ]);
+  });
+
+  it("treats an unterminated $$ as ordinary characters", () => {
+    expect(splitQueries("SELECT $$; SELECT 2")).toEqual([
+      "SELECT $$",
+      "SELECT 2",
+    ]);
+  });
+
+  it("does not treat other bare words starting with $ as heredocs", () => {
+    expect(splitQueries("SELECT $foo, '$bar'; SELECT 2")).toEqual([
+      "SELECT $foo, '$bar'",
+      "SELECT 2",
+    ]);
+  });
+
+  it("handles a backslash-escaped backtick in a quoted identifier", () => {
+    // 04097_tuple_inside_nullable_subcolumn_descendants: the escaped backtick
+    // used to close the identifier early, leaving the identifier open and
+    // merging the rest of the file into one statement.
+    const sql =
+      "CREATE TABLE t (tup Nullable(Tuple(`a\\`b` Nullable(String)))) ENGINE = Memory;" +
+      " SELECT 1; SELECT 2";
+    expect(splitQueries(sql)).toEqual([
+      "CREATE TABLE t (tup Nullable(Tuple(`a\\`b` Nullable(String)))) ENGINE = Memory",
+      "SELECT 1",
+      "SELECT 2",
+    ]);
+  });
 });

tests/clickhouse-test-runner/src/split-queries.ts

@@ -1,3 +1,34 @@
+/** `[A-Za-z0-9_]`, matching the server lexer's `isWordCharASCII`. */
+function isWordChar(code: number): boolean {
+  return (
+    (code >= 0x61 && code <= 0x7a) || // a-z
+    (code >= 0x41 && code <= 0x5a) || // A-Z
+    (code >= 0x30 && code <= 0x39) || // 0-9
+    code === 0x5f // _
+  );
+}
+
+/**
+ * Returns the exclusive end offset of the dollar-quoted heredoc that starts at
+ * `start` (which points at a `$`), or `-1` when that `$` does not open one.
+ *
+ * Mirrors the server lexer: the opening delimiter is `$`, an optional tag of
+ * word characters, and `$` again (so `$$` and `$tag$`); the heredoc runs to the
+ * next occurrence of the same delimiter. An unterminated delimiter is not a
+ * heredoc, and its `$` is then an ordinary character.
+ */
+function findHeredocEnd(sql: string, start: number): number {
+  const tagEnd = sql.indexOf("$", start + 1);
+  if (tagEnd === -1) return -1;
+  for (let i = start + 1; i < tagEnd; i++) {
+    if (!isWordChar(sql.charCodeAt(i))) return -1;
+  }
+  const delimiter = sql.slice(start, tagEnd + 1);
+  const end = sql.indexOf(delimiter, tagEnd + 1);
+  if (end === -1) return -1;
+  return end + delimiter.length;
+}
+
 export function splitQueries(sql: string): string[] {
   const queries: string[] = [];
   let current = "";
@@ -16,7 +47,7 @@ export function splitQueries(sql: string): string[] {
       continue;
     }
 
-    if ((inSingleQuote || inDoubleQuote) && ch === "\\") {
+    if ((inSingleQuote || inDoubleQuote || inBacktick) && ch === "\\") {
       current += ch;
       escaping = true;
       continue;
@@ -41,6 +72,17 @@ export function splitQueries(sql: string): string[] {
         i = end - 1;
         continue;
       }
+      // Dollar-quoted heredoc: $$ ... $$ or $tag$ ... $tag$. The server lexer
+      // reads the whole region as a single token, so quotes, comment markers
+      // and semicolons inside it must not affect the split.
+      if (ch === "$") {
+        const heredocEnd = findHeredocEnd(sql, i);
+        if (heredocEnd !== -1) {
+          current += sql.slice(i, heredocEnd);
+          i = heredocEnd - 1;
+          continue;
+        }
+      }
     }
 
     if (!inDoubleQuote && !inBacktick && ch === "'") {