diff --git a/pyathena/formatter.py b/pyathena/formatter.py index 694b6a9a..5759642e 100644 --- a/pyathena/formatter.py +++ b/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: diff --git a/tests/pyathena/test_formatter.py b/tests/pyathena/test_formatter.py index 631a1163..7af39e4c 100644 --- a/tests/pyathena/test_formatter.py +++ b/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} + )