This is an automated email from the ASF dual-hosted git repository.
henry3260 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 7fed0f07cf0 Fix Hive metastore partition sensor initialization (#69612)
7fed0f07cf0 is described below
commit 7fed0f07cf0b79159296b1d6e6ec8641a94d6ec6
Author: Henry Chen <[email protected]>
AuthorDate: Wed Jul 15 20:00:26 2026 +0800
Fix Hive metastore partition sensor initialization (#69612)
---
.../apache/hive/sensors/metastore_partition.py | 53 ++++++++++++----------
.../hive/sensors/test_metastore_partition.py | 51 ++++++++++++++++++++-
2 files changed, 79 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 6b3abee572c..d6c375795de 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 @@ if TYPE_CHECKING:
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,23 @@ class MetastorePartitionSensor(SqlSensor):
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)
+ _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:
- 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}';
- """
+ 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,
+ "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 743bdc03a53..58d5ea7d7f5 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
airflow.providers.apache.hive.sensors.metastore_partition import MetastoreP
from unit.apache.hive import DEFAULT_DATE, DEFAULT_DATE_DS, MockDBConnection,
TestHiveEnvironment
[email protected](
+ ("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,
