From e42e0397648d024920ebeb422ea4e7e7649e120e Mon Sep 17 00:00:00 2001 From: Petr Date: Sat, 11 Apr 2026 20:19:26 +0200 Subject: [PATCH] fix: preserve multi-element script[] arrays in sync pull/push roundtrip sync pull/push was collapsing transformation script[] arrays from N elements to 1, causing Snowflake runtime errors (MULTI_STATEMENT_COUNT=1). Root cause: extraction joined all script elements without boundaries, and merge always returned a single-element array. Fix: use a state-machine SQL splitter (split on semicolons respecting strings, comments, dollar-quotes) matching the approach of the old keboola-as-code CLI and the Keboola UI. The SQL file stays clean with no artificial markers. Closes #119 --- plugins/kbagent/.claude-plugin/plugin.json | 2 +- pyproject.toml | 2 +- src/keboola_agent_cli/sync/code_extraction.py | 36 +++-- src/keboola_agent_cli/sync/sql_split.py | 134 +++++++++++++++++ tests/test_sql_split.py | 141 ++++++++++++++++++ tests/test_sync_code_extraction.py | 76 +++++++++- uv.lock | 2 +- 7 files changed, 373 insertions(+), 20 deletions(-) create mode 100644 src/keboola_agent_cli/sync/sql_split.py create mode 100644 tests/test_sql_split.py diff --git a/plugins/kbagent/.claude-plugin/plugin.json b/plugins/kbagent/.claude-plugin/plugin.json index bfa09dca..128b52a1 100644 --- a/plugins/kbagent/.claude-plugin/plugin.json +++ b/plugins/kbagent/.claude-plugin/plugin.json @@ -1,6 +1,6 @@ { "name": "kbagent", - "version": "0.17.4", + "version": "0.17.5", "description": "AI-friendly interface to Keboola Connection projects — explore configs, jobs, lineage, call MCP tools, manage dev branches, and debug SQL in workspaces", "author": { "name": "Keboola", diff --git a/pyproject.toml b/pyproject.toml index f5e52602..8f2a5a56 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "keboola-agent-cli" -version = "0.17.4" +version = "0.17.5" description = "AI-friendly CLI for managing Keboola projects" readme = "README.md" requires-python = ">=3.12" diff --git a/src/keboola_agent_cli/sync/code_extraction.py b/src/keboola_agent_cli/sync/code_extraction.py index 1130be80..3ffc6d96 100644 --- a/src/keboola_agent_cli/sync/code_extraction.py +++ b/src/keboola_agent_cli/sync/code_extraction.py @@ -9,6 +9,8 @@ from pathlib import Path from typing import Any +from keboola_agent_cli.sync.sql_split import split_statements + def _strip_trailing_empty(lines: list[str]) -> list[str]: """Remove trailing empty lines but preserve leading whitespace.""" @@ -18,18 +20,20 @@ def _strip_trailing_empty(lines: list[str]) -> list[str]: return result -def _lines_to_script(lines: list[str]) -> list[str]: - """Convert collected lines into a single-string script element. +def _lines_to_script(lines: list[str], *, is_sql: bool = False) -> list[str]: + """Convert collected lines back into the ``script[]`` array. - The Keboola transformation runner treats each element of the ``script`` - array as a separate executable statement. Joining all lines of a CODE - block into one string ensures multi-line SQL/Python is executed as a - single unit. + For SQL transformations: splits on semicolons using a state machine, + producing one element per statement (matching Keboola runtime semantics). + For Python/other: joins all lines into a single element. """ stripped = _strip_trailing_empty(lines) if not stripped: return [] - return ["\n".join(stripped)] + content = "\n".join(stripped) + if is_sql: + return split_statements(content) + return [content] # Component patterns that contain SQL transformations @@ -137,7 +141,9 @@ def _extract_sql_transformation(config_data: dict[str, Any], config_dir: Path) - lines.append(SQL_CODE_MARKER.format(name=code_name)) scripts = code.get("script") or [] - for script in scripts: + for si, script in enumerate(scripts): + if si > 0: + lines.append("") # blank line between statements if isinstance(script, str) and "\n" in script: lines.extend(script.split("\n")) else: @@ -184,7 +190,7 @@ def _parse_sql_blocks(content: str) -> list[dict[str, Any]]: if stripped.startswith("/* ===== BLOCK:") and stripped.endswith("===== */"): # Save previous code if any if current_code is not None and current_block is not None: - current_code["script"] = _lines_to_script(current_script_lines) + current_code["script"] = _lines_to_script(current_script_lines, is_sql=True) current_block.setdefault("codes", []).append(current_code) current_code = None current_script_lines = [] @@ -198,7 +204,7 @@ def _parse_sql_blocks(content: str) -> list[dict[str, Any]]: if stripped.startswith("/* ===== CODE:") and stripped.endswith("===== */"): # Save previous code if any if current_code is not None and current_block is not None: - current_code["script"] = _lines_to_script(current_script_lines) + current_code["script"] = _lines_to_script(current_script_lines, is_sql=True) current_block.setdefault("codes", []).append(current_code) current_script_lines = [] @@ -212,7 +218,7 @@ def _parse_sql_blocks(content: str) -> list[dict[str, Any]]: # Don't forget the last code block if current_code is not None and current_block is not None: - current_code["script"] = _lines_to_script(current_script_lines) + current_code["script"] = _lines_to_script(current_script_lines, is_sql=True) current_block.setdefault("codes", []).append(current_code) # If no markers found, treat entire content as single block/code @@ -220,7 +226,9 @@ def _parse_sql_blocks(content: str) -> list[dict[str, Any]]: blocks = [ { "name": "Block 1", - "codes": [{"name": "Code 1", "script": _lines_to_script(content.split("\n"))}], + "codes": [ + {"name": "Code 1", "script": _lines_to_script(content.split("\n"), is_sql=True)} + ], } ] @@ -247,7 +255,9 @@ def _extract_python_transformation(config_data: dict[str, Any], config_dir: Path lines.append(PYTHON_CODE_MARKER.format(name=code_name)) scripts = code.get("script") or [] - for script in scripts: + for si, script in enumerate(scripts): + if si > 0: + lines.append("") # blank line between script elements if isinstance(script, str) and "\n" in script: lines.extend(script.split("\n")) else: diff --git a/src/keboola_agent_cli/sync/sql_split.py b/src/keboola_agent_cli/sync/sql_split.py new file mode 100644 index 00000000..88a7339b --- /dev/null +++ b/src/keboola_agent_cli/sync/sql_split.py @@ -0,0 +1,134 @@ +"""Split and join SQL statements using a state machine. + +Splits on semicolons while respecting: +- Single-quoted strings ('...') +- Double-quoted identifiers ("...") +- Dollar-quoted blocks ($$...$$) +- Line comments (--, #, //) +- Block comments (/* ... */) + +Compatible with the Keboola UI splitter and the old keboola-as-code CLI. +""" + +from __future__ import annotations + +from enum import Enum, auto + + +class _State(Enum): + NORMAL = auto() + SINGLE_QUOTE = auto() + DOUBLE_QUOTE = auto() + DOLLAR_QUOTE = auto() + LINE_COMMENT = auto() + BLOCK_COMMENT = auto() + + +def split_statements(sql: str) -> list[str]: + """Split SQL text into individual statements on semicolons. + + Returns a list of stripped, non-empty statements. Trailing + semicolons are preserved on each statement. + """ + sql = sql.rstrip() + if not sql: + return [] + + state = _State.NORMAL + statements: list[str] = [] + buf: list[str] = [] + i = 0 + n = len(sql) + + while i < n: + ch = sql[i] + nxt = sql[i + 1] if i + 1 < n else "" + two = ch + nxt + + if state == _State.NORMAL: + if ch == ";": + buf.append(ch) + stmt = "".join(buf).strip() + if stmt.strip(";"): + statements.append(stmt) + buf = [] + elif ch == "'": + buf.append(ch) + state = _State.SINGLE_QUOTE + elif ch == '"': + buf.append(ch) + state = _State.DOUBLE_QUOTE + elif two == "$$": + buf.append(two) + state = _State.DOLLAR_QUOTE + i += 1 + elif two in ("--", "//"): + buf.append(two) + state = _State.LINE_COMMENT + i += 1 + elif ch == "#": + buf.append(ch) + state = _State.LINE_COMMENT + elif two == "/*": + buf.append(two) + state = _State.BLOCK_COMMENT + i += 1 + else: + buf.append(ch) + + elif state == _State.SINGLE_QUOTE: + if ch == "\\" and nxt: + buf.append(two) + i += 1 + elif ch == "'": + buf.append(ch) + state = _State.NORMAL + else: + buf.append(ch) + + elif state == _State.DOUBLE_QUOTE: + if ch == "\\" and nxt: + buf.append(two) + i += 1 + elif ch == '"': + buf.append(ch) + state = _State.NORMAL + else: + buf.append(ch) + + elif state == _State.DOLLAR_QUOTE: + if two == "$$": + buf.append(two) + state = _State.NORMAL + i += 1 + else: + buf.append(ch) + + elif state == _State.LINE_COMMENT: + buf.append(ch) + if ch == "\n": + state = _State.NORMAL + + elif state == _State.BLOCK_COMMENT: + if two == "*/": + buf.append(two) + state = _State.NORMAL + i += 1 + else: + buf.append(ch) + + i += 1 + + # Remaining content without trailing semicolon + remaining = "".join(buf).strip() + if remaining: + statements.append(remaining) + + return statements + + +def join_statements(statements: list[str]) -> str: + """Join SQL statements with double newlines (matching Keboola convention).""" + if not statements: + return "" + return "\n\n".join(s.rstrip() for s in statements) diff --git a/tests/test_sql_split.py b/tests/test_sql_split.py new file mode 100644 index 00000000..7e2c4ff1 --- /dev/null +++ b/tests/test_sql_split.py @@ -0,0 +1,141 @@ +"""Tests for the SQL statement splitter state machine. + +Test cases ported from keboola-as-code (Go) sql_test.go and extended. +""" + +import pytest + +from keboola_agent_cli.sync.sql_split import join_statements, split_statements + + +class TestSplitStatements: + """Tests for split_statements state machine.""" + + def test_empty(self) -> None: + assert split_statements("") == [] + + def test_whitespace_only(self) -> None: + assert split_statements(" \n\n\n ") == [] + + def test_one_statement(self) -> None: + assert split_statements("SELECT * FROM bar") == ["SELECT * FROM bar"] + + def test_one_statement_with_semicolon(self) -> None: + assert split_statements("SELECT * FROM bar;") == ["SELECT * FROM bar;"] + + def test_one_statement_whitespace_padding(self) -> None: + assert split_statements(" \n\n\nSELECT * FROM bar\t\n ") == ["SELECT * FROM bar"] + + def test_multiple_statements(self) -> None: + sql = "SELECT 1;\nINSERT INTO bar VALUES('x', 'y');\nTRUNCATE records;" + result = split_statements(sql) + assert result == [ + "SELECT 1;", + "INSERT INTO bar VALUES('x', 'y');", + "TRUNCATE records;", + ] + + def test_multiple_with_extra_whitespace(self) -> None: + sql = " \n\n\nSELECT * FROM [bar];\t\n INSERT INTO bar VALUES('x', 'y'); TRUNCATE records;;;" + result = split_statements(sql) + assert result == [ + "SELECT * FROM [bar];", + "INSERT INTO bar VALUES('x', 'y');", + "TRUNCATE records;", + ] + + def test_split_simple_queries(self) -> None: + sql = "SELECT 1;\nSelect 2;\nSELECT 3;" + result = split_statements(sql) + assert result == ["SELECT 1;", "Select 2;", "SELECT 3;"] + + def test_block_comment(self) -> None: + sql = "SELECT 1;\n/*\n Select 2;\n*/\nSELECT 3;" + result = split_statements(sql) + assert result == ["SELECT 1;", "/*\n Select 2;\n*/\nSELECT 3;"] + + def test_line_comment_dash(self) -> None: + sql = "SELECT 1;\n-- Select 2;\nSELECT 3;" + result = split_statements(sql) + assert result == ["SELECT 1;", "-- Select 2;\nSELECT 3;"] + + def test_line_comment_hash(self) -> None: + sql = "SELECT 1;\n# Select 2;\nSELECT 3;" + result = split_statements(sql) + assert result == ["SELECT 1;", "# Select 2;\nSELECT 3;"] + + def test_line_comment_double_slash(self) -> None: + sql = "SELECT 1;\n// Select 2;\nSELECT 3;" + result = split_statements(sql) + assert result == ["SELECT 1;", "// Select 2;\nSELECT 3;"] + + def test_dollar_quoted_block(self) -> None: + sql = "SELECT 1;\nexecute immediate $$\n SELECT 2;\n SELECT 3;\n$$;" + result = split_statements(sql) + assert result == [ + "SELECT 1;", + "execute immediate $$\n SELECT 2;\n SELECT 3;\n$$;", + ] + + def test_single_quoted_string_with_semicolon(self) -> None: + sql = "SELECT 'hello; world';\nSELECT 2;" + result = split_statements(sql) + assert result == ["SELECT 'hello; world';", "SELECT 2;"] + + def test_double_quoted_identifier_with_semicolon(self) -> None: + sql = 'SELECT "col;name" FROM t;\nSELECT 2;' + result = split_statements(sql) + assert result == ['SELECT "col;name" FROM t;', "SELECT 2;"] + + def test_escaped_quote_in_string(self) -> None: + sql = "SELECT 'it\\'s a test; yes';\nSELECT 2;" + result = split_statements(sql) + assert result == ["SELECT 'it\\'s a test; yes';", "SELECT 2;"] + + def test_no_trailing_semicolon(self) -> None: + sql = "SELECT 1;\nSELECT 2" + result = split_statements(sql) + assert result == ["SELECT 1;", "SELECT 2"] + + def test_multiline_create_table(self) -> None: + sql = "CREATE TABLE foo AS\n SELECT col1\n FROM bar;\n\nINSERT INTO foo VALUES (1);" + result = split_statements(sql) + assert result == [ + "CREATE TABLE foo AS\n SELECT col1\n FROM bar;", + "INSERT INTO foo VALUES (1);", + ] + + +class TestJoinStatements: + """Tests for join_statements.""" + + def test_empty(self) -> None: + assert join_statements([]) == "" + + def test_single(self) -> None: + assert join_statements(["SELECT 1;"]) == "SELECT 1;" + + def test_multiple(self) -> None: + result = join_statements(["SELECT 1;", "SELECT 2;", "SELECT 3;"]) + assert result == "SELECT 1;\n\nSELECT 2;\n\nSELECT 3;" + + def test_strips_trailing_whitespace(self) -> None: + result = join_statements(["SELECT 1; ", "SELECT 2;\n"]) + assert result == "SELECT 1;\n\nSELECT 2;" + + +class TestRoundTrip: + """Test that split -> join -> split is idempotent.""" + + @pytest.mark.parametrize( + "statements", + [ + ["SELECT 1;", "SELECT 2;", "SELECT 3;"], + ["CREATE TABLE foo AS\n SELECT col1\n FROM bar;", "INSERT INTO foo VALUES (1);"], + ["SELECT 'semicolon; inside';", "SELECT 2;"], + ], + ) + def test_split_join_roundtrip(self, statements: list[str]) -> None: + joined = join_statements(statements) + split_back = split_statements(joined) + assert split_back == statements diff --git a/tests/test_sync_code_extraction.py b/tests/test_sync_code_extraction.py index 80d342ed..222db49a 100644 --- a/tests/test_sync_code_extraction.py +++ b/tests/test_sync_code_extraction.py @@ -226,10 +226,10 @@ def test_sql_merge_no_markers(self, tmp_path: Path) -> None: assert blocks[0]["name"] == "Block 1" assert blocks[0]["codes"][0]["name"] == "Code 1" script = blocks[0]["codes"][0]["script"] - # Must be a single joined string, not per-line - assert len(script) == 1 - assert "SELECT 1;" in script[0] - assert "SELECT 2;" in script[0] + # SQL splitter splits on semicolons -> 2 statements + assert len(script) == 2 + assert script[0] == "SELECT 1;" + assert script[1] == "SELECT 2;" def test_multiline_sql_produces_single_string(self, tmp_path: Path) -> None: """Multi-line SQL statement is joined into a single script element.""" @@ -330,6 +330,74 @@ def test_whitespace_only_code_block(self, tmp_path: Path) -> None: assert codes[0]["name"] == "Spaces" assert codes[0]["script"] == [] + def test_sql_roundtrip_multi_script(self, tmp_path: Path) -> None: + """Round-trip preserves multi-element script[] arrays (issue #119).""" + config_data = { + "parameters": { + "blocks": [ + { + "name": "Block 1", + "codes": [ + { + "name": "multi-stmt", + "script": [ + "CREATE OR REPLACE TABLE a AS SELECT 1;", + "INSERT INTO a VALUES (2);", + "UPDATE a SET x = 3 WHERE x = 2;", + ], + } + ], + } + ], + }, + } + original_scripts = copy.deepcopy( + config_data["parameters"]["blocks"][0]["codes"][0]["script"] + ) + config_dir = tmp_path / "sql-multi-script" + + extract_code_files("keboola.snowflake-transformation", config_data, config_dir) + + # File should be clean SQL with no artificial markers + content = (config_dir / "transform.sql").read_text(encoding="utf-8") + assert "STATEMENT" not in content + + merge_code_files("keboola.snowflake-transformation", config_data, config_dir) + + scripts = config_data["parameters"]["blocks"][0]["codes"][0]["script"] + assert scripts == original_scripts + + def test_sql_roundtrip_multiline_multi_script(self, tmp_path: Path) -> None: + """Round-trip with multi-line statements in multi-element script[].""" + config_data = { + "parameters": { + "blocks": [ + { + "name": "Block 1", + "codes": [ + { + "name": "complex", + "script": [ + "CREATE TABLE foo AS\n SELECT col1\n FROM bar;", + "INSERT INTO foo\n SELECT col2\n FROM baz;", + ], + } + ], + } + ], + }, + } + original_scripts = copy.deepcopy( + config_data["parameters"]["blocks"][0]["codes"][0]["script"] + ) + config_dir = tmp_path / "sql-multi-multiline" + + extract_code_files("keboola.snowflake-transformation", config_data, config_dir) + merge_code_files("keboola.snowflake-transformation", config_data, config_dir) + + scripts = config_data["parameters"]["blocks"][0]["codes"][0]["script"] + assert scripts == original_scripts + # =================================================================== # Python Transformation Tests diff --git a/uv.lock b/uv.lock index 01f80a68..dda3d789 100644 --- a/uv.lock +++ b/uv.lock @@ -408,7 +408,7 @@ wheels = [ [[package]] name = "keboola-agent-cli" -version = "0.17.1" +version = "0.17.5" source = { editable = "." } dependencies = [ { name = "httpx" },