From dedecb93722b28bbc7586400277568dad7cd0886 Mon Sep 17 00:00:00 2001 From: henry3260 Date: Wed, 8 Jul 2026 20:59:32 +0800 Subject: [PATCH 1/2] Fix Hive metastore partition sensor initialization --- .../hive/sensors/metastore_partition.py | 46 ++++++++--------- .../hive/sensors/test_metastore_partition.py | 51 ++++++++++++++++++- 2 files changed, 72 insertions(+), 25 deletions(-) diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py index 6b3abee572c49..b8cd865265bc9 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py @@ -26,6 +26,18 @@ from airflow.providers.common.compat.sdk import Context +_METASTORE_PARTITION_SQL = """ +SELECT 'X' +FROM PARTITIONS A0 +LEFT OUTER JOIN TBLS B0 ON A0.TBL_ID = B0.TBL_ID +LEFT OUTER JOIN DBS C0 ON B0.DB_ID = C0.DB_ID +WHERE + B0.TBL_NAME = %(table)s AND + C0.NAME = %(schema)s AND + A0.PART_NAME = %(partition_name)s; +""" + + class MetastorePartitionSensor(SqlSensor): """ An alternative to the HivePartitionSensor that talk directly to the MySQL db. @@ -43,7 +55,7 @@ class MetastorePartitionSensor(SqlSensor): :param mysql_conn_id: a reference to the MySQL conn_id for the metastore """ - template_fields: Sequence[str] = ("partition_name", "table", "schema") + template_fields: Sequence[str] = (*SqlSensor.template_fields, "partition_name", "table", "schema") ui_color = "#8da7be" def __init__( @@ -58,28 +70,16 @@ def __init__( self.partition_name = partition_name self.table = table self.schema = schema - self.first_poke = True - self.conn_id = mysql_conn_id - # TODO(aoen): We shouldn't be using SqlSensor here but MetastorePartitionSensor. - # The problem is the way apply_defaults works isn't compatible with inheritance. - # The inheritance model needs to be reworked in order to support overriding args/ - # kwargs with arguments here, then 'conn_id' and 'sql' can be passed into the - # constructor below and apply_defaults will no longer throw an exception. - super().__init__(**kwargs) + super().__init__(conn_id=mysql_conn_id, sql=_METASTORE_PARTITION_SQL, **kwargs) def poke(self, context: Context) -> Any: - if self.first_poke: - self.first_poke = False - if "." in self.table: - self.schema, self.table = self.table.split(".") - self.sql = f""" - SELECT 'X' - FROM PARTITIONS A0 - LEFT OUTER JOIN TBLS B0 ON A0.TBL_ID = B0.TBL_ID - LEFT OUTER JOIN DBS C0 ON B0.DB_ID = C0.DB_ID - WHERE - B0.TBL_NAME = '{self.table}' AND - C0.NAME = '{self.schema}' AND - A0.PART_NAME = '{self.partition_name}'; - """ + schema = self.schema + table = self.table + if "." in table: + schema, table = table.split(".", maxsplit=1) + self.parameters = { + "table": table, + "schema": schema, + "partition_name": self.partition_name, + } return super().poke(context) diff --git a/providers/apache/hive/tests/unit/apache/hive/sensors/test_metastore_partition.py b/providers/apache/hive/tests/unit/apache/hive/sensors/test_metastore_partition.py index 743bdc03a538f..58d5ea7d7f56c 100644 --- a/providers/apache/hive/tests/unit/apache/hive/sensors/test_metastore_partition.py +++ b/providers/apache/hive/tests/unit/apache/hive/sensors/test_metastore_partition.py @@ -27,6 +27,54 @@ from unit.apache.hive import DEFAULT_DATE, DEFAULT_DATE_DS, MockDBConnection, TestHiveEnvironment +@pytest.mark.parametrize( + ("table", "schema", "expected_parameters"), + [ + ( + "airflow.static_babynames_partitioned", + "default", + { + "table": "static_babynames_partitioned", + "schema": "airflow", + "partition_name": f"ds={DEFAULT_DATE_DS}", + }, + ), + ( + "static_babynames_partitioned", + "airflow", + { + "table": "static_babynames_partitioned", + "schema": "airflow", + "partition_name": f"ds={DEFAULT_DATE_DS}", + }, + ), + ], +) +def test_metastore_partition_sensor_uses_documented_constructor_and_parameters( + table, schema, expected_parameters +): + op = MetastorePartitionSensor( + task_id="hive_partition_check", + mysql_conn_id="test_connection_id", + table=table, + schema=schema, + partition_name=f"ds={DEFAULT_DATE_DS}", + ) + hook = mock.MagicMock() + hook.get_records.return_value = [["test_record"]] + op._get_hook = mock.MagicMock(return_value=hook) + + assert op.conn_id == "test_connection_id" + assert op.poke({}) is True + + hook.get_records.assert_called_once() + sql, parameters = hook.get_records.call_args.args + assert "B0.TBL_NAME = %(table)s" in sql + assert "C0.NAME = %(schema)s" in sql + assert "A0.PART_NAME = %(partition_name)s" in sql + assert parameters == expected_parameters + + @pytest.mark.skipif( "AIRFLOW_RUNALL_TESTS" not in os.environ, reason="Skipped because AIRFLOW_RUNALL_TESTS is not set" ) @@ -34,8 +82,7 @@ class TestHivePartitionSensor(TestHiveEnvironment): def test_hive_metastore_sql_sensor(self): op = MetastorePartitionSensor( task_id="hive_partition_check", - conn_id="test_connection_id", - sql="test_sql", + mysql_conn_id="test_connection_id", table="airflow.static_babynames_partitioned", partition_name=f"ds={DEFAULT_DATE_DS}", dag=self.dag, From 2682772cef73253aa3e8316fa11478f4ddf78bde Mon Sep 17 00:00:00 2001 From: Henry Chen Date: Tue, 14 Jul 2026 20:57:54 +0800 Subject: [PATCH 2/2] Update providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py Co-authored-by: Wei Lee --- .../apache/hive/sensors/metastore_partition.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py index b8cd865265bc9..d6c375795de1c 100644 --- a/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py +++ b/providers/apache/hive/src/airflow/providers/apache/hive/sensors/metastore_partition.py @@ -70,13 +70,20 @@ def __init__( self.partition_name = partition_name self.table = table self.schema = schema - super().__init__(conn_id=mysql_conn_id, sql=_METASTORE_PARTITION_SQL, **kwargs) + _kwargs: dict[str, Any] = {"conn_id": mysql_conn_id, "sql": _METASTORE_PARTITION_SQL} + if kwargs: + _kwargs |= kwargs + super().__init__(**_kwargs) def poke(self, context: Context) -> Any: - schema = self.schema - table = self.table - if "." in table: - schema, table = table.split(".", maxsplit=1) + if "." in self.table: + parts = self.table.split(".") + if len(parts) != 2: + raise ValueError(f"Expected 'schema.table' format, got: {self.table!r}") + schema, table = parts + else: + schema, table = self.schema, self.table + self.parameters = { "table": table, "schema": schema,