Skip to content

Commit 810d467

Browse files
moiseenkovmoiseenkov
authored
Implement MetastoreHivePartitionSensor (#31016)
1 parent 9cc72bb commit 810d467

9 files changed

Lines changed: 599 additions & 3 deletions

File tree

‎airflow/providers/google/cloud/hooks/dataproc_metastore.py‎

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,15 @@ def get_dataproc_metastore_client(self) -> DataprocMetastoreClient:
5353
credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
5454
)
5555

56+
def get_dataproc_metastore_client_v1beta(self):
57+
"""Returns DataprocMetastoreClient (from v1 beta)."""
58+
from google.cloud.metastore_v1beta import DataprocMetastoreClient
59+
60+
client_options = ClientOptions(api_endpoint="metastore.googleapis.com:443")
61+
return DataprocMetastoreClient(
62+
credentials=self.get_credentials(), client_info=CLIENT_INFO, client_options=client_options
63+
)
64+
5665
def wait_for_operation(self, timeout: float | None, operation: Operation):
5766
"""Waits for long-lasting operation to complete."""
5867
try:
@@ -638,3 +647,49 @@ def update_service(
638647
metadata=metadata,
639648
)
640649
return result
650+
651+
@GoogleBaseHook.fallback_to_default_project_id
652+
def list_hive_partitions(
653+
self,
654+
project_id: str,
655+
service_id: str,
656+
region: str,
657+
table: str,
658+
partition_names: list[str] | None = None,
659+
) -> Operation:
660+
"""
661+
Lists Hive partitions.
662+
663+
:param project_id: Optional. The ID of a dbt Cloud project.
664+
:param service_id: Required. Dataproc Metastore service id.
665+
:param region: Required. The ID of the Google Cloud region that the service belongs to.
666+
:param table: Required. Name of the partitioned table
667+
:param partition_names: Optional. List of table partitions to wait for.
668+
A name of a partition should look like "ds=1", or "a=1/b=2" in case of multiple partitions.
669+
Note that you cannot use logical or comparison operators as in HivePartitionSensor.
670+
If not specified then the sensor will wait for at least one partition regardless its name.
671+
"""
672+
# Remove duplicates from the `partition_names` and preserve elements order
673+
# because dictionaries are ordered since Python 3.7+
674+
_partitions = list(dict.fromkeys(partition_names)) if partition_names else []
675+
676+
query = f"""
677+
SELECT *
678+
FROM PARTITIONS
679+
INNER JOIN TBLS
680+
ON PARTITIONS.TBL_ID = TBLS.TBL_ID
681+
WHERE
682+
TBLS.TBL_NAME = '{table}'"""
683+
if _partitions:
684+
query += f"""
685+
AND PARTITIONS.PART_NAME IN ({', '.join(f"'{p}'" for p in _partitions)})"""
686+
query += ";"
687+
688+
client = self.get_dataproc_metastore_client_v1beta()
689+
result = client.query_metadata(
690+
request={
691+
"service": f"projects/{project_id}/locations/{region}/services/{service_id}",
692+
"query": query,
693+
}
694+
)
695+
return result

‎airflow/providers/google/cloud/hooks/gcs.py‎

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020

2121
import functools
2222
import gzip as gz
23+
import json
2324
import os
2425
import shutil
2526
import time
@@ -29,12 +30,12 @@
2930
from io import BytesIO
3031
from os import path
3132
from tempfile import NamedTemporaryFile
32-
from typing import IO, Callable, Generator, Sequence, TypeVar, cast, overload
33+
from typing import IO, Any, Callable, Generator, Sequence, TypeVar, cast, overload
3334
from urllib.parse import urlsplit
3435

3536
from aiohttp import ClientSession
3637
from gcloud.aio.storage import Storage
37-
from google.api_core.exceptions import NotFound
38+
from google.api_core.exceptions import GoogleAPICallError, NotFound
3839
from google.api_core.retry import Retry
3940

4041
# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
@@ -1232,6 +1233,36 @@ def gcs_object_is_directory(bucket: str) -> bool:
12321233
return len(blob) == 0 or blob.endswith("/")
12331234

