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

Author: kukushking

awswrangler/athena/_write_iceberg.py

@@ -39,6 +39,16 @@ def _escape_athena_string_literal(value: Any) -> str:
     return str(value).replace("'", "''")
 
 
+def _escape_athena_identifier(name: Any) -> str:
+    # Used for identifiers (column/table names) spliced inside double-quote delimited
+    # identifiers, e.g. "<name>". Column names can originate from the Glue Data Catalog
+    # during automatic schema reconciliation, so they are not necessarily caller-trusted.
+    #
+    # Trino/Athena delimited identifiers escape an embedded double-quote by doubling it.
+    # Without this, a name containing " closes the identifier and appends arbitrary SQL.
+    return str(name).replace('"', '""')
+
+
 def _create_iceberg_table(
     df: pd.DataFrame,
     database: str,
@@ -216,9 +226,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_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_identifier(table)}" ADD COLUMNS ({add_cols_str})']
 
 
 def _alter_iceberg_table_change_columns_sql(
@@ -228,7 +240,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_identifier(table)
+        escaped_col = _escape_athena_identifier(col_name)
+        sql_statements.append(f'ALTER TABLE "{escaped_table}" CHANGE COLUMN "{escaped_col}" "{escaped_col}" {col_type}')
 
     return sql_statements
 
@@ -290,7 +304,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_identifier(_extract_column_from_partition_transform(col))}"' for col in partition_cols
+    ]
     return f"ORDER BY {', '.join(order_cols)}"
 
 
@@ -359,34 +375,40 @@ def _merge_iceberg(
     """
     wg_config: _WorkGroupConfig = _get_workgroup_config(session=boto3_session, workgroup=workgroup)
 
+    esc_database = _escape_athena_identifier(database)
+    esc_table = _escape_athena_identifier(table)
+    esc_source_table = _escape_athena_identifier(source_table)
+    esc_columns = [_escape_athena_identifier(x) for x in df.columns]
+
     sql_statement: str
     if merge_cols:
+        esc_merge_cols = [_escape_athena_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}
         """
 
@@ -829,10 +851,14 @@ def delete_from_iceberg_table(
             index=False,
         )
 
+        esc_database = _escape_athena_identifier(database)
+        esc_table = _escape_athena_identifier(table)
+        esc_temp_table = _escape_athena_identifier(temp_table)
+        esc_merge_cols = [_escape_athena_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,52 @@ 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_identifier_doubles_quotes() -> None:
+    from awswrangler.athena._write_iceberg import _escape_athena_identifier
+
+    assert _escape_athena_identifier("name") == "name"
+    assert _escape_athena_identifier('a"b') == 'a""b'
+    assert _escape_athena_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_alter_iceberg_add_columns_escapes_identifier() -> None:
+    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