Source code for airflow.providers.apache.hive.sensors.metastore_partition

#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
from __future__ import annotations

from collections.abc import Sequence
from typing import TYPE_CHECKING, Any

from airflow.providers.common.sql.sensors.sql import SqlSensor

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;
"""



[docs]
class MetastorePartitionSensor(SqlSensor):
    """
    An alternative to the HivePartitionSensor that talk directly to the MySQL db.

    This was created as a result of observing sub optimal queries generated by the
    Metastore thrift service when hitting subpartitioned tables. The Thrift service's
    queries were written in a way that would not leverage the indexes.

    :param schema: the schema
    :param table: the table
    :param partition_name: the partition name, as defined in the PARTITIONS
        table of the Metastore. Order of the fields does matter.
        Examples: ``ds=2016-01-01`` or
        ``ds=2016-01-01/sub=foo`` for a sub partitioned table
    :param mysql_conn_id: a reference to the MySQL conn_id for the metastore
    """


[docs]
    template_fields: Sequence[str] = (*SqlSensor.template_fields, "partition_name", "table", "schema")


[docs]
    ui_color = "#8da7be"


    def __init__(
        self,
        *,
        table: str,
        partition_name: str,
        schema: str = "default",
        mysql_conn_id: str = "metastore_mysql",
        **kwargs: Any,
    ):

[docs]
        self.partition_name = partition_name


[docs]
        self.table = table


[docs]
        self.schema = schema

        _kwargs: dict[str, Any] = {"conn_id": mysql_conn_id, "sql": _METASTORE_PARTITION_SQL}
        if kwargs:
            _kwargs |= kwargs
        super().__init__(**_kwargs)


[docs]
    def poke(self, context: Context) -> Any:
        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)