mirror of
https://github.com/Canner/WrenAI.git
synced 2026-08-30 18:00:36 +08:00
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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user