fix(athena): strip trailing semicolon before EXPLAIN in dry_run (#2421)

core/wren/src/wren/connector/athena.py

@@ -11,6 +11,7 @@ from __future__ import annotations
 import contextlib
 import datetime as dtlib
 import json
+import re
 from decimal import Decimal as PyDecimal
 from typing import Any
 
@@ -19,6 +20,22 @@ import pyarrow as pa
 from wren.connector.base import ConnectorABC
 from wren.model.error import DIALECT_SQL, ErrorCode, ErrorPhase, WrenError
 
+_TRAILING_SEMICOLONS_RE = re.compile(r"[;\s]+\Z")
+
+
+def _strip_trailing_semicolon(sql: str) -> str:
+    """Strip any trailing ``;`` characters and surrounding whitespace.
+
+    ``dry_run`` runs ``EXPLAIN {sql}``; Athena (Trino-flavoured) rejects a
+    trailing semicolon there (``EXPLAIN SELECT 1;`` → syntax error). Only the
+    terminating run of semicolons/whitespace is stripped, so semicolons inside
+    string literals (e.g. ``SELECT 'a;b'``) are preserved. Mirrors the
+    trino/postgres/mysql connectors, which already strip before EXPLAIN or
+    subquery-wrapping.
+    """
+    return _TRAILING_SEMICOLONS_RE.sub("", sql)
+
+
 # Athena's DB-API cursor returns Trino-style type names. We delegate the
 # lexing to sqlglot so we get nested type support (array<row<a int, b varchar>>,
 # decimal(p, s), map<K, V>, etc.) for free.
@@ -318,7 +335,7 @@ class AthenaConnector(ConnectorABC):
     def dry_run(self, sql: str) -> None:
         try:
             with contextlib.closing(self.connection.cursor()) as cursor:
-                cursor.execute(f"EXPLAIN {sql}")
+                cursor.execute(f"EXPLAIN {_strip_trailing_semicolon(sql)}")
         except (WrenError, TimeoutError):
             raise
         except Exception as e:

core/wren/tests/unit/test_athena_semicolon.py

@@ -0,0 +1,45 @@
+"""Trailing-semicolon stripping for the Athena connector (mocked, no live DB).
+
+``AthenaConnector.dry_run`` runs ``EXPLAIN {sql}``. Athena's engine is
+Trino-flavoured and rejects a trailing semicolon there
+(``EXPLAIN SELECT 1;`` -> syntax error). These tests use a mocked cursor and
+assert on the executed SQL, so no AWS connection is required.
+"""
+
+from unittest.mock import MagicMock
+
+from wren.connector.athena import AthenaConnector, _strip_trailing_semicolon
+
+
+def _make_mock_connector() -> tuple[AthenaConnector, MagicMock]:
+    """Build an AthenaConnector bypassing __init__ (no real connection)."""
+    connector = AthenaConnector.__new__(AthenaConnector)
+    cursor = MagicMock()
+    conn = MagicMock()
+    conn.cursor.return_value = cursor
+    connector.connection = conn
+    return connector, cursor
+
+
+def test_dry_run_strips_trailing_semicolon_before_explain() -> None:
+    connector, cursor = _make_mock_connector()
+    connector.dry_run("SELECT 1;")
+    (sent,), _ = cursor.execute.call_args
+    assert sent == "EXPLAIN SELECT 1"
+    assert not sent.endswith(";")
+
+
+def test_dry_run_strips_trailing_semicolon_and_whitespace() -> None:
+    connector, cursor = _make_mock_connector()
+    connector.dry_run("SELECT 1;  \n")
+    (sent,), _ = cursor.execute.call_args
+    assert sent == "EXPLAIN SELECT 1"
+
+
+def test_helper_preserves_semicolon_inside_string_literal() -> None:
+    sql = "SELECT 'a;b' AS x"
+    assert _strip_trailing_semicolon(sql) == sql
+
+
+def test_helper_no_trailing_semicolon_unchanged() -> None:
+    assert _strip_trailing_semicolon("SELECT 1") == "SELECT 1"