From 2e8dd2cafe25cb63386cb651e710a27ef225639e Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 7 Aug 2026 00:16:41 -0400 Subject: [PATCH 1/3] refactor(snowflake): strip trailing semicolon on unlimited query path Keep LIMIT pushdown/subquery rationale comments and coerce_limit. Align unlimited execute with dry_run and other connectors (#2595). --- core/wren/src/wren/connector/snowflake.py | 6 ++++-- .../wren/tests/unit/test_snowflake_limit_pushdown.py | 12 ++++++++++++ 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/core/wren/src/wren/connector/snowflake.py b/core/wren/src/wren/connector/snowflake.py index 842c0d5478..a5dc9ffc54 100644 --- a/core/wren/src/wren/connector/snowflake.py +++ b/core/wren/src/wren/connector/snowflake.py @@ -56,18 +56,20 @@ def __init__(self, connection_info): def query(self, sql: str, limit: int | None = None) -> pa.Table: limit = coerce_limit(limit) + # Align unlimited execute with dry_run and other connectors (mysql/ + # bigquery/duckdb/redshift): strip a terminating `;` before send. + executed = strip_trailing_semicolon(sql) # Push LIMIT into Snowflake when requested so we do not download a # full result set only to slice it in Python. Wrap as a subquery so a # trailing semicolon in the user SQL cannot break composition, and so # statements that already contain an ORDER BY keep their ordering # under the outer LIMIT. - executed = sql if limit is not None: # Place the user SQL on its own line so a trailing line comment # (`-- ...`) cannot swallow the closing paren, alias, or LIMIT. executed = ( "SELECT * FROM (\n" - f"{strip_trailing_semicolon(sql)}\n" + f"{executed}\n" f") AS _wren_sub LIMIT {limit}" ) try: diff --git a/core/wren/tests/unit/test_snowflake_limit_pushdown.py b/core/wren/tests/unit/test_snowflake_limit_pushdown.py index 910b07ac49..f4f7ab0972 100644 --- a/core/wren/tests/unit/test_snowflake_limit_pushdown.py +++ b/core/wren/tests/unit/test_snowflake_limit_pushdown.py @@ -60,6 +60,18 @@ def test_query_without_limit_runs_original_sql(): cursor.execute.assert_called_once_with("SELECT 1") +def test_query_without_limit_strips_trailing_semicolon(): + connector = SnowflakeConnector.__new__(SnowflakeConnector) + connector.connection = MagicMock() + cursor = MagicMock() + connector.connection.cursor.return_value.__enter__.return_value = cursor + cursor.fetch_arrow_all.return_value = pa.table({}) + + connector.query("SELECT 1;") + + cursor.execute.assert_called_once_with("SELECT 1") + + def test_dry_run_strips_trailing_semicolon_before_describe(): connector = SnowflakeConnector.__new__(SnowflakeConnector) connector.connection = MagicMock() From 9fa78bea35252269c9f165a7787dbab77fcb9749 Mon Sep 17 00:00:00 2001 From: Bartok9 Date: Fri, 7 Aug 2026 00:20:10 -0400 Subject: [PATCH 2/3] style(snowflake): ruff-format limit wrap string --- core/wren/src/wren/connector/snowflake.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/core/wren/src/wren/connector/snowflake.py b/core/wren/src/wren/connector/snowflake.py index a5dc9ffc54..79ff7aabd5 100644 --- a/core/wren/src/wren/connector/snowflake.py +++ b/core/wren/src/wren/connector/snowflake.py @@ -67,11 +67,7 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: if limit is not None: # Place the user SQL on its own line so a trailing line comment # (`-- ...`) cannot swallow the closing paren, alias, or LIMIT. - executed = ( - "SELECT * FROM (\n" - f"{executed}\n" - f") AS _wren_sub LIMIT {limit}" - ) + executed = f"SELECT * FROM (\n{executed}\n) AS _wren_sub LIMIT {limit}" try: with self.connection.cursor() as cursor: cursor.execute(executed) From d005b2267870e6292b15ce58e79a1a80f0a1b606 Mon Sep 17 00:00:00 2001 From: Bartok9 <259807879+Bartok9@users.noreply.github.com> Date: Fri, 7 Aug 2026 02:49:38 -0400 Subject: [PATCH 3/3] refactor(snowflake): align LIMIT-wrap comment with pre-strip ordering Drop the trailing-semicolon composition rationale from the subquery wrap comment now that strip_trailing_semicolon runs before the wrap; keep the ORDER BY / pushdown reason. --- core/wren/src/wren/connector/snowflake.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/core/wren/src/wren/connector/snowflake.py b/core/wren/src/wren/connector/snowflake.py index 79ff7aabd5..6fcf5d8c00 100644 --- a/core/wren/src/wren/connector/snowflake.py +++ b/core/wren/src/wren/connector/snowflake.py @@ -60,10 +60,9 @@ def query(self, sql: str, limit: int | None = None) -> pa.Table: # bigquery/duckdb/redshift): strip a terminating `;` before send. executed = strip_trailing_semicolon(sql) # Push LIMIT into Snowflake when requested so we do not download a - # full result set only to slice it in Python. Wrap as a subquery so a - # trailing semicolon in the user SQL cannot break composition, and so + # full result set only to slice it in Python. Wrap as a subquery so # statements that already contain an ORDER BY keep their ordering - # under the outer LIMIT. + # under the outer LIMIT. (Trailing `;` is already stripped above.) if limit is not None: # Place the user SQL on its own line so a trailing line comment # (`-- ...`) cannot swallow the closing paren, alias, or LIMIT.