Merge pull request #745 from pyathena-dev/fix/formatter-escaper-sqli

Fix escaper selection for DELETE and CTAS to prevent SQL injection (CVE-2026-65321)

Author: laughingman7743

pyathena/formatter.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import re
 import textwrap
 import uuid
 from abc import ABCMeta, abstractmethod
@@ -154,6 +155,57 @@ def _escape_presto(val: str) -> str:
     return f"'{escaped}'"
 
 
+_LEADING_COMMENT_PATTERN = re.compile(
+    r"^(?:\s*(?:/\*.*?\*/|--[^\n]*(?:\n|$)))+",
+    re.DOTALL,
+)
+
+# Statements executed by the Hive DDL engine, where the backslash escaping
+# convention of _escape_hive() is the correct one. Everything else is assumed
+# to be executed by Trino, which requires single quotes to be doubled and
+# treats a backslash inside a string literal as an ordinary character.
+#
+# Defaulting to the Trino escaper is deliberate: it is the fail-safe
+# direction. Doubling a quote in a Hive statement is at worst a formatting
+# bug, because the Hive lexer concatenates adjacent string literals, while
+# backslash-escaping a quote in a Trino statement lets the parameter value
+# terminate the string literal and inject SQL.
+#
+# CREATE TABLE is only treated as Hive DDL when it is not a CTAS statement.
+# CREATE TABLE ... AS SELECT and CREATE VIEW are executed by Trino.
+_HIVE_STATEMENT_PATTERN = re.compile(
+    r"""^(?:
+          ALTER\s+(?:DATABASE|SCHEMA|TABLE)\b
+        | CREATE\s+(?:DATABASE|SCHEMA)\b
+        | CREATE\s+(?:EXTERNAL\s+)?TABLE\b(?!.*\bAS\b)
+        | DROP\s+(?:DATABASE|SCHEMA|TABLE)\b
+        | MSCK\s+REPAIR\b
+        | SHOW\b
+        | DESC(?:RIBE)?\b
+    )""",
+    re.IGNORECASE | re.VERBOSE | re.DOTALL,
+)
+
+
+def _strip_leading_comments(operation: str) -> str:
+    """Remove leading SQL comments and whitespace.
+
+    Statement type detection must not be defeated by a leading comment such
+    as ``/* generated by ... */ DELETE FROM ...``.
+    """
+    match = _LEADING_COMMENT_PATTERN.match(operation)
+    if match:
+        operation = operation[match.end() :]
+    return operation.lstrip()
+
+
+def _get_escaper(operation: str) -> Callable[[str], str]:
+    """Select the escaper matching the engine that will parse the statement."""
+    if _HIVE_STATEMENT_PATTERN.match(_strip_leading_comments(operation)):
+        return _escape_hive
+    return _escape_presto
+
+
 def _escape_hive(val: str) -> str:
     escaped = (
         val.replace("\\", "\\\\")
@@ -268,11 +320,7 @@ def format(self, operation: str, parameters: dict[str, Any] | None = None) -> st
             raise ProgrammingError("Query is none or empty.")
         operation = operation.strip()
 
-        operation_upper = operation.upper()
-        if operation_upper.startswith(("SELECT", "WITH", "INSERT", "UPDATE", "MERGE")):
-            escaper = _escape_presto
-        else:
-            escaper = _escape_hive
+        escaper = _get_escaper(operation)
 
         kwargs: dict[str, Any] | None = None
         if parameters is not None:

tests/pyathena/test_formatter.py

@@ -6,6 +6,61 @@
 
 from pyathena.error import ProgrammingError
 
+# A value whose single quote must not be allowed to terminate the string
+# literal. Trino and Athena do not treat a backslash as an escape character
+# inside a string literal, so a statement parsed by Trino must double the quote
+# ('') rather than backslash-escape it (\').
+HOSTILE = "a' OR 1=1 --"
+
+# Statements parsed by the Trino/Athena engine, where single quotes must be
+# doubled. DELETE and CTAS (CREATE TABLE ... AS SELECT) are the paths that
+# CVE-2026-65321 left routed through the Hive escaper.
+TRINO_STATEMENTS = [
+    "SELECT * FROM t WHERE c = %(v)s",
+    "WITH x AS (SELECT 1) SELECT * FROM t WHERE c = %(v)s",
+    "INSERT INTO t VALUES (%(v)s)",
+    "UPDATE t SET c = %(v)s",
+    "MERGE INTO t USING s ON t.c = %(v)s",
+    "DELETE FROM t WHERE c = %(v)s",
+    "delete from t where c = %(v)s",
+    "CREATE TABLE t WITH (format = 'PARQUET') AS SELECT * FROM s WHERE c = %(v)s",
+    "CREATE VIEW v AS SELECT * FROM s WHERE c = %(v)s",
+    "EXPLAIN DELETE FROM t WHERE c = %(v)s",
+    "UNLOAD (SELECT * FROM t WHERE c = %(v)s) TO 's3://b/p/' WITH (format = 'PARQUET')",
+    "/* generated by etl */ DELETE FROM t WHERE c = %(v)s",
+    "-- nightly compaction\nDELETE FROM t WHERE c = %(v)s",
+    "/* a */\n-- b\nDELETE FROM t WHERE c = %(v)s",
+    "/* multi\nline */\nDELETE FROM t WHERE c = %(v)s",
+    "\n\t DELETE FROM t WHERE c = %(v)s",
+    # Near-miss statements: their prefixes are adjacent to the Hive allowlist
+    # (ALTER/DROP/CREATE ... VIEW, VALUES) but Trino executes them, so they must
+    # NOT be misrouted to the Hive escaper.
+    "ALTER VIEW v AS SELECT name FROM s WHERE c = %(v)s",
+    "DROP VIEW v -- %(v)s",
+    "CREATE OR REPLACE VIEW v AS SELECT %(v)s",
+    "VALUES (%(v)s)",
+]
+
+# Statements parsed by the Hive DDL engine, where backslash escaping is correct.
+HIVE_STATEMENTS = [
+    "CREATE EXTERNAL TABLE t (c string) LOCATION %(v)s",
+    "CREATE TABLE t (c string) LOCATION %(v)s",
+    "CREATE DATABASE d LOCATION %(v)s",
+    "ALTER TABLE t SET LOCATION %(v)s",
+    "DROP TABLE t PURGE -- %(v)s",
+    "MSCK REPAIR TABLE t -- %(v)s",
+    "SHOW PARTITIONS t -- %(v)s",
+    "DESCRIBE t -- %(v)s",
+    "CREATE SCHEMA s LOCATION %(v)s",
+    "DROP DATABASE d -- %(v)s",
+    "ALTER DATABASE d SET DBPROPERTIES ('k' = %(v)s)",
+    # Leading comments must be stripped before detection on the Hive side too,
+    # otherwise a commented Hive DDL statement would silently lose its escaping.
+    "-- cleanup\nDROP TABLE t LOCATION %(v)s",
+    "/* migrate */ ALTER TABLE t SET LOCATION %(v)s",
+    "-- a\n-- b\nCREATE EXTERNAL TABLE t (c string) LOCATION %(v)s",
+]
+
 
 class TestDefaultParameterFormatter:
     def test_add_partition(self, formatter):
@@ -468,3 +523,33 @@ def test_format_bad_parameter(self, formatter):
                 ["a string"],
             ),
         )
