refactor(snowflake): strip trailing semicolon on unlimited query path for connector consistency (#2593)

Co-authored-by: Bartok9 <bartok9@users.noreply.github.com>
This commit is contained in:
Bartok
2026-08-10 22:48:42 -04:00
committed by GitHub
parent b1387d89ec
commit ec85b1e158
2 changed files with 18 additions and 9 deletions
+6 -9
View File
@@ -56,20 +56,17 @@ class SnowflakeConnector(ConnectorABC):
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
# 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.
executed = sql
# 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.
executed = (
"SELECT * FROM (\n"
f"{strip_trailing_semicolon(sql)}\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)
@@ -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()