apache-airflow-providers-google 18.0.0__py3-none-any.whl → 18.1.0rc1__py3-none-any.whl

airflow/providers/google/__init__.py

@@ -29,7 +29,7 @@ from airflow import __version__ as airflow_version
 
 __all__ = ["__version__"]
 
-__version__ = "18.0.0"
+__version__ = "18.1.0"
 
 if packaging.version.parse(packaging.version.parse(airflow_version).base_version) < packaging.version.parse(
     "2.10.0"

airflow/providers/google/ads/hooks/ads.py

@@ -28,14 +28,14 @@ from google.ads.googleads.errors import GoogleAdsException
 from google.auth.exceptions import GoogleAuthError
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers.google.common.hooks.base_google import get_field
-from airflow.providers.google.version_compat import BaseHook
 
 if TYPE_CHECKING:
-    from google.ads.googleads.v20.services.services.customer_service import CustomerServiceClient
-    from google.ads.googleads.v20.services.services.google_ads_service import GoogleAdsServiceClient
-    from google.ads.googleads.v20.services.services.google_ads_service.pagers import SearchPager
-    from google.ads.googleads.v20.services.types.google_ads_service import GoogleAdsRow
+    from google.ads.googleads.v21.services.services.customer_service import CustomerServiceClient
+    from google.ads.googleads.v21.services.services.google_ads_service import GoogleAdsServiceClient
+    from google.ads.googleads.v21.services.services.google_ads_service.pagers import SearchPager
+    from google.ads.googleads.v21.services.types.google_ads_service import GoogleAdsRow
 
 
 class GoogleAdsHook(BaseHook):

airflow/providers/google/assets/gcs.py

@@ -18,23 +18,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import Asset
 from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 
 if TYPE_CHECKING:
     from urllib.parse import SplitResult
 
-    from airflow.providers.common.compat.assets import Asset
     from airflow.providers.common.compat.openlineage.facet import Dataset as OpenLineageDataset
-else:
-    try:
-        from airflow.providers.common.compat.assets import Asset
-    except ImportError:
-        if AIRFLOW_V_3_0_PLUS:
-            from airflow.sdk.definitions.asset import Asset
-        else:
-            # dataset is renamed to asset since Airflow 3.0
-            from airflow.datasets import Dataset as Asset
 
 
 def create_asset(*, bucket: str, key: str, extra: dict | None = None) -> Asset:

airflow/providers/google/cloud/bundles/__init__.py

@@ -0,0 +1,16 @@
+# 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.

airflow/providers/google/cloud/bundles/gcs.py

@@ -0,0 +1,161 @@
+# 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
+
+import os
+from pathlib import Path
+
+import structlog
+from google.api_core.exceptions import NotFound
+
+from airflow.dag_processing.bundles.base import BaseDagBundle
+from airflow.exceptions import AirflowException
+from airflow.providers.google.cloud.hooks.gcs import GCSHook
+from airflow.providers.google.common.hooks.base_google import GoogleBaseHook
+
+
+class GCSDagBundle(BaseDagBundle):
+    """
+    GCS Dag bundle - exposes a directory in GCS as a Dag bundle.
+
+    This allows Airflow to load Dags directly from a GCS bucket.
+
+    :param gcp_conn_id: Airflow connection ID for GCS.  Defaults to GoogleBaseHook.default_conn_name.
+    :param bucket_name: The name of the GCS bucket containing the Dag files.
+    :param prefix:  Optional subdirectory within the GCS bucket where the Dags are stored.
+                    If None, Dags are assumed to be at the root of the bucket (Optional).
+    """
+
+    supports_versioning = False
+
+    def __init__(
+        self,
+        *,
+        gcp_conn_id: str = GoogleBaseHook.default_conn_name,
+        bucket_name: str,
+        prefix: str = "",
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.gcp_conn_id = gcp_conn_id
+        self.bucket_name = bucket_name
+        self.prefix = prefix
+        # Local path where GCS Dags are downloaded
+        self.gcs_dags_dir: Path = self.base_dir
+
+        log = structlog.get_logger(__name__)
+        self._log = log.bind(
+            bundle_name=self.name,
+            version=self.version,
+            bucket_name=self.bucket_name,
+            prefix=self.prefix,
+            gcp_conn_id=self.gcp_conn_id,
+        )
+        self._gcs_hook: GCSHook | None = None
+
+    def _initialize(self):
+        with self.lock():
+            if not self.gcs_dags_dir.exists():
+                self._log.info("Creating local Dags directory: %s", self.gcs_dags_dir)
+                os.makedirs(self.gcs_dags_dir)
+
+            if not self.gcs_dags_dir.is_dir():
+                raise NotADirectoryError(f"Local Dags path: {self.gcs_dags_dir} is not a directory.")
+
+            try:
+                self.gcs_hook.get_bucket(bucket_name=self.bucket_name)
+            except NotFound:
+                raise ValueError(f"GCS bucket '{self.bucket_name}' does not exist.")
+
+            if self.prefix:
+                # don't check when prefix is ""
+                if not self.gcs_hook.list(bucket_name=self.bucket_name, prefix=self.prefix):
+                    raise ValueError(f"GCS prefix 'gs://{self.bucket_name}/{self.prefix}' does not exist.")
+            self.refresh()
+
+    def initialize(self) -> None:
+        self._initialize()
+        super().initialize()
+
+    @property
+    def gcs_hook(self):
+        if self._gcs_hook is None:
+            try:
+                self._gcs_hook: GCSHook = GCSHook(gcp_conn_id=self.gcp_conn_id)  # Initialize GCS hook.
+            except AirflowException as e:
+                self._log.warning("Could not create GCSHook for connection %s: %s", self.gcp_conn_id, e)
+        return self._gcs_hook
+
+    def __repr__(self):
+        return (
+            f"<GCSDagBundle("
+            f"name={self.name!r}, "
+            f"bucket_name={self.bucket_name!r}, "
+            f"prefix={self.prefix!r}, "
+            f"version={self.version!r}"
+            f")>"
+        )
+
+    def get_current_version(self) -> str | None:
+        """Return the current version of the Dag bundle. Currently not supported."""
+        return None
+
+    @property
+    def path(self) -> Path:
+        """Return the local path to the Dag files."""
+        return self.gcs_dags_dir  # Path where Dags are downloaded.
+
+    def refresh(self) -> None:
+        """Refresh the Dag bundle by re-downloading the Dags from GCS."""
+        if self.version:
+            raise ValueError("Refreshing a specific version is not supported")
+
+        with self.lock():
+            self._log.debug(
+                "Downloading Dags from gs://%s/%s to %s", self.bucket_name, self.prefix, self.gcs_dags_dir
+            )
+            self.gcs_hook.sync_to_local_dir(
+                bucket_name=self.bucket_name,
+                prefix=self.prefix,
+                local_dir=self.gcs_dags_dir,
+                delete_stale=True,
+            )
+
+    def view_url(self, version: str | None = None) -> str | None:
+        """
+        Return a URL for viewing the Dags in GCS. Currently, versioning is not supported.
+
+        This method is deprecated and will be removed when the minimum supported Airflow version is 3.1.
+        Use `view_url_template` instead.
+        """
+        return self.view_url_template()
+
+    def view_url_template(self) -> str | None:
+        """Return a URL for viewing the Dags in GCS. Currently, versioning is not supported."""
+        if self.version:
+            raise ValueError("GCS url with version is not supported")
+        if hasattr(self, "_view_url_template") and self._view_url_template:
+            # Because we use this method in the view_url method, we need to handle
+            # backward compatibility for Airflow versions that doesn't have the
+            # _view_url_template attribute. Should be removed when we drop support for Airflow 3.0
+            return self._view_url_template
+        # https://console.cloud.google.com/storage/browser/<bucket-name>/<prefix>
+        url = f"https://console.cloud.google.com/storage/browser/{self.bucket_name}"
+        if self.prefix:
+            url += f"/{self.prefix}"
+
+        return url

airflow/providers/google/cloud/hooks/bigquery.py

@@ -20,17 +20,18 @@
 
 from __future__ import annotations
 
-import asyncio
 import json
 import logging
 import re
 import time
 import uuid
+import warnings
 from collections.abc import Iterable, Mapping, Sequence
 from copy import deepcopy
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any, Literal, NoReturn, cast, overload
 
+import pendulum
 from aiohttp import ClientSession as ClientSession
 from gcloud.aio.bigquery import Job, Table as Table_async
 from google.cloud.bigquery import (
@@ -75,6 +76,7 @@ from airflow.providers.google.common.hooks.base_google import (
     GoogleBaseHook,
     get_field,
 )
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.hashlib_wrapper import md5
 from airflow.utils.helpers import convert_camel_to_snake
 from airflow.utils.log.logging_mixin import LoggingMixin
@@ -86,6 +88,8 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from requests import Session
 
+    from airflow.sdk import Context
+
 log = logging.getLogger(__name__)
 
 BigQueryJob = CopyJob | QueryJob | LoadJob | ExtractJob
@@ -1274,7 +1278,16 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
             job_api_repr.result(timeout=timeout, retry=retry)
         return job_api_repr
 
-    def generate_job_id(self, job_id, dag_id, task_id, logical_date, configuration, force_rerun=False) -> str:
+    def generate_job_id(
+        self,
+        job_id: str | None,
+        dag_id: str,
+        task_id: str,
+        logical_date: datetime | None,
+        configuration: dict,
+        run_after: pendulum.DateTime | None = None,
+        force_rerun: bool = False,
+    ) -> str:
         if force_rerun:
             hash_base = str(uuid.uuid4())
         else:
@@ -1285,10 +1298,35 @@ class BigQueryHook(GoogleBaseHook, DbApiHook):
         if job_id:
             return f"{job_id}_{uniqueness_suffix}"
 
-        exec_date = logical_date.isoformat()
-        job_id = f"airflow_{dag_id}_{task_id}_{exec_date}_{uniqueness_suffix}"
+        if logical_date is not None:
+            if AIRFLOW_V_3_0_PLUS:
+                warnings.warn(
+                    "The 'logical_date' parameter is deprecated. Please use 'run_after' instead.",
+                    AirflowProviderDeprecationWarning,
+                    stacklevel=1,
+                )
+            job_id_timestamp = logical_date
+        elif run_after is not None:
+            job_id_timestamp = run_after
+        else:
+            job_id_timestamp = pendulum.now("UTC")
+
+        job_id = f"airflow_{dag_id}_{task_id}_{job_id_timestamp.isoformat()}_{uniqueness_suffix}"
         return re.sub(r"[:\-+.]", "_", job_id)
 
+    def get_run_after_or_logical_date(self, context: Context) -> pendulum.DateTime:
+        if AIRFLOW_V_3_0_PLUS:
+            if dag_run := context.get("dag_run"):
+                run_after = pendulum.instance(dag_run.run_after)
+            else:
+                run_after = pendulum.now("UTC")
+        else:
+            if logical_date := context.get("logical_date"):
+                run_after = logical_date
+            else:
+                run_after = pendulum.now("UTC")
+        return run_after
+
     def split_tablename(
         self, table_input: str, default_project_id: str, var_name: str | None = None
     ) -> tuple[str, str, str]:
@@ -1975,46 +2013,11 @@ class BigQueryAsyncHook(GoogleBaseAsyncHook):
     async def _get_job(
         self, job_id: str | None, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None
     ) -> BigQueryJob | UnknownJob:
-        """
-        Get BigQuery job by its ID, project ID and location.
-
-        WARNING.
-        This is a temporary workaround  for issues below, and it's not intended to be used elsewhere!
-            https://github.com/apache/airflow/issues/35833
-            https://github.com/talkiq/gcloud-aio/issues/584
-
-        This method was developed, because neither the `google-cloud-bigquery` nor the `gcloud-aio-bigquery`
-        provides asynchronous access to a BigQuery jobs with location parameter. That's why this method wraps
-        synchronous client call with the event loop's run_in_executor() method.
-
-        This workaround must be deleted along with the method _get_job_sync() and replaced by more robust and
-        cleaner solution in one of two cases:
-            1. The `google-cloud-bigquery` library provides async client with get_job method, that supports
-            optional parameter `location`
-            2. The `gcloud-aio-bigquery` library supports the `location` parameter in get_job() method.
-        """
-        loop = asyncio.get_event_loop()
-        job = await loop.run_in_executor(None, self._get_job_sync, job_id, project_id, location)
+        """Get BigQuery job by its ID, project ID and location."""
+        sync_hook = await self.get_sync_hook()
+        job = sync_hook.get_job(job_id=job_id, project_id=project_id, location=location)
         return job
 
-    def _get_job_sync(self, job_id, project_id, location):
-        """
-        Get BigQuery job by its ID, project ID and location synchronously.
-
-        WARNING
-        This is a temporary workaround  for issues below, and it's not intended to be used elsewhere!
-            https://github.com/apache/airflow/issues/35833
-            https://github.com/talkiq/gcloud-aio/issues/584
-
-        This workaround must be deleted along with the method _get_job() and replaced by more robust and
-        cleaner solution in one of two cases:
-            1. The `google-cloud-bigquery` library provides async client with get_job method, that supports
-            optional parameter `location`
-            2. The `gcloud-aio-bigquery` library supports the `location` parameter in get_job() method.
-        """
-        hook = BigQueryHook(**self._hook_kwargs)
-        return hook.get_job(job_id=job_id, project_id=project_id, location=location)
-
     async def get_job_status(
         self, job_id: str | None, project_id: str = PROVIDE_PROJECT_ID, location: str | None = None
     ) -> dict[str, str]:

airflow/providers/google/cloud/hooks/cloud_composer.py

@@ -24,9 +24,10 @@ from collections.abc import MutableSequence, Sequence
 from typing import TYPE_CHECKING, Any
 from urllib.parse import urljoin
 
+from aiohttp import ClientSession
 from google.api_core.client_options import ClientOptions
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.auth.transport.requests import AuthorizedSession
+from google.auth.transport.requests import AuthorizedSession, Request
 from google.cloud.orchestration.airflow.service_v1 import (
     EnvironmentsAsyncClient,
     EnvironmentsClient,
@@ -472,6 +473,38 @@ class CloudComposerHook(GoogleBaseHook, OperationHelper):
 
         return response.json()
 
+    def get_dag_runs(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of dag runs for provided DAG.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :param timeout: The timeout for this request.
+        """
+        response = self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns",
+            timeout=timeout,
+        )
+
+        if response.status_code != 200:
+            self.log.error(
+                "Failed to get DAG runs for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response.status_code,
+                response.text,
+            )
+            response.raise_for_status()
+
+        return response.json()
+
 
 class CloudComposerAsyncHook(GoogleBaseAsyncHook):
     """Hook for Google Cloud Composer async APIs."""
@@ -489,6 +522,42 @@ class CloudComposerAsyncHook(GoogleBaseAsyncHook):
             client_options=self.client_options,
         )
 
+    async def make_composer_airflow_api_request(
+        self,
+        method: str,
+        airflow_uri: str,
+        path: str,
+        data: Any | None = None,
+        timeout: float | None = None,
+    ):
+        """
+        Make a request to Cloud Composer environment's web server.
+
+        :param method: The request method to use ('GET', 'OPTIONS', 'HEAD', 'POST', 'PUT', 'PATCH', 'DELETE').
+        :param airflow_uri: The URI of the Apache Airflow Web UI hosted within this environment.
+        :param path: The path to send the request.
+        :param data: Dictionary, list of tuples, bytes, or file-like object to send in the body of the request.
+        :param timeout: The timeout for this request.
+        """
+        sync_hook = await self.get_sync_hook()
+        credentials = sync_hook.get_credentials()
+
+        if not credentials.valid:
+            credentials.refresh(Request())
+
+        async with ClientSession() as session:
+            async with session.request(
+                method=method,
+                url=urljoin(airflow_uri, path),
+                data=data,
+                headers={
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {credentials.token}",
+                },
+                timeout=timeout,
+            ) as response:
+                return await response.json(), response.status
+
     def get_environment_name(self, project_id, region, environment_id):
         return f"projects/{project_id}/locations/{region}/environments/{environment_id}"
 
@@ -594,6 +663,35 @@ class CloudComposerAsyncHook(GoogleBaseAsyncHook):
             metadata=metadata,
         )
 
+    @GoogleBaseHook.fallback_to_default_project_id
+    async def get_environment(
+        self,
+        project_id: str,
+        region: str,
+        environment_id: str,
+        retry: AsyncRetry | _MethodDefault = DEFAULT,
+        timeout: float | None = None,
+        metadata: Sequence[tuple[str, str]] = (),
+    ) -> Environment:
+        """
+        Get an existing environment.
+
+        :param project_id: Required. The ID of the Google Cloud project that the service belongs to.
+        :param region: Required. The ID of the Google Cloud region that the service belongs to.
+        :param environment_id: Required. The ID of the Google Cloud environment that the service belongs to.
+        :param retry: Designation of what errors, if any, should be retried.
+        :param timeout: The timeout for this request.
+        :param metadata: Strings which should be sent along with the request as metadata.
+        """
+        client = await self.get_environment_client()
+
+        return await client.get_environment(
+            request={"name": self.get_environment_name(project_id, region, environment_id)},
+            retry=retry,
+            timeout=timeout,
+            metadata=metadata,
+        )
+
     @GoogleBaseHook.fallback_to_default_project_id
     async def execute_airflow_command(
         self,
@@ -719,3 +817,35 @@ class CloudComposerAsyncHook(GoogleBaseAsyncHook):
 
             self.log.info("Sleeping for %s seconds.", poll_interval)
             await asyncio.sleep(poll_interval)
+
+    async def get_dag_runs(
+        self,
+        composer_airflow_uri: str,
+        composer_dag_id: str,
+        timeout: float | None = None,
+    ) -> dict:
+        """
+        Get the list of dag runs for provided DAG.
+
+        :param composer_airflow_uri: The URI of the Apache Airflow Web UI hosted within Composer environment.
+        :param composer_dag_id: The ID of DAG.
+        :param timeout: The timeout for this request.
+        """
+        response_body, response_status_code = await self.make_composer_airflow_api_request(
+            method="GET",
+            airflow_uri=composer_airflow_uri,
+            path=f"/api/v1/dags/{composer_dag_id}/dagRuns",
+            timeout=timeout,
+        )
+
+        if response_status_code != 200:
+            self.log.error(
+                "Failed to get DAG runs for dag_id=%s from %s (status=%s): %s",
+                composer_dag_id,
+                composer_airflow_uri,
+                response_status_code,
+                response_body["title"],
+            )
+            raise AirflowException(response_body["title"])
+
+        return response_body