fix(athena): properly escape identifiers in Iceberg SQL generation (#3423)

* fix(athena): escape identifiers in Iceberg SQL to prevent injection via Glue column names

* fix(athena): escape DDL identifiers (backticks) in CREATE/ALTER and overwrite DELETE

* refactor(athena): rename identifier escaper to _escape_athena_dml_identifier for symmetry with DDL helper

Author: kukushking

awswrangler/athena/_write_iceberg.py

@@ -39,6 +39,27 @@ def _escape_athena_string_literal(value: Any) -> str:
     return str(value).replace("'", "''")
 
 
+def _escape_athena_dml_identifier(name: Any) -> str:
+    # For identifiers (column/table names) spliced inside DOUBLE-QUOTE delimited
+    # identifiers in DML statements (SELECT/INSERT/MERGE/DELETE), e.g. "<name>".
+    # Column names can originate from the Glue Data Catalog during automatic schema
+    # reconciliation, so they are not necessarily caller-trusted.
+    #
+    # Athena's Trino-based DML engine escapes an embedded double-quote in a delimited
+    # identifier by doubling it. Without this, a name containing " closes the
+    # identifier and the remainder is parsed as SQL.
+    return str(name).replace('"', '""')
+
+
+def _escape_athena_ddl_identifier(name: Any) -> str:
+    # For identifiers spliced inside BACKTICK delimited identifiers in DDL statements
+    # (CREATE TABLE / ALTER TABLE), e.g. `<name>`. Athena's Hive-based DDL engine uses
+    # backticks (not double quotes) for identifier quoting and escapes an embedded
+    # backtick by doubling it. Mixing the two families is a syntax error, so DDL and
+    # DML splices use separate helpers.
+    return str(name).replace("`", "``")
+
+
 def _create_iceberg_table(
     df: pd.DataFrame,
     database: str,
@@ -63,12 +84,14 @@ def _create_iceberg_table(
     columns_types, _ = catalog.extract_athena_types(df=df, index=index, dtype=dtype)
     cols_str: str = ", ".join(
         [
-            f"{k} {v}"
+            f"`{_escape_athena_ddl_identifier(k)}` {v}"
             if (columns_comments is None or columns_comments.get(k) is None)
-            else f"{k} {v} COMMENT '{_escape_athena_string_literal(columns_comments[k])}'"
+            else f"`{_escape_athena_ddl_identifier(k)}` {v} COMMENT '{_escape_athena_string_literal(columns_comments[k])}'"
             for k, v in columns_types.items()
         ]
     )
+    # partition_cols may be partition transform expressions (e.g. "day(ts)", "truncate(10, col)"),
+    # not plain identifiers, so they are spliced verbatim rather than quoted as identifiers.
     partition_cols_str: str = f"PARTITIONED BY ({', '.join([col for col in partition_cols])})" if partition_cols else ""
     table_properties_str: str = (
         ", "
@@ -83,9 +106,9 @@ def _create_iceberg_table(
     )
 
     create_sql: str = (
-        f"CREATE TABLE IF NOT EXISTS `{table}` ({cols_str}) "
+        f"CREATE TABLE IF NOT EXISTS `{_escape_athena_ddl_identifier(table)}` ({cols_str}) "
         f"{partition_cols_str} "
-        f"LOCATION '{path}' "
+        f"LOCATION '{_escape_athena_string_literal(path)}' "
         f"TBLPROPERTIES ('table_type' ='ICEBERG', 'format'='parquet'{table_properties_str})"
     )
 
@@ -216,9 +239,11 @@ def _alter_iceberg_table_add_columns_sql(
     table: str,
     columns_to_add: dict[str, str],
 ) -> list[str]:
-    add_cols_str = ", ".join([f"{col_name} {columns_to_add[col_name]}" for col_name in columns_to_add])
+    add_cols_str = ", ".join(
+        [f"`{_escape_athena_ddl_identifier(col_name)}` {columns_to_add[col_name]}" for col_name in columns_to_add]
+    )
 
-    return [f"ALTER TABLE {table} ADD COLUMNS ({add_cols_str})"]
+    return [f"ALTER TABLE `{_escape_athena_ddl_identifier(table)}` ADD COLUMNS ({add_cols_str})"]
 
 
 def _alter_iceberg_table_change_columns_sql(
@@ -228,7 +253,9 @@ def _alter_iceberg_table_change_columns_sql(
     sql_statements = []
 
     for col_name, col_type in columns_to_change.items():
-        sql_statements.append(f"ALTER TABLE {table} CHANGE COLUMN {col_name} {col_name} {col_type}")
+        escaped_table = _escape_athena_ddl_identifier(table)
+        escaped_col = _escape_athena_ddl_identifier(col_name)
+        sql_statements.append(f"ALTER TABLE `{escaped_table}` CHANGE COLUMN `{escaped_col}` `{escaped_col}` {col_type}")
 
     return sql_statements
 
@@ -290,7 +317,9 @@ def _build_order_by_clause(partition_cols: list[str] | None) -> str:
     if not partition_cols:
         return ""
 
-    order_cols = [f'"{_extract_column_from_partition_transform(col)}"' for col in partition_cols]
+    order_cols = [
+        f'"{_escape_athena_dml_identifier(_extract_column_from_partition_transform(col))}"' for col in partition_cols
+    ]
     return f"ORDER BY {', '.join(order_cols)}"
 
 
@@ -359,34 +388,40 @@ def _merge_iceberg(
     """
     wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
 
+    esc_database = _escape_athena_dml_identifier(database)
+    esc_table = _escape_athena_dml_identifier(table)
+    esc_source_table = _escape_athena_dml_identifier(source_table)
+    esc_columns = [_escape_athena_dml_identifier(x) for x in df.columns]
+
     sql_statement: str
     if merge_cols:
+        esc_merge_cols = [_escape_athena_dml_identifier(x) for x in merge_cols]
         if merge_condition == "update":
             match_condition = f"""WHEN MATCHED THEN
-                UPDATE SET {", ".join([f'"{x}" = source."{x}"' for x in df.columns])}"""
+                UPDATE SET {", ".join([f'"{x}" = source."{x}"' for x in esc_columns])}"""
         else:
             match_condition = ""
 
         if merge_match_nulls:
-            merge_conditions = [f'(target."{x}" IS NOT DISTINCT FROM source."{x}")' for x in merge_cols]
+            merge_conditions = [f'(target."{x}" IS NOT DISTINCT FROM source."{x}")' for x in esc_merge_cols]
         else:
-            merge_conditions = [f'(target."{x}" = source."{x}")' for x in merge_cols]
+            merge_conditions = [f'(target."{x}" = source."{x}")' for x in esc_merge_cols]
 
         sql_statement = f"""
-            MERGE INTO "{database}"."{table}" target
-            USING "{database}"."{source_table}" source
+            MERGE INTO "{esc_database}"."{esc_table}" target
+            USING "{esc_database}"."{esc_source_table}" source
             ON {" AND ".join(merge_conditions)}
             {match_condition}
             WHEN NOT MATCHED THEN
-                INSERT ({", ".join([f'"{x}"' for x in df.columns])})
-                VALUES ({", ".join([f'source."{x}"' for x in df.columns])})
+                INSERT ({", ".join([f'"{x}"' for x in esc_columns])})
+                VALUES ({", ".join([f'source."{x}"' for x in esc_columns])})
         """
     else:
         order_by_clause = _build_order_by_clause(partition_cols)
         sql_statement = f"""
-        INSERT INTO "{database}"."{table}" ({", ".join([f'"{x}"' for x in df.columns])})
-        SELECT {", ".join([f'"{x}"' for x in df.columns])}
-          FROM "{database}"."{source_table}"
+        INSERT INTO "{esc_database}"."{esc_table}" ({", ".join([f'"{x}"' for x in esc_columns])})
+        SELECT {", ".join([f'"{x}"' for x in esc_columns])}
+          FROM "{esc_database}"."{esc_source_table}"
           {order_by_clause}
         """
 
@@ -644,7 +679,7 @@ def to_iceberg(  # noqa: PLR0913
             )
         # if mode == "overwrite", delete whole data from table (but not table itself)
         elif mode == "overwrite":
-            delete_sql_statement = f"DELETE FROM {table}"
+            delete_sql_statement = f'DELETE FROM "{_escape_athena_dml_identifier(table)}"'
             delete_query_execution_id: str = _start_query_execution(
                 sql=delete_sql_statement,
                 workgroup=workgroup,
@@ -829,10 +864,14 @@ def delete_from_iceberg_table(
             index=False,
         )
 
+        esc_database = _escape_athena_dml_identifier(database)
+        esc_table = _escape_athena_dml_identifier(table)
+        esc_temp_table = _escape_athena_dml_identifier(temp_table)
+        esc_merge_cols = [_escape_athena_dml_identifier(x) for x in merge_cols]
         sql_statement = f"""
-            MERGE INTO "{database}"."{table}" target
-            USING "{database}"."{temp_table}" source
-            ON {" AND ".join([f'target."{x}" = source."{x}"' for x in merge_cols])}
+            MERGE INTO "{esc_database}"."{esc_table}" target
+            USING "{esc_database}"."{esc_temp_table}" source
+            ON {" AND ".join([f'target."{x}" = source."{x}"' for x in esc_merge_cols])}
             WHEN MATCHED THEN
                 DELETE
         """

tests/unit/test_athena_iceberg.py

@@ -1270,3 +1270,101 @@ def test_build_order_by_clause_multiple() -> None:
 
     result = _build_order_by_clause(["name", "day(ts)"])
     assert result == 'ORDER BY "name", "ts"'
+
+
+def test_escape_athena_dml_identifier_doubles_quotes() -> None:
+    from awswrangler.athena._write_iceberg import _escape_athena_dml_identifier
+
+    assert _escape_athena_dml_identifier("name") == "name"
+    assert _escape_athena_dml_identifier('a"b') == 'a""b'
+    assert _escape_athena_dml_identifier('x") DROP TABLE t --') == 'x"") DROP TABLE t --'
+
+
+def test_build_order_by_clause_escapes_identifier() -> None:
+    from awswrangler.athena._write_iceberg import _build_order_by_clause
+
+    result = _build_order_by_clause(['a"b'])
+    assert result == 'ORDER BY "a""b"'
+
+
+def test_escape_athena_ddl_identifier_doubles_backticks() -> None:
+    from awswrangler.athena._write_iceberg import _escape_athena_ddl_identifier
+
+    assert _escape_athena_ddl_identifier("name") == "name"
+    assert _escape_athena_ddl_identifier("a`b") == "a``b"
+    assert _escape_athena_ddl_identifier("x`) DROP TABLE t --") == "x``) DROP TABLE t --"
+
+
+def test_alter_iceberg_add_columns_escapes_identifier() -> None:
+    # ALTER TABLE is Hive-based DDL: identifiers are backtick-quoted, escaped by doubling `.
+    from awswrangler.athena._write_iceberg import _alter_iceberg_table_add_columns_sql
+
+    result = _alter_iceberg_table_add_columns_sql(table="t`1", columns_to_add={"c`) x --": "bigint"})
+    assert result == ["ALTER TABLE `t``1` ADD COLUMNS (`c``) x --` bigint)"]
+
+
+def test_alter_iceberg_change_columns_escapes_identifier() -> None:
+    from awswrangler.athena._write_iceberg import _alter_iceberg_table_change_columns_sql
+
+    result = _alter_iceberg_table_change_columns_sql(table="t`1", columns_to_change={"c`x": "bigint"})
+    assert result == ["ALTER TABLE `t``1` CHANGE COLUMN `c``x` `c``x` bigint"]
+
+
+def test_merge_iceberg_escapes_malicious_column_name() -> None:
+    """Column names from the Glue catalog must be escaped before being spliced into SQL."""
+    from unittest import mock
+
+    from awswrangler.athena import _write_iceberg
+
+    malicious = 'id") ; DROP TABLE victim --'
+    df = pd.DataFrame({malicious: [1], "v": [2]})
+
+    with mock.patch.object(_write_iceberg, "_get_workgroup_config", return_value=mock.MagicMock()), mock.patch.object(
+        _write_iceberg, "_start_query_execution", return_value="qid"
+    ) as start, mock.patch.object(_write_iceberg, "wait_query"):
+        _write_iceberg._merge_iceberg(df=df, database="db", table="t", source_table="src")
+
+    sql = start.call_args.kwargs["sql"]
+    # The injected identifier must appear only in doubled-quote form, never as a raw closing quote.
+    assert '"id"") ; DROP TABLE victim --"' in sql
+    assert '"id") ;' not in sql
+
+
+def test_create_iceberg_table_escapes_ddl_identifiers() -> None:
+    """Table/column names spliced into CREATE TABLE DDL must be backtick-escaped."""
+    from unittest import mock
+
+    from awswrangler.athena import _write_iceberg
+
+    df = pd.DataFrame({"c`0": [1]})
+
+    with mock.patch.object(
+        _write_iceberg.catalog,
+        "extract_athena_types",
+        return_value=({"c`0": "bigint"}, {}),
+    ), mock.patch.object(_write_iceberg, "_start_query_execution", return_value="qid") as start, mock.patch.object(
+        _write_iceberg, "wait_query"
+    ):
+        _write_iceberg._create_iceberg_table(
+            df=df,
+            database="db",
+            table="t`1",
+            path="s3://bucket/t/",
+            wg_config=mock.MagicMock(),
+            partition_cols=None,
+            additional_table_properties=None,
+        )
+
+    sql = start.call_args.kwargs["sql"]
+    assert "CREATE TABLE IF NOT EXISTS `t``1`" in sql
+    assert "`c``0` bigint" in sql
+
+
+def test_delete_from_iceberg_overwrite_escapes_table() -> None:
+    """The DELETE FROM issued for mode='overwrite' must quote/escape the table name."""
+    from awswrangler.athena._write_iceberg import _escape_athena_dml_identifier
+
+    # Mirrors the splice in to_iceberg's overwrite branch.
+    table = 'v" ; DROP TABLE x --'
+    stmt = f'DELETE FROM "{_escape_athena_dml_identifier(table)}"'
+    assert stmt == 'DELETE FROM "v"" ; DROP TABLE x --"'