This is an automated email from the ASF dual-hosted git repository.
potiuk pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 2e30d73489c Use parameterized queries in HiveStatsCollectionOperator
(#66751)
2e30d73489c is described below
commit 2e30d73489c8dde4ea3944e80ad62647223310f5
Author: Harish Kolla <[email protected]>
AuthorDate: Wed Aug 19 06:26:14 2026 -0600
Use parameterized queries in HiveStatsCollectionOperator (#66751)
* 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, 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, partition values are passed as bound parameters
using the hook's declared placeholder (PrestoHook overrides the
DbApiHook default to "?"). Identifiers (table name and partition
column names) cannot be parameterized in standard SQL — they are
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.
* Use ValueError for invalid-identifier rejection in
HiveStatsCollectionOperator
The static check `check-no-new-airflow-exceptions` flagged the two new
`raise AirflowException` calls added in this PR (the table-identifier
allowlist and the partition-column allowlist).
Switch both to `raise ValueError`, matching the existing input-validation
pattern in the same provider (sensors/named_hive_partition.py:81 raises
`ValueError` for an invalid partition string).
This keeps the pre-existing `raise AirflowException("The query returned
None")` at line 161 untouched (that one is in the allowlist already) and
restores the AirflowException count for this file to its prior value.
* Quote Presto identifiers in HiveStatsCollectionOperator instead of
rejecting them
The previous approach rejected any table or column identifier that was not a
plain word, which broke valid Hive identifiers containing dashes or spaces.
Instead, quote every identifier interpolated into the Presto stats query --
the
table in FROM (per dotted component, so database.table still works), the
partition columns in WHERE, and the metastore columns and their aliases
projected in the SELECT -- using double quotes with embedded quotes doubled.
Plain identifiers are still emitted bare, so existing Dags are unaffected.
Partition values and the MySQL bookkeeping values stay bound as query
parameters.
* Preserve caller-quoted identifiers in HiveStatsCollectionOperator
An identifier the caller already double-quoted (e.g. "weird-col") was
re-escaped into """weird-col""", breaking the Presto stats query for users
who quoted identifiers themselves. Recognize a well-formed double-quoted
identifier and pass it through unchanged; anything else is still
double-quoted with embedded quotes doubled, so injection payloads stay
inert as a single identifier.
* Re-trigger CI (transient Docker Hub timeout)
* Split the Presto table name only on dots outside quoted identifiers
Quoting the table name per `self.table.split(".")` component treated a dot
inside a quoted identifier as a catalog/schema separator, so each fragment
matched neither the plain-identifier nor the quoted-identifier regex and
both
halves were re-escaped: `"my.table"` became `"""my"."table"""` and
`db."odd.name"` became `db."""odd"."name"""`. Plain `db.tbl` and bare `tbl`
were unaffected, so this only hit the quoted-with-dot case, which is exactly
the case the quoting support exists for.
Move the assembly into `_quote_presto_table`, which splits on the dots
*outside* a quoted identifier, leaving a quoted component that contains a
dot
intact.
Add `test_quote_presto_table` pinning the pass-through behaviour (bare,
qualified-plain, special-character, pre-quoted, embedded doubled quote, and
quoted-containing-a-dot), so a future refactor cannot quietly reintroduce
the
double-escaping, plus the dotted cases at the emitted-SQL level in
`test_execute_quotes_table_identifier`.
---
.../providers/apache/hive/operators/hive_stats.py | 101 ++++++--
.../unit/apache/hive/operators/test_hive_stats.py | 258 +++++++++++++++++++--
2 files changed, 320 insertions(+), 39 deletions(-)
diff --git
a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py
b/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py
index 443ce5038df..1d36982b54e 100644
---
a/providers/apache/hive/src/airflow/providers/apache/hive/operators/hive_stats.py
+++
b/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,50 @@ from airflow.providers.presto.hooks.presto import PrestoHook
if TYPE_CHECKING:
from airflow.providers.common.compat.sdk import Context
+# The table, the partition columns, and the metastore columns projected in the
Presto
+# stats SELECT are interpolated as identifiers, which cannot be bound as SQL
parameters.
+# Plain word identifiers are emitted unchanged, and identifiers the caller
already
+# double-quoted correctly are passed through as-is (so a pre-quoted name such
as
+# ``"weird-col"`` is not re-escaped into ``"""weird-col"""``); anything else is
+# double-quoted with embedded quotes doubled (how Presto/Trino escape
identifiers).
+# Kept local rather than reusing common.sql's ``Dialect.escape_word``, which
needs a
+# live connection and does not double embedded quotes.
+_PLAIN_IDENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*")
+# A fully and correctly double-quoted identifier: opening and closing quotes
with every
+# embedded quote doubled (e.g. ``"a""b"``). Used to detect identifiers the
caller has
+# already escaped so they are left untouched instead of being double-escaped.
+_QUOTED_IDENT_RE = re.compile(r'"(?:[^"]|"")*"')
+# One component of a qualified ``<catalog>.<schema>.<table>`` name: either a
quoted
+# identifier (which may itself contain dots) or a run of characters up to the
next
+# separating dot. Matching in that order means only the dots *outside* a quoted
+# identifier separate components, so ``"my.table"`` stays a single identifier.
+_QUALIFIED_NAME_PART_RE = re.compile(r'"(?:[^"]|"")*"|[^.]+')
+
+
+def _quote_presto_identifier(identifier: str) -> str:
+ """
+ Quote a Presto/Trino identifier.
+
+ Plain word identifiers and identifiers the caller already double-quoted
+ correctly are returned unchanged; anything else is wrapped in double quotes
+ with embedded quotes doubled.
+ """
+ if _PLAIN_IDENT_RE.fullmatch(identifier) or
_QUOTED_IDENT_RE.fullmatch(identifier):
+ return identifier
+ return '"' + identifier.replace('"', '""') + '"'
+
+
+def _quote_presto_table(table: str) -> str:
+ """
+ Quote a possibly qualified Presto/Trino table name.
+
+ The name is split into components on the dots that separate catalog,
schema and
+ table, and each component is quoted with :func:`_quote_presto_identifier`.
Dots
+ inside a quoted identifier are part of the name rather than separators, so
+ ``db."odd.name"`` keeps its two components instead of being split into
three.
+ """
+ return ".".join(_quote_presto_identifier(part) for part in
_QUALIFIED_NAME_PART_RE.findall(table))
+
class HiveStatsCollectionOperator(BaseOperator):
"""
@@ -96,18 +141,21 @@ class HiveStatsCollectionOperator(BaseOperator):
"""Get default expressions."""
if col in self.excluded_columns:
return {}
- exp = {(col, "non_null"): f"COUNT({col})"}
+ # Quote only the interpolated identifier in the SQL value; the dict
key keeps the
+ # bare column name, which is what gets stored in the
``hive_stats.col`` column.
+ quoted_col = _quote_presto_identifier(col)
+ exp = {(col, "non_null"): f"COUNT({quoted_col})"}
if col_type in {"double", "int", "bigint", "float"}:
- exp[(col, "sum")] = f"SUM({col})"
- exp[(col, "min")] = f"MIN({col})"
- exp[(col, "max")] = f"MAX({col})"
- exp[(col, "avg")] = f"AVG({col})"
+ exp[(col, "sum")] = f"SUM({quoted_col})"
+ exp[(col, "min")] = f"MIN({quoted_col})"
+ exp[(col, "max")] = f"MAX({quoted_col})"
+ exp[(col, "avg")] = f"AVG({quoted_col})"
elif col_type == "boolean":
- exp[(col, "true")] = f"SUM(CASE WHEN {col} THEN 1 ELSE 0 END)"
- exp[(col, "false")] = f"SUM(CASE WHEN NOT {col} THEN 1 ELSE 0 END)"
+ exp[(col, "true")] = f"SUM(CASE WHEN {quoted_col} THEN 1 ELSE 0
END)"
+ exp[(col, "false")] = f"SUM(CASE WHEN NOT {quoted_col} THEN 1 ELSE
0 END)"
elif col_type == "string":
- exp[(col, "len")] = f"SUM(CAST(LENGTH({col}) AS BIGINT))"
- exp[(col, "approx_distinct")] = f"APPROX_DISTINCT({col})"
+ exp[(col, "len")] = f"SUM(CAST(LENGTH({quoted_col}) AS BIGINT))"
+ exp[(col, "approx_distinct")] = f"APPROX_DISTINCT({quoted_col})"
return exp
@@ -126,15 +174,20 @@ class HiveStatsCollectionOperator(BaseOperator):
assign_exprs = self.get_default_exprs(col, col_type)
exprs.update(assign_exprs)
exprs.update(self.extra_exprs)
- exprs_str = ",\n ".join(f"{v} AS {k[0]}__{k[1]}" for k, v in
exprs.items())
+ exprs_str = ",\n ".join(
+ f"{v} AS {_quote_presto_identifier(f'{k[0]}__{k[1]}')}" for k, v
in exprs.items()
+ )
- where_clause_ = [f"{k} = '{v}'" for k, v in self.partition.items()]
+ presto = PrestoHook(presto_conn_id=self.presto_conn_id)
+ # Build the WHERE clause against the hook's declared parameter
placeholder
+ # (PrestoHook defaults to `?`; a connection may override it via the
`placeholder` extra).
+ placeholder = presto.placeholder
+ where_clause_ = [f"{_quote_presto_identifier(k)} = {placeholder}" for
k in self.partition.keys()]
where_clause = " AND\n ".join(where_clause_)
- sql = f"SELECT {exprs_str} FROM {self.table} WHERE {where_clause};"
+ sql = f"SELECT {exprs_str} FROM {_quote_presto_table(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 +196,23 @@ class HiveStatsCollectionOperator(BaseOperator):
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 = [
diff --git
a/providers/apache/hive/tests/unit/apache/hive/operators/test_hive_stats.py
b/providers/apache/hive/tests/unit/apache/hive/operators/test_hive_stats.py
index 336803814e7..5d88cd04871 100644
--- a/providers/apache/hive/tests/unit/apache/hive/operators/test_hive_stats.py
+++ b/providers/apache/hive/tests/unit/apache/hive/operators/test_hive_stats.py
@@ -23,7 +23,11 @@ from unittest.mock import MagicMock, patch
import pytest
-from airflow.providers.apache.hive.operators.hive_stats import
HiveStatsCollectionOperator
+from airflow.providers.apache.hive.operators.hive_stats import (
+ HiveStatsCollectionOperator,
+ _quote_presto_identifier,
+ _quote_presto_table,
+)
from airflow.providers.common.compat.sdk import AirflowException
from airflow.providers.presto.hooks.presto import PrestoHook
@@ -60,6 +64,63 @@ class MockPrestoHook(PrestoHook):
return self.conn
[email protected](
+ ("identifier", "expected"),
+ [
+ # Plain word identifiers are emitted unchanged.
+ ("plain_col", "plain_col"),
+ ("_underscore", "_underscore"),
+ # Anything that is not a plain word identifier is double-quoted, with
any
+ # embedded quote doubled so an injection payload stays a single inert
name.
+ ("weird-col", '"weird-col"'),
+ ("evil col", '"evil col"'),
+ ('a"b', '"a""b"'),
+ # An identifier the caller already double-quoted correctly is passed
through
+ # untouched rather than re-escaped into '"""weird-col"""'.
+ ('"weird-col"', '"weird-col"'),
+ ('"test-123-quoted"', '"test-123-quoted"'),
+ ('"a""b"', '"a""b"'),
+ # A dot inside a pre-quoted identifier is part of the name; the helper
sees one
+ # identifier and leaves it alone (splitting on dots is the caller's
job, and is
+ # covered by test_quote_presto_table below).
+ ('"my.table"', '"my.table"'),
+ # A malformed pre-quoted identifier (unbalanced / lone inner quote) is
not a
+ # valid quoted identifier, so it is escaped as a literal name instead.
+ ('"a"b"', '"""a""b"""'),
+ ],
+)
+def test_quote_presto_identifier(identifier, expected):
+ assert _quote_presto_identifier(identifier) == expected
+
+
[email protected](
+ ("table", "expected"),
+ [
+ # Bare and qualified plain names are emitted unchanged.
+ ("tbl", "tbl"),
+ ("db.tbl", "db.tbl"),
+ # Components that are not plain word identifiers are quoted one by one.
+ ("my-db.my-tbl", '"my-db"."my-tbl"'),
+ ("db.odd tbl", 'db."odd tbl"'),
+ # Pre-quoted components are passed through rather than re-escaped,
including
+ # ones carrying a correctly doubled embedded quote.
+ ('"weird-tbl"', '"weird-tbl"'),
+ ('db."a""b"', 'db."a""b"'),
+ # Only the dots outside a quoted identifier separate components, so a
quoted
+ # name containing a dot stays whole instead of being torn into two
fragments
+ # and re-escaped into nonsense such as '"""my"."table"""'.
+ ('"my.table"', '"my.table"'),
+ ('db."odd.name"', 'db."odd.name"'),
+ ('"my.db"."odd.name"', '"my.db"."odd.name"'),
+ # Trailing SQL after a closing quote is not a valid quoted identifier,
so it is
+ # escaped into one inert component instead of breaking out of the FROM
clause.
+ ('db."x"; DROP TABLE users--', 'db."x"."; DROP TABLE users--"'),
+ ],
+)
+def test_quote_presto_table(table, expected):
+ assert _quote_presto_table(table) == expected
+
+
class TestHiveStatsCollectionOperator(TestHiveEnvironment):
def setup_method(self, method):
self.kwargs = dict(
@@ -131,6 +192,7 @@ class TestHiveStatsCollectionOperator(TestHiveEnvironment):
def test_execute(self, mock_hive_metastore_hook, mock_presto_hook,
mock_mysql_hook, mock_json_dumps):
mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
hive_stats_collection_operator =
HiveStatsCollectionOperator(**self.kwargs)
hive_stats_collection_operator.execute(context={})
@@ -187,6 +249,7 @@ class TestHiveStatsCollectionOperator(TestHiveEnvironment):
self.kwargs.update(dict(assignment_func=assignment_func))
mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
hive_stats_collection_operator =
HiveStatsCollectionOperator(**self.kwargs)
hive_stats_collection_operator.execute(context={})
@@ -234,6 +297,7 @@ class TestHiveStatsCollectionOperator(TestHiveEnvironment):
self.kwargs.update(dict(assignment_func=assignment_func))
mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
hive_stats_collection_operator =
HiveStatsCollectionOperator(**self.kwargs)
hive_stats_collection_operator.execute(context={})
@@ -275,6 +339,7 @@ class TestHiveStatsCollectionOperator(TestHiveEnvironment):
mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
mock_mysql_hook.return_value.get_records.return_value = False
mock_presto_hook.return_value.get_first.return_value = None
+ mock_presto_hook.return_value.placeholder = "?"
with pytest.raises(AirflowException):
HiveStatsCollectionOperator(**self.kwargs).execute(context={})
@@ -288,18 +353,177 @@ class
TestHiveStatsCollectionOperator(TestHiveEnvironment):
):
mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
mock_mysql_hook.return_value.get_records.return_value = True
+ mock_presto_hook.return_value.placeholder = "?"
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,
+ ),
+ )
+
+ @pytest.mark.parametrize(
+ ("table", "expected_from"),
+ [
+ # Plain identifiers (including a qualified <db>.<table>) are
emitted
+ # unquoted exactly as before, so existing Dags are unaffected.
+ ("plain_table", "FROM plain_table"),
+ ("db.tbl", "FROM db.tbl"),
+ # Identifiers with other characters are double-quoted; a qualified
+ # name is quoted per dotted component.
+ ("weird-table", 'FROM "weird-table"'),
+ ("my-db.my-tbl", 'FROM "my-db"."my-tbl"'),
+ # A dot inside a quoted component is part of the identifier, not a
+ # catalog/schema separator, so the name reaches the FROM clause
intact.
+ ('"my.table"', 'FROM "my.table"'),
+ ('db."odd.name"', 'FROM db."odd.name"'),
+ # An embedded double quote is doubled, so an injection payload is
+ # contained as a single inert identifier instead of breaking out of
+ # the FROM clause (and is no longer rejected outright).
+ ('evil"; DROP TABLE users--', 'FROM "evil""; DROP TABLE users--"'),
+ ],
+ )
+ @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_quotes_table_identifier(
+ self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook,
table, expected_from
+ ):
+ mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
+ mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
+
+ self.kwargs["table"] = table
+ HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+ presto_sql = mock_presto_hook.return_value.get_first.call_args.args[0]
+ assert expected_from in presto_sql
+
+ @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_quotes_partition_column(
+ self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook
+ ):
+ # A partition key that is not a plain identifier (here it contains a
+ # space) is double-quoted in the WHERE clause while its value is still
+ # bound as a parameter, so special-character columns keep working
+ # without relying on the caller to escape them.
+ mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
+ mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
+
+ self.kwargs["partition"] = {"evil col": "value"}
+ HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+ presto_call = mock_presto_hook.return_value.get_first.call_args
+ assert '"evil col" = ?' in presto_call.args[0]
+ assert presto_call.kwargs["parameters"] == ("value",)
+
+ @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_quotes_special_char_projection_column(
+ self, mock_hive_metastore_hook, mock_presto_hook, mock_mysql_hook
+ ):
+ # A metastore column whose name is not a plain identifier (here it
contains
+ # a hyphen) is double-quoted everywhere it is interpolated into the
Presto
+ # SELECT -- both in the stat expressions and in their aliases -- so the
+ # projection no longer emits invalid SQL like COUNT(weird-col). The
bare
+ # column name is still what gets stored in hive_stats.col; the quoting
must
+ # not leak into the inserted data (the dict key stays unquoted).
+ mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[
+ _FakeCol("weird-col", "string")
+ ]
+ mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
+ # A string column yields three default exprs (non_null, len,
approx_distinct)
+ # plus the COUNT(*) entry; supply matching positional results so the
pivot
+ # that builds the inserted rows is exercised rather than zipping to
empty.
+ mock_presto_hook.return_value.get_first.return_value = [1, 2, 3, 4]
+
+ HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+ presto_sql = mock_presto_hook.return_value.get_first.call_args.args[0]
+ assert 'COUNT("weird-col")' in presto_sql
+ assert 'AS "weird-col__non_null"' in presto_sql
+ assert "COUNT(weird-col)" not in presto_sql
+
+ inserted_rows =
mock_mysql_hook.return_value.insert_rows.call_args.kwargs["rows"]
+ col_values = {row[4] for row in inserted_rows}
+ assert "weird-col" in col_values # the bare column name is stored, ...
+ assert '"weird-col"' not in col_values # ... the SQL quoting did not
leak into the data
+
+ @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. We use distinctive values for table
+ # and dttm so the absence-from-SQL assertion is not satisfied by
+ # accidental substrings of keywords like "table_name".
+ mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
+ mock_mysql_hook.return_value.get_records.return_value = True
+ mock_presto_hook.return_value.placeholder = "?"
+
+ self.kwargs["table"] = "audit_db.audit_stats_table"
+ self.kwargs["dttm"] = "audit-dttm-marker-2099"
+ 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 "audit_db.audit_stats_table" not in select_sql
+ assert "audit_db.audit_stats_table" not in delete_sql
+ assert "audit-dttm-marker-2099" not in select_sql
+ assert "audit-dttm-marker-2099" 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. PrestoHook uses
+ # `?` as its driver placeholder (not the default `%s`).
+ mock_hive_metastore_hook.return_value.get_table.return_value.sd.cols =
[fake_col]
+ mock_mysql_hook.return_value.get_records.return_value = False
+ mock_presto_hook.return_value.placeholder = "?"
+
+ self.kwargs["partition"] = {"col": "value"}
+ HiveStatsCollectionOperator(**self.kwargs).execute(context={})
+
+ presto_call = mock_presto_hook.return_value.get_first.call_args
+ assert "col = ?" 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 +550,27 @@ class
TestHiveStatsCollectionOperator(TestHiveEnvironment):
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 = ?;"
)
- 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 = [
(