12341235

1236+
def parse_json_from_gcs(gcp_conn_id: str, file_uri: str) -> Any:
1237+
"""
1238+
Downloads and parses json file from Google cloud Storage.
1239+
1240+
:param gcp_conn_id: Airflow Google Cloud connection ID.
1241+
:param file_uri: full path to json file
1242+
example: ``gs://test-bucket/dir1/dir2/file``
1243+
"""
1244+
gcs_hook = GCSHook(gcp_conn_id=gcp_conn_id)
1245+
bucket, blob = _parse_gcs_url(file_uri)
1246+
with NamedTemporaryFile(mode="w+b") as file:
1247+
try:
1248+
gcs_hook.download(bucket_name=bucket, object_name=blob, filename=file.name)
1249+
except GoogleAPICallError as ex:
1250+
raise AirflowException(f"Failed to download file with query result: {ex}")
1251+
1252+
file.seek(0)
1253+
try:
1254+
json_data = file.read()
1255+
except (ValueError, OSError, RuntimeError) as ex:
1256+
raise AirflowException(f"Failed to read file: {ex}")
1257+
1258+
try:
1259+
result = json.loads(json_data)
1260+
except json.JSONDecodeError as ex:
1261+
raise AirflowException(f"Failed to decode query result from bytes to json: {ex}")
1262+
1263+
return result
1264+
1265+
12351266
def _parse_gcs_url(gsurl: str) -> tuple[str, str]:
12361267
"""
12371268
Given a Google Cloud Storage URL (gs://<bucket>/<blob>), returns a

‎airflow/providers/google/cloud/sensors/dataproc_metastore.py‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
#
2+
# Licensed to the Apache Software Foundation (ASF) under one
3+
# or more contributor license agreements. See the NOTICE file
4+
# distributed with this work for additional information
5+
# regarding copyright ownership. The ASF licenses this file
6+
# to you under the Apache License, Version 2.0 (the
7+
# "License"); you may not use this file except in compliance
8+
# with the License. You may obtain a copy of the License at
9+
#
10+
# http://www.apache.org/licenses/LICENSE-2.0
11+
#
12+
# Unless required by applicable law or agreed to in writing,
13+
# software distributed under the License is distributed on an
14+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
15+
# KIND, either express or implied. See the License for the
16+
# specific language governing permissions and limitations
17+
# under the License.
18+
from __future__ import annotations
19+
20+
from typing import TYPE_CHECKING, Sequence
21+
22+
from google.api_core.operation import Operation
23+
24+
from airflow import AirflowException
25+
from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
26+
from airflow.providers.google.cloud.hooks.gcs import parse_json_from_gcs
27+
from airflow.sensors.base import BaseSensorOperator
28+
29+
if TYPE_CHECKING:
30+
from airflow.utils.context import Context
31+
32+
33+
class MetastoreHivePartitionSensor(BaseSensorOperator):
34+
"""
35+
Waits for partitions to show up in Hive.
36+
This sensor uses Google Cloud SDK and passes requests via gRPC.
37+
38+
:param service_id: Required. Dataproc Metastore service id.
39+
:param region: Required. The ID of the Google Cloud region that the service belongs to.
40+
:param table: Required. Name of the partitioned table
41+
:param partitions: List of table partitions to wait for.
42+
A name of a partition should look like "ds=1", or "a=1/b=2" in case of nested partitions.
43+
Note that you cannot use logical or comparison operators as in HivePartitionSensor.
44+
If not specified then the sensor will wait for at least one partition regardless its name.
45+
:param gcp_conn_id: Airflow Google Cloud connection ID.
46+
:param impersonation_chain: Optional service account to impersonate using short-term
47+
credentials, or chained list of accounts required to get the access_token
48+
of the last account in the list, which will be impersonated in the request.
49+
If set as a string, the account must grant the originating account
50+
the Service Account Token Creator IAM role.
51+
If set as a sequence, the identities from the list must grant
52+
Service Account Token Creator IAM role to the directly preceding identity, with first
53+
account from the list granting this role to the originating account.
54+
"""
55+
56+
template_fields: Sequence[str] = (
57+
"service_id",
58+
"region",
59+
"table",
60+
"partitions",
61+
"impersonation_chain",
62+
)
63+
64+
def __init__(
65+
self,
66+
service_id: str,
67+
region: str,
68+
table: str,
69+
partitions: list[str] | None,
70+
gcp_conn_id: str = "google_cloud_default",
71+
impersonation_chain: str | Sequence[str] | None = None,
72+
*args,
73+
**kwargs,
74+
):
75+
super().__init__(*args, **kwargs)
76+
self.service_id = service_id
77+
self.region = region
78+
self.table = table
79+
self.partitions = partitions or []
80+
self.gcp_conn_id = gcp_conn_id
81+
self.impersonation_chain = impersonation_chain
82+
83+
def poke(self, context: Context) -> bool:
84+
hook = DataprocMetastoreHook(
85+
gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain
86+
)
87+
operation: Operation = hook.list_hive_partitions(
88+
region=self.region, service_id=self.service_id, table=self.table, partition_names=self.partitions
89+
)
90+
metadata = hook.wait_for_operation(timeout=self.timeout, operation=operation)
91+
result_manifest_uri: str = metadata.result_manifest_uri
92+
self.log.info("Received result manifest URI: %s", result_manifest_uri)
93+
94+
self.log.info("Extracting result manifest")
95+
manifest: dict = parse_json_from_gcs(gcp_conn_id=self.gcp_conn_id, file_uri=result_manifest_uri)
96+
if not (manifest and isinstance(manifest, dict)):
97+
raise AirflowException(
98+
f"Failed to extract result manifest. "
99+
f"Expected not empty dict, but this was received: {manifest}"
100+
)
101+
102+
if manifest.get("status", {}).get("code") != 0:
103+
raise AirflowException(f"Request failed: {manifest.get('message')}")
104+
105+
# Extract actual query results
106+
result_base_uri = result_manifest_uri.rsplit("/", 1)[0]
107+
results = (f"{result_base_uri}//{filename}" for filename in manifest.get("filenames", []))
108+
found_partitions = sum(
109+
len(parse_json_from_gcs(gcp_conn_id=self.gcp_conn_id, file_uri=uri).get("rows", []))
110+
for uri in results
111+
)
112+
113+
# Return True if we got all requested partitions.
114+
# If no partitions were given in the request, then we expect to find at least one.
115+
return found_partitions > 0 and found_partitions >= len(set(self.partitions))

‎airflow/providers/google/provider.yaml‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -628,6 +628,9 @@ sensors:
628628
- integration-name: Google Dataproc
629629
python-modules:
630630
- airflow.providers.google.cloud.sensors.dataproc
631+
- integration-name: Google Dataproc Metastore
632+
python-modules:
633+
- airflow.providers.google.cloud.sensors.dataproc_metastore
631634
- integration-name: Google Cloud Storage (GCS)
632635
python-modules:
633636
- airflow.providers.google.cloud.sensors.gcs

‎docs/apache-airflow-providers-google/operators/cloud/dataproc_metastore.rst‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,3 +194,15 @@ To list backups you can use:
194194
:dedent: 4
195195
:start-after: [START how_to_cloud_dataproc_metastore_list_backups_operator]
196196
:end-before: [END how_to_cloud_dataproc_metastore_list_backups_operator]
197+
198+
Check Hive partitions existence
199+
-------------------------------
200+
201+
To check that Hive partitions have been created in the Metastore for a given table you can use:
202+
:class:`~airflow.providers.google.cloud.sensors.dataproc_metastore.MetastoreHivePartitionSensor`
203+
204+
.. exampleinclude:: /../../tests/system/providers/google/cloud/dataproc_metastore/example_dataproc_metastore_hive_partition_sensor.py
205+
:language: python
206+
:dedent: 4
207+
:start-after: [START how_to_cloud_dataproc_metastore_hive_partition_sensor]
208+
:end-before: [END how_to_cloud_dataproc_metastore_hive_partition_sensor]

0 commit comments

Comments
 (0)