fix: Refactor SQL identifiers to mitigate injection risks (#2543)

Author: kukushking

awswrangler/_sql_utils.py

@@ -0,0 +1,40 @@
+"""SQL utilities."""
+import re
+
+from awswrangler import exceptions
+
+
+def identifier(sql: str, sql_mode: str = "mysql") -> str:
+    """
+    Turn the input into an escaped SQL identifier, such as the name of a table or column.
+
+    sql: str
+        Identifier to use in SQL.
+    sql_mode: str
+        "mysql" for default MySQL identifiers (backticks), "ansi" for ANSI-compatible identifiers (double quotes), or
+        "mssql" for MSSQL identifiers (square brackets).
+
+    Returns
+    -------
+    str
+        Escaped SQL identifier.
+    """
+    if not isinstance(sql, str):
+        raise exceptions.InvalidArgumentValue("identifier must be a str")
+
+    if len(sql) == 0:
+        raise exceptions.InvalidArgumentValue("identifier must be > 0 characters in length")
+
+    if re.search(r"[^a-zA-Z0-9-_ ]", sql):
+        raise exceptions.InvalidArgumentValue(
+            "identifier must contain only alphanumeric characters, spaces, underscores, or hyphens"
+        )
+
+    if sql_mode == "mysql":
+        return f"`{sql}`"
+    elif sql_mode == "ansi":
+        return f'"{sql}"'
+    elif sql_mode == "mssql":
+        return f"[{sql}]"
+
+    raise ValueError(f"Unknown SQL MODE: {sql_mode}")

awswrangler/data_api/rds.py

@@ -1,7 +1,6 @@
 """RDS Data API Connector."""
 import datetime as dt
 import logging
-import re
 import time
 import uuid
 from decimal import Decimal
@@ -12,6 +11,7 @@
 
 import awswrangler.pandas as pd
 from awswrangler import _data_types, _databases, _utils, exceptions
+from awswrangler._sql_utils import identifier
 from awswrangler.data_api import _connector
 
 if TYPE_CHECKING:
@@ -228,19 +228,6 @@ def _get_statement_result(self, request_id: str) -> pd.DataFrame:
         return dataframe
 
 
-def escape_identifier(identifier: str, sql_mode: str = "mysql") -> str:
-    """Escape identifiers. Uses MySQL-compatible backticks by default."""
-    if not isinstance(identifier, str):
-        raise TypeError("SQL identifier must be a string")
-    if re.search(r"\W", identifier):
-        raise TypeError(f"SQL identifier contains invalid characters: {identifier}")
-    if sql_mode == "mysql":
-        return f"`{identifier}`"
-    elif sql_mode == "ansi":
-        return f'"{identifier}"'
-    raise ValueError(f"Unknown SQL MODE: {sql_mode}")
-
-
 def connect(
     resource_arn: str, database: str, secret_arn: str = "", boto3_session: Optional[boto3.Session] = None, **kwargs: Any
 ) -> RdsDataApi:
@@ -286,7 +273,7 @@ def read_sql_query(sql: str, con: RdsDataApi, database: Optional[str] = None) ->
 
 
 def _drop_table(con: RdsDataApi, table: str, database: str, transaction_id: str, sql_mode: str) -> None:
-    sql = f"DROP TABLE IF EXISTS {escape_identifier(table, sql_mode=sql_mode)}"
+    sql = f"DROP TABLE IF EXISTS {identifier(table, sql_mode=sql_mode)}"
     _logger.debug("Drop table query:\n%s", sql)
     con.execute(sql, database=database, transaction_id=transaction_id)
 
@@ -329,8 +316,8 @@ def _create_table(
         varchar_lengths=varchar_lengths,
         converter_func=_data_types.pyarrow2mysql,
     )
-    cols_str: str = "".join([f"{escape_identifier(k, sql_mode=sql_mode)} {v},\n" for k, v in mysql_types.items()])[:-2]
-    sql = f"CREATE TABLE IF NOT EXISTS {escape_identifier(table, sql_mode=sql_mode)} (\n{cols_str})"
+    cols_str: str = "".join([f"{identifier(k, sql_mode=sql_mode)} {v},\n" for k, v in mysql_types.items()])[:-2]
+    sql = f"CREATE TABLE IF NOT EXISTS {identifier(table, sql_mode=sql_mode)} (\n{cols_str})"
 
     _logger.debug("Create table query:\n%s", sql)
     con.execute(sql, database=database, transaction_id=transaction_id)
@@ -443,6 +430,8 @@ def to_sql(
         inserted into the database columns `col1` and `col3`.
     chunksize: int
         Number of rows which are inserted with each SQL query. Defaults to inserting 200 rows per query.
+    sql_mode: str
+        "mysql" for default MySQL identifiers (backticks) or "ansi" for ANSI-compatible identifiers (double quotes).
     """
     if df.empty is True:
         raise exceptions.EmptyDataFrame("DataFrame cannot be empty.")
@@ -470,15 +459,13 @@ def to_sql(
             df = df.reset_index(level=df.index.names)
 
         if use_column_names:
-            insertion_columns = (
-                "(" + ", ".join([f"{escape_identifier(col, sql_mode=sql_mode)}" for col in df.columns]) + ")"
-            )
+            insertion_columns = "(" + ", ".join([f"{identifier(col, sql_mode=sql_mode)}" for col in df.columns]) + ")"
         else:
             insertion_columns = ""
 
         placeholders = ", ".join([f":{col}" for col in df.columns])
 
-        sql = f"INSERT INTO {escape_identifier(table, sql_mode=sql_mode)} {insertion_columns} VALUES ({placeholders})"
+        sql = f"INSERT INTO {identifier(table, sql_mode=sql_mode)} {insertion_columns} VALUES ({placeholders})"
         parameter_sets = _generate_parameter_sets(df)
 
         for parameter_sets_chunk in _utils.chunkify(parameter_sets, max_length=chunksize):

awswrangler/mysql.py

@@ -12,6 +12,7 @@
 from awswrangler import _data_types, _utils, exceptions
 from awswrangler import _databases as _db_utils
 from awswrangler._config import apply_configs
+from awswrangler._sql_utils import identifier
 
 if TYPE_CHECKING:
     try:
@@ -37,15 +38,19 @@ def _validate_connection(con: "Connection[Any]") -> None:
 
 
 def _drop_table(cursor: "Cursor", schema: Optional[str], table: str) -> None:
-    schema_str = f"`{schema}`." if schema else ""
-    sql = f"DROP TABLE IF EXISTS {schema_str}`{table}`"
+    schema_str = f"{identifier(schema)}." if schema else ""
+    sql = f"DROP TABLE IF EXISTS {schema_str}{identifier(table)}"
     _logger.debug("Drop table query:\n%s", sql)
     cursor.execute(sql)
 
 
 def _does_table_exist(cursor: "Cursor", schema: Optional[str], table: str) -> bool:
-    schema_str = f"TABLE_SCHEMA = '{schema}' AND" if schema else ""
-    cursor.execute(f"SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE {schema_str} TABLE_NAME = %s", args=[table])
+    if schema:
+        cursor.execute(
+            "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s", args=[schema, table]
+        )
+    else:
+        cursor.execute("SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = %s", args=[table])
     return len(cursor.fetchall()) > 0
 
 
@@ -71,8 +76,8 @@ def _create_table(
         varchar_lengths=varchar_lengths,
         converter_func=_data_types.pyarrow2mysql,
     )
-    cols_str: str = "".join([f"`{k}` {v},\n" for k, v in mysql_types.items()])[:-2]
-    sql = f"CREATE TABLE IF NOT EXISTS `{schema}`.`{table}` (\n{cols_str})"
+    cols_str: str = "".join([f"{identifier(k)} {v},\n" for k, v in mysql_types.items()])[:-2]
+    sql = f"CREATE TABLE IF NOT EXISTS {identifier(schema)}.{identifier(table)} (\n{cols_str})"
     _logger.debug("Create table query:\n%s", sql)
     cursor.execute(sql)
 
@@ -419,7 +424,11 @@ def read_sql_table(
     >>> con.close()
 
     """
-    sql: str = f"SELECT * FROM `{table}`" if schema is None else f"SELECT * FROM `{schema}`.`{table}`"
+    sql: str = (
+        f"SELECT * FROM {identifier(table)}"
+        if schema is None
+        else f"SELECT * FROM {identifier(schema)}.{identifier(table)}"
+    )
     return read_sql_query(
         sql=sql,
         con=con,
@@ -551,29 +560,35 @@ def to_sql(
             upsert_str = ""
             ignore_str = " IGNORE" if mode == "ignore" else ""
             if use_column_names:
-                insertion_columns = f"(`{'`, `'.join(df.columns)}`)"
+                insertion_columns = f"({', '.join([identifier(col) for col in df.columns])})"
             if mode == "upsert_duplicate_key":
-                upsert_columns = ", ".join(df.columns.map(lambda column: f"`{column}`=VALUES(`{column}`)"))
+                upsert_columns = ", ".join(df.columns.map(lambda col: f"{identifier(col)}=VALUES({identifier(col)})"))
                 upsert_str = f" ON DUPLICATE KEY UPDATE {upsert_columns}"
             placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
                 df=df, column_placeholders=column_placeholders, chunksize=chunksize
             )
             sql: str
             for placeholders, parameters in placeholder_parameter_pair_generator:
                 if mode == "upsert_replace_into":
-                    sql = f"REPLACE INTO `{schema}`.`{table}` {insertion_columns} VALUES {placeholders}"
+                    sql = f"REPLACE INTO {identifier(schema)}.{identifier(table)} {insertion_columns} VALUES {placeholders}"
                 else:
-                    sql = f"""INSERT{ignore_str} INTO `{schema}`.`{table}` {insertion_columns}
+                    sql = f"""INSERT{ignore_str} INTO {identifier(schema)}.{identifier(table)} {insertion_columns}
 VALUES {placeholders}{upsert_str}"""
                 _logger.debug("sql: %s", sql)
                 cursor.executemany(sql, (parameters,))
             con.commit()
             if mode == "upsert_distinct":
                 temp_table = f"{table}_{uuid.uuid4().hex}"
-                cursor.execute(f"CREATE TABLE `{schema}`.`{temp_table}` LIKE `{schema}`.`{table}`")
-                cursor.execute(f"INSERT INTO `{schema}`.`{temp_table}` SELECT DISTINCT * FROM `{schema}`.`{table}`")
-                cursor.execute(f"DROP TABLE IF EXISTS `{schema}`.`{table}`")
-                cursor.execute(f"ALTER TABLE `{schema}`.`{temp_table}` RENAME TO `{table}`")
+                cursor.execute(
+                    f"CREATE TABLE {identifier(schema)}.{identifier(temp_table)} LIKE {identifier(schema)}.{identifier(table)}"
+                )
+                cursor.execute(
+                    f"INSERT INTO {identifier(schema)}.{identifier(temp_table)} SELECT DISTINCT * FROM {identifier(schema)}.{identifier(table)}"
+                )
+                cursor.execute(f"DROP TABLE IF EXISTS {identifier(schema)}.{identifier(table)}")
+                cursor.execute(
+                    f"ALTER TABLE {identifier(schema)}.{identifier(temp_table)} RENAME TO {identifier(table)}"
+                )
                 con.commit()
 
     except Exception as ex:

awswrangler/oracle.py

@@ -24,6 +24,7 @@
 from awswrangler import _data_types, _utils, exceptions
 from awswrangler import _databases as _db_utils
 from awswrangler._config import apply_configs
+from awswrangler._sql_utils import identifier
 
 __all__ = ["connect", "read_sql_query", "read_sql_table", "to_sql"]
 
@@ -43,8 +44,8 @@ def _validate_connection(con: "oracledb.Connection") -> None:
 
 
 def _get_table_identifier(schema: Optional[str], table: str) -> str:
-    schema_str = f'"{schema}".' if schema else ""
-    table_identifier = f'{schema_str}"{table}"'
+    schema_str = f'{identifier(schema, sql_mode="ansi")}.' if schema else ""
+    table_identifier = f'{schema_str}{identifier(table, sql_mode="ansi")}'
     return table_identifier
 
 
@@ -65,8 +66,14 @@ def _drop_table(cursor: "oracledb.Cursor", schema: Optional[str], table: str) ->
 
 
 def _does_table_exist(cursor: "oracledb.Cursor", schema: Optional[str], table: str) -> bool:
-    schema_str = f"OWNER = '{schema}' AND" if schema else ""
-    cursor.execute(f"SELECT * FROM ALL_TABLES WHERE {schema_str} TABLE_NAME = '{table}'")
+    if schema:
+        cursor.execute(
+            "SELECT * FROM ALL_TABLES WHERE OWNER = :db_schema AND TABLE_NAME = :db_table",
+            db_schema=schema,
+            db_table=table,
+        )
+    else:
+        cursor.execute("SELECT * FROM ALL_TABLES WHERE TABLE_NAME = :tbl", tbl=table)
     return len(cursor.fetchall()) > 0
 
 
@@ -93,10 +100,10 @@ def _create_table(
         varchar_lengths=varchar_lengths,
         converter_func=_data_types.pyarrow2oracle,
     )
-    cols_str: str = "".join([f'"{k}" {v},\n' for k, v in oracle_types.items()])[:-2]
+    cols_str: str = "".join([f'{identifier(k, sql_mode="ansi")} {v},\n' for k, v in oracle_types.items()])[:-2]
 
     if primary_keys:
-        primary_keys_str = ", ".join([f'"{k}"' for k in primary_keys])
+        primary_keys_str = ", ".join([f'{identifier(k, sql_mode="ansi")}' for k in primary_keys])
     else:
         primary_keys_str = None
 
@@ -450,7 +457,7 @@ def _generate_insert_statement(
     column_placeholders: str = f"({', '.join([':' + str(i + 1) for i in range(len(df.columns))])})"
 
     if use_column_names:
-        insertion_columns = "(" + ", ".join('"' + column + '"' for column in df.columns) + ")"
+        insertion_columns = "(" + ", ".join(identifier(column, sql_mode="ansi") for column in df.columns) + ")"
     else:
         insertion_columns = ""
 
@@ -470,14 +477,19 @@ def _generate_upsert_statement(
 
     non_primary_key_columns = [key for key in df.columns if key not in set(primary_keys)]
 
-    primary_keys_str = ", ".join([f'"{key}"' for key in primary_keys])
-    columns_str = ", ".join([f'"{key}"' for key in non_primary_key_columns])
+    primary_keys_str = ", ".join([f'{identifier(key, sql_mode="ansi")}' for key in primary_keys])
+    columns_str = ", ".join([f'{identifier(key, sql_mode="ansi")}' for key in non_primary_key_columns])
 
     column_placeholders: str = f"({', '.join([':' + str(i + 1) for i in range(len(df.columns))])})"
 
-    primary_key_condition_str = " AND ".join([f'"{key}" = :{i+1}' for i, key in enumerate(primary_keys)])
+    primary_key_condition_str = " AND ".join(
+        [f'{identifier(key, sql_mode="ansi")} = :{i+1}' for i, key in enumerate(primary_keys)]
+    )
     assignment_str = ", ".join(
-        [f'"{col}" = :{i + len(primary_keys) + 1}' for i, col in enumerate(non_primary_key_columns)]
+        [
+            f'{identifier(col, sql_mode="ansi")} = :{i + len(primary_keys) + 1}'
+            for i, col in enumerate(non_primary_key_columns)
+        ]
     )
 
     return f"""

awswrangler/postgresql.py

@@ -69,7 +69,7 @@ def _create_table(
         varchar_lengths=varchar_lengths,
         converter_func=_data_types.pyarrow2postgresql,
     )
-    cols_str: str = "".join([f'"{k}" {v},\n' for k, v in postgresql_types.items()])[:-2]
+    cols_str: str = "".join([f"{pg8000_native.identifier(k)} {v},\n" for k, v in postgresql_types.items()])[:-2]
     sql = f"CREATE TABLE IF NOT EXISTS {pg8000_native.identifier(schema)}.{pg8000_native.identifier(table)} (\n{cols_str})"
     _logger.debug("Create table query:\n%s", sql)
     cursor.execute(sql)
@@ -584,7 +584,7 @@ def to_sql(
             if index:
                 df.reset_index(level=df.index.names, inplace=True)
             column_placeholders: str = ", ".join(["%s"] * len(df.columns))
-            column_names = [f'"{column}"' for column in df.columns]
+            column_names = [pg8000_native.identifier(column) for column in df.columns]
             insertion_columns = ""
             upsert_str = ""
             if use_column_names:

awswrangler/sqlserver.py

@@ -24,6 +24,7 @@
 from awswrangler import _data_types, _utils, exceptions
 from awswrangler import _databases as _db_utils
 from awswrangler._config import apply_configs
+from awswrangler._sql_utils import identifier
 
 __all__ = ["connect", "read_sql_query", "read_sql_table", "to_sql"]
 
@@ -50,9 +51,10 @@ def _validate_connection(con: "pyodbc.Connection") -> None:
 
 
 def _get_table_identifier(schema: Optional[str], table: str) -> str:
-    schema_str = f'"{schema}".' if schema else ""
-    table_identifier = f'{schema_str}"{table}"'
-    return table_identifier
+    if schema:
+        return f"{identifier(schema, sql_mode='mssql')}.{identifier(table, sql_mode='mssql')}"
+    else:
+        return identifier(table, sql_mode="mssql")
 
 
 def _drop_table(cursor: "Cursor", schema: Optional[str], table: str) -> None:
@@ -63,8 +65,12 @@ def _drop_table(cursor: "Cursor", schema: Optional[str], table: str) -> None:
 
 
 def _does_table_exist(cursor: "Cursor", schema: Optional[str], table: str) -> bool:
-    schema_str = f"TABLE_SCHEMA = '{schema}' AND" if schema else ""
-    cursor.execute(f"SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE {schema_str} TABLE_NAME = ?", table)
+    if schema:
+        cursor.execute(
+            "SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA = ? AND TABLE_NAME = ?", (schema, table)
+        )
+    else:
+        cursor.execute("SELECT * FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_NAME = ?", table)
     return len(cursor.fetchall()) > 0
 
 
@@ -90,7 +96,7 @@ def _create_table(
         varchar_lengths=varchar_lengths,
         converter_func=_data_types.pyarrow2sqlserver,
     )
-    cols_str: str = "".join([f'"{k}" {v},\n' for k, v in sqlserver_types.items()])[:-2]
+    cols_str: str = "".join([f"{identifier(k, sql_mode='mssql')} {v},\n" for k, v in sqlserver_types.items()])[:-2]
     table_identifier = _get_table_identifier(schema, table)
     sql = (
         f"IF OBJECT_ID(N'{table_identifier}', N'U') IS NULL BEGIN CREATE TABLE {table_identifier} (\n{cols_str}); END;"
@@ -529,7 +535,7 @@ def to_sql(
             table_identifier = _get_table_identifier(schema, table)
             insertion_columns = ""
             if use_column_names:
-                quoted_columns = ", ".join(f'"{col}"' for col in df.columns)
+                quoted_columns = ", ".join(f"{identifier(col, sql_mode='mssql')}" for col in df.columns)
                 insertion_columns = f"({quoted_columns})"
             placeholder_parameter_pair_generator = _db_utils.generate_placeholder_parameter_pairs(
                 df=df, column_placeholders=column_placeholders, chunksize=chunksize