+
+    @pytest.mark.parametrize("operation", TRINO_STATEMENTS)
+    def test_trino_statements_use_quote_doubling(self, formatter, operation):
+        result = formatter.format(operation, {"v": HOSTILE})
+        assert "\\'" not in result, f"backslash escaping leaked into: {result}"
+        assert "''" in result
+
+    @pytest.mark.parametrize("operation", HIVE_STATEMENTS)
+    def test_hive_ddl_keeps_backslash_escaping(self, formatter, operation):
+        assert "\\'" in formatter.format(operation, {"v": HOSTILE})
+
+    def test_quote_cannot_terminate_literal_in_delete(self, formatter):
+        result = formatter.format("DELETE FROM t WHERE c = %(v)s", {"v": HOSTILE})
+        # Every quote in the rendered value is doubled, so the literal opened by
+        # the formatter is still open until its own closing quote.
+        assert result.count("'") % 2 == 0
+        assert result == "DELETE FROM t WHERE c = 'a'' OR 1=1 --'"
+
+    def test_ctas_detection_is_case_and_whitespace_insensitive(self, formatter):
+        for operation in (
+            "create  table  t  as  select %(v)s",
+            "CREATE\nTABLE\nt\nAS\nSELECT %(v)s",
+        ):
+            assert "\\'" not in formatter.format(operation, {"v": HOSTILE})
+
+    def test_table_name_starting_with_as_is_still_hive_ddl(self, formatter):
+        # `as_of` must not be mistaken for the CTAS `AS` keyword.
+        assert "\\'" in formatter.format(
+            "CREATE EXTERNAL TABLE as_of (c string) LOCATION %(v)s", {"v": HOSTILE}
+        )