Use parameterized queries in HiveStatsCollectionOperator

HiveStatsCollectionOperator builds its bookkeeping SQL (the SELECT
and DELETE against hive_stats in the MySQL metastore, and the
SELECT ... FROM <table> WHERE <partition_key> = '<value>' against
Presto) by f-string-interpolating template-rendered fields (table,
partition, dttm) directly into raw SQL strings.

Per the security model in airflow-core/docs/security/security_model.rst
and airflow-core/docs/security/sql.rst, this is not a vulnerability --
Dag authors are trusted users responsible for sanitizing input before
passing it to operators. The change here is defense-in-depth so that
the operator does not rely on each Dag author to sanitize.

The MySQL bookkeeping SELECT and DELETE now use %s placeholders with
the parameters= kwarg of MySqlHook.get_records / .run, so the values
are bound by the driver instead of interpolated.

For the Presto SELECT, the table name and partition column names
cannot be parameterized in standard SQL. They are now validated against
^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$ for the table
and ^[A-Za-z_][A-Za-z0-9_]*$ for partition keys, raising AirflowException
on mismatch. Partition values are passed as bound parameters.

Author: Har1sh-k

providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py

@@ -18,6 +18,7 @@
 from __future__ import annotations
 
 import json
+import re
 from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Any
 
@@ -29,6 +30,13 @@
 if TYPE_CHECKING:
     from airflow.providers.common.compat.sdk import Context
 
+# Hive table names may be qualified as `<database>.<table>`; identifiers must
+# be plain word characters so they can be safely interpolated into the Presto
+# query that selects partition stats. Identifiers cannot be bound as parameters
+# in standard SQL.
+_HIVE_TABLE_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?$")
+_HIVE_COLUMN_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
+
 
 class HiveStatsCollectionOperator(BaseOperator):
     """
@@ -112,6 +120,16 @@ def get_default_exprs(self, col: str, col_type: str) -> dict[Any, Any]:
         return exp
 
     def execute(self, context: Context) -> None:
+        if not _HIVE_TABLE_RE.match(self.table):
+            raise AirflowException(
+                f"Invalid Hive table identifier: {self.table!r}. Must match {_HIVE_TABLE_RE.pattern}."
+            )
+        for partition_key in self.partition.keys():
+            if not _HIVE_COLUMN_RE.match(partition_key):
+                raise AirflowException(
+                    f"Invalid partition column name: {partition_key!r}. Must match {_HIVE_COLUMN_RE.pattern}."
+                )
+
         metastore = HiveMetastoreHook(metastore_conn_id=self.metastore_conn_id)
         table = metastore.get_table(table_name=self.table)
         field_types = {col.name: col.type for col in table.sd.cols}
@@ -128,13 +146,13 @@ def execute(self, context: Context) -> None:
         exprs.update(self.extra_exprs)
         exprs_str = ",\n        ".join(f"{v} AS {k[0]}__{k[1]}" for k, v in exprs.items())
 
-        where_clause_ = [f"{k} = '{v}'" for k, v in self.partition.items()]
+        where_clause_ = [f"{k} = %s" for k in self.partition.keys()]
         where_clause = " AND\n        ".join(where_clause_)
         sql = f"SELECT {exprs_str} FROM {self.table} WHERE {where_clause};"
 
         presto = PrestoHook(presto_conn_id=self.presto_conn_id)
         self.log.info("Executing SQL check: %s", sql)
-        row = presto.get_first(sql)
+        row = presto.get_first(sql, parameters=tuple(self.partition.values()))
         self.log.info("Record: %s", row)
         if not row:
             raise AirflowException("The query returned None")
@@ -143,23 +161,23 @@ def execute(self, context: Context) -> None:
 
         self.log.info("Deleting rows from previous runs if they exist")
         mysql = MySqlHook(self.mysql_conn_id)
-        sql = f"""
+        sql = """
         SELECT 1 FROM hive_stats
         WHERE
-            table_name='{self.table}' AND
-            partition_repr='{part_json}' AND
-            dttm='{self.dttm}'
+            table_name = %s AND
+            partition_repr = %s AND
+            dttm = %s
         LIMIT 1;
         """
-        if mysql.get_records(sql):
-            sql = f"""
+        if mysql.get_records(sql, parameters=(self.table, part_json, self.dttm)):
+            sql = """
             DELETE FROM hive_stats
             WHERE
-                table_name='{self.table}' AND
-                partition_repr='{part_json}' AND
-                dttm='{self.dttm}';
+                table_name = %s AND
+                partition_repr = %s AND
+                dttm = %s;
             """
-            mysql.run(sql)
+            mysql.run(sql, parameters=(self.table, part_json, self.dttm))
 
         self.log.info("Pivoting and loading cells into the Airflow db")
         rows = [

providers/apache/hive/tests/unit/apache/hive/operators/test_hive_stats.py

@@ -292,14 +292,97 @@ def test_execute_delete_previous_runs_rows(
         hive_stats_collection_operator = HiveStatsCollectionOperator(**self.kwargs)
         hive_stats_collection_operator.execute(context={})
 
-        sql = f"""
+        expected_sql = """
             DELETE FROM hive_stats
             WHERE
-                table_name='{hive_stats_collection_operator.table}' AND
-                partition_repr='{mock_json_dumps.return_value}' AND
-                dttm='{hive_stats_collection_operator.dttm}';
+                table_name = %s AND
+                partition_repr = %s AND
+                dttm = %s;
             """
-        mock_mysql_hook.return_value.run.assert_called_once_with(sql)
+        mock_mysql_hook.return_value.run.assert_called_once_with(
+            expected_sql,
+            parameters=(
+                hive_stats_collection_operator.table,
+                mock_json_dumps.return_value,
+                hive_stats_collection_operator.dttm,
+            ),
+        )
+
+    @patch("airflow.providers.apache.hive.operators.hive_stats.MySqlHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.PrestoHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.HiveMetastoreHook")
+    def test_execute_rejects_invalid_table_identifier(
+        self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook
+    ):
+        # The Presto SELECT interpolates the table identifier; the operator
+        # rejects any value that does not match the <db>.<table> allowlist
+        # so callers cannot smuggle whitespace or punctuation into the
+        # identifier position.
+        self.kwargs["table"] = "evil; DROP TABLE users--"
+        with pytest.raises(AirflowException, match="Invalid Hive table identifier"):
+            HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+    @patch("airflow.providers.apache.hive.operators.hive_stats.MySqlHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.PrestoHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.HiveMetastoreHook")
+    def test_execute_rejects_invalid_partition_column(
+        self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook
+    ):
+        # Partition keys reach the SELECT clause as column identifiers and
+        # are validated against the same allowlist.
+        self.kwargs["partition"] = {"evil col": "value"}
+        with pytest.raises(AirflowException, match="Invalid partition column name"):
+            HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+    @patch("airflow.providers.apache.hive.operators.hive_stats.json.dumps")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.MySqlHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.PrestoHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.HiveMetastoreHook")
+    def test_execute_parameterizes_mysql_bookkeeping_queries(
+        self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook, mock_json_dumps
+    ):
+        # The bookkeeping SELECT and DELETE against hive_stats bind table,
+        # partition_repr, and dttm as %s parameters instead of interpolating
+        # them into the SQL body, so the operator does not rely on the
+        # caller to escape those values.
+        mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols = [fake_col]
+        mock_mysql_hook.return_value.get_records.return_value = True
+
+        op = HiveStatsCollectionOperator(**self.kwargs)
+        op.execute(context={})
+
+        select_call = mock_mysql_hook.return_value.get_records.call_args
+        delete_call = mock_mysql_hook.return_value.run.call_args
+
+        select_sql = select_call.args[0]
+        delete_sql = delete_call.args[0]
+        assert "%s" in select_sql
+        assert "%s" in delete_sql
+        assert op.table not in select_sql
+        assert op.table not in delete_sql
+
+        expected_params = (op.table, mock_json_dumps.return_value, op.dttm)
+        assert select_call.kwargs["parameters"] == expected_params
+        assert delete_call.kwargs["parameters"] == expected_params
+
+    @patch("airflow.providers.apache.hive.operators.hive_stats.MySqlHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.PrestoHook")
+    @patch("airflow.providers.apache.hive.operators.hive_stats.HiveMetastoreHook")
+    def test_execute_parameterizes_presto_partition_values(
+        self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook
+    ):
+        # Partition values cannot influence the Presto SQL body — they are
+        # passed as bound parameters alongside the SELECT.
+        mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols = [fake_col]
+        mock_mysql_hook.return_value.get_records.return_value = False
+
+        self.kwargs["partition"] = {"col": "value"}
+        HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+        presto_call = mock_presto_hook.return_value.get_first.call_args
+        assert "col = %s" in presto_call.args[0]
+        assert "'value'" not in presto_call.args[0]
+        assert presto_call.kwargs["parameters"] == ("value",)
 
     @pytest.mark.skipif(
         "AIRFLOW_RUNALL_TESTS" not in os.environ, reason="Skipped because AIRFLOW_RUNALL_TESTS is not set"
@@ -326,23 +409,27 @@ def test_runs_for_hive_stats(self, mock_hive_metastore_hook):
                 op.run(start_date=DEFAULT_DATE, end_date=DEFAULT_DATE, ignore_ti_state=True)
 
         select_count_query = (
-            "SELECT COUNT(*) AS __count FROM airflow.static_babynames_partitioned WHERE ds = '2015-01-01';"
+            "SELECT COUNT(*) AS __count FROM airflow.static_babynames_partitioned WHERE ds = %s;"
         )
-        mock_presto_hook.get_first.assert_called_with(hql=select_count_query)
+        presto_call = mock_presto_hook.get_first.call_args
+        actual_presto_query = re.sub(r"\s{2,}", " ", presto_call.args[0]).strip()
+        assert actual_presto_query == select_count_query
+        assert presto_call.kwargs["parameters"] == ("2015-01-01",)
 
         expected_stats_select_query = (
-            "SELECT 1 "
-            "FROM hive_stats "
-            "WHERE table_name='airflow.static_babynames_partitioned' "
-            '  AND partition_repr=\'{"ds": "2015-01-01"}\' '
-            "  AND dttm='2015-01-01T00:00:00+00:00' "
-            "LIMIT 1;"
+            "SELECT 1 FROM hive_stats WHERE table_name = %s AND partition_repr = %s AND dttm = %s LIMIT 1;"
         )
 
-        raw_stats_select_query = mock_mysql_hook.get_records.call_args_list[0][0][0]
+        stats_select_call = mock_mysql_hook.get_records.call_args_list[0]
+        raw_stats_select_query = stats_select_call[0][0]
         actual_stats_select_query = re.sub(r"\s{2,}", " ", raw_stats_select_query).strip()
 
         assert expected_stats_select_query == actual_stats_select_query
+        assert stats_select_call.kwargs["parameters"] == (
+            "airflow.static_babynames_partitioned",
+            '{"ds": "2015-01-01"}',
+            "2015-01-01T00:00:00+00:00",
+        )
 
         insert_rows_val = [
             (