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

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

@@ -25,6 +25,7 @@ from datetime import datetime
 from typing import Any
 
 from dateutil import parser
+from google.api_core.exceptions import NotFound
 from google.cloud.orchestration.airflow.service_v1.types import ExecuteAirflowCommandResponse
 
 from airflow.exceptions import AirflowException
@@ -188,6 +189,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         impersonation_chain: str | Sequence[str] | None = None,
         poll_interval: int = 10,
         composer_airflow_version: int = 2,
+        use_rest_api: bool = False,
     ):
         super().__init__()
         self.project_id = project_id
@@ -202,6 +204,7 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
         self.impersonation_chain = impersonation_chain
         self.poll_interval = poll_interval
         self.composer_airflow_version = composer_airflow_version
+        self.use_rest_api = use_rest_api
 
     def serialize(self) -> tuple[str, dict[str, Any]]:
         return (
@@ -219,31 +222,55 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
                 "impersonation_chain": self.impersonation_chain,
                 "poll_interval": self.poll_interval,
                 "composer_airflow_version": self.composer_airflow_version,
+                "use_rest_api": self.use_rest_api,
             },
         )
 
     async def _pull_dag_runs(self) -> list[dict]:
         """Pull the list of dag runs."""
-        cmd_parameters = (
-            ["-d", self.composer_dag_id, "-o", "json"]
-            if self.composer_airflow_version < 3
-            else [self.composer_dag_id, "-o", "json"]
-        )
-        dag_runs_cmd = await self.gcp_hook.execute_airflow_command(
-            project_id=self.project_id,
-            region=self.region,
-            environment_id=self.environment_id,
-            command="dags",
-            subcommand="list-runs",
-            parameters=cmd_parameters,
-        )
-        cmd_result = await self.gcp_hook.wait_command_execution_result(
-            project_id=self.project_id,
-            region=self.region,
-            environment_id=self.environment_id,
-            execution_cmd_info=ExecuteAirflowCommandResponse.to_dict(dag_runs_cmd),
-        )
-        dag_runs = json.loads(cmd_result["output"][0]["content"])
+        if self.use_rest_api:
+            try:
+                environment = await self.gcp_hook.get_environment(
+                    project_id=self.project_id,
+                    region=self.region,
+                    environment_id=self.environment_id,
+                )
+            except NotFound as not_found_err:
+                self.log.info("The Composer environment %s does not exist.", self.environment_id)
+                raise AirflowException(not_found_err)
+            composer_airflow_uri = environment.config.airflow_uri
+
+            self.log.info(
+                "Pulling the DAG %s runs from the %s environment...",
+                self.composer_dag_id,
+                self.environment_id,
+            )
+            dag_runs_response = await self.gcp_hook.get_dag_runs(
+                composer_airflow_uri=composer_airflow_uri,
+                composer_dag_id=self.composer_dag_id,
+            )
+            dag_runs = dag_runs_response["dag_runs"]
+        else:
+            cmd_parameters = (
+                ["-d", self.composer_dag_id, "-o", "json"]
+                if self.composer_airflow_version < 3
+                else [self.composer_dag_id, "-o", "json"]
+            )
+            dag_runs_cmd = await self.gcp_hook.execute_airflow_command(
+                project_id=self.project_id,
+                region=self.region,
+                environment_id=self.environment_id,
+                command="dags",
+                subcommand="list-runs",
+                parameters=cmd_parameters,
+            )
+            cmd_result = await self.gcp_hook.wait_command_execution_result(
+                project_id=self.project_id,
+                region=self.region,
+                environment_id=self.environment_id,
+                execution_cmd_info=ExecuteAirflowCommandResponse.to_dict(dag_runs_cmd),
+            )
+            dag_runs = json.loads(cmd_result["output"][0]["content"])
         return dag_runs
 
     def _check_dag_runs_states(
@@ -271,7 +298,10 @@ class CloudComposerDAGRunTrigger(BaseTrigger):
 
     def _check_composer_dag_run_id_states(self, dag_runs: list[dict]) -> bool:
         for dag_run in dag_runs:
-            if dag_run["run_id"] == self.composer_dag_run_id and dag_run["state"] in self.allowed_states:
+            if (
+                dag_run["dag_run_id" if self.use_rest_api else "run_id"] == self.composer_dag_run_id
+                and dag_run["state"] in self.allowed_states
+            ):
                 return True
         return False
 

airflow/providers/google/cloud/triggers/cloud_run.py

@@ -134,7 +134,7 @@ class CloudRunJobFinishedTrigger(BaseTrigger):
 
         yield TriggerEvent(
             {
-                "status": RunJobStatus.TIMEOUT,
+                "status": RunJobStatus.TIMEOUT.value,
                 "job_name": self.job_name,
             }
         )

airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py

@@ -23,6 +23,7 @@ from typing import Any
 
 from google.api_core.exceptions import GoogleAPIError
 from google.cloud.storage_transfer_v1.types import TransferOperation
+from google.protobuf.json_format import MessageToDict
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
@@ -231,3 +232,92 @@ class CloudStorageTransferServiceCheckJobStatusTrigger(BaseTrigger):
         except Exception as e:
             self.log.exception("Exception occurred while checking for query completion")
             yield TriggerEvent({"status": "error", "message": str(e)})
+
+
+class CloudDataTransferServiceRunJobTrigger(BaseTrigger):
+    """
+    CloudDataTransferServiceRunJobTrigger run on the trigger worker to run Cloud Storage Transfer job.
+
+    :param job_name: The name of the transfer job
+    :param project_id: The ID of the project that owns the Transfer Job.
+    :param poke_interval: Polling period in seconds to check for the status
+    :param gcp_conn_id: The connection ID used to connect to Google Cloud.
+    :param impersonation_chain: Optional service account to impersonate using short-term
+        credentials, or chained list of accounts required to get the access_token
+        of the last account in the list, which will be impersonated in the request.
+        If set as a string, the account must grant the originating account
+        the Service Account Token Creator IAM role.
+        If set as a sequence, the identities from the list must grant
+        Service Account Token Creator IAM role to the directly preceding identity, with first
+        account from the list granting this role to the originating account (templated).
+    """
+
+    def __init__(
+        self,
+        job_name: str,
+        project_id: str = PROVIDE_PROJECT_ID,
+        poke_interval: float = 10.0,
+        gcp_conn_id: str = "google_cloud_default",
+        impersonation_chain: str | Sequence[str] | None = None,
+    ):
+        super().__init__()
+        self.job_name = job_name
+        self.project_id = project_id
+        self.poke_interval = poke_interval
+        self.gcp_conn_id = gcp_conn_id
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize CloudDataTransferServiceRunJobTrigger arguments and classpath."""
+        return (
+            f"{self.__class__.__module__}.{self.__class__.__qualname__}",
+            {
+                "job_name": self.job_name,
+                "project_id": self.project_id,
+                "poke_interval": self.poke_interval,
+                "gcp_conn_id": self.gcp_conn_id,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    def _get_async_hook(self) -> CloudDataTransferServiceAsyncHook:
+        return CloudDataTransferServiceAsyncHook(
+            project_id=self.project_id,
+            gcp_conn_id=self.gcp_conn_id,
+            impersonation_chain=self.impersonation_chain,
+        )
+
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        """Run the transfer job and yield a TriggerEvent."""
+        hook = self._get_async_hook()
+
+        try:
+            job_operation = await hook.run_transfer_job(self.job_name)
+            while True:
+                job_completed = await job_operation.done()
+                if job_completed:
+                    yield TriggerEvent(
+                        {
+                            "status": "success",
+                            "message": "Transfer operation run completed successfully",
+                            "job_result": {
+                                "name": job_operation.operation.name,
+                                "metadata": MessageToDict(
+                                    job_operation.operation.metadata, preserving_proto_field_name=True
+                                ),
+                                "response": MessageToDict(
+                                    job_operation.operation.response, preserving_proto_field_name=True
+                                ),
+                            },
+                        }
+                    )
+                    return
+
+                self.log.info(
+                    "Sleeping for %s seconds.",
+                    self.poke_interval,
+                )
+                await asyncio.sleep(self.poke_interval)
+        except Exception as e:
+            self.log.exception("Exception occurred while running transfer job")
+            yield TriggerEvent({"status": "error", "message": str(e)})

airflow/providers/google/cloud/triggers/pubsub.py

@@ -86,26 +86,22 @@ class PubsubPullTrigger(BaseTrigger):
         )
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
-        try:
-            while True:
-                if pulled_messages := await self.hook.pull(
-                    project_id=self.project_id,
-                    subscription=self.subscription,
-                    max_messages=self.max_messages,
-                    return_immediately=True,
-                ):
-                    if self.ack_messages:
-                        await self.message_acknowledgement(pulled_messages)
+        while True:
+            if pulled_messages := await self.hook.pull(
+                project_id=self.project_id,
+                subscription=self.subscription,
+                max_messages=self.max_messages,
+                return_immediately=True,
+            ):
+                if self.ack_messages:
+                    await self.message_acknowledgement(pulled_messages)
 
-                    messages_json = [ReceivedMessage.to_dict(m) for m in pulled_messages]
+                messages_json = [ReceivedMessage.to_dict(m) for m in pulled_messages]
 
-                    yield TriggerEvent({"status": "success", "message": messages_json})
-                    return
-                self.log.info("Sleeping for %s seconds.", self.poke_interval)
-                await asyncio.sleep(self.poke_interval)
-        except Exception as e:
-            yield TriggerEvent({"status": "error", "message": str(e)})
-            return
+                yield TriggerEvent({"status": "success", "message": messages_json})
+                return
+            self.log.info("Sleeping for %s seconds.", self.poke_interval)
+            await asyncio.sleep(self.poke_interval)
 
     async def message_acknowledgement(self, pulled_messages):
         await self.hook.acknowledge(

airflow/providers/google/common/hooks/base_google.py

@@ -50,12 +50,12 @@ from requests import Session
 
 from airflow import version
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseHook
 from airflow.providers.google.cloud.utils.credentials_provider import (
     _get_scopes,
     _get_target_principal_and_delegates,
     get_credentials_and_project_id,
 )
-from airflow.providers.google.version_compat import BaseHook
 from airflow.utils.process_utils import patch_environ
 
 if TYPE_CHECKING:

airflow/providers/google/get_provider_info.py

@@ -472,6 +472,12 @@ def get_provider_info():
                 ],
                 "tags": ["gcp"],
             },
+            {
+                "integration-name": "Google Cloud Generative AI",
+                "external-doc-url": "https://cloud.google.com/generative-ai-studio",
+                "how-to-guide": ["/docs/apache-airflow-providers-google/operators/cloud/gen_ai.rst"],
+                "tags": ["gcp"],
+            },
         ],
         "operators": [
             {
@@ -693,6 +699,10 @@ def get_provider_info():
                 "integration-name": "Google Cloud Logging Sink",
                 "python-modules": ["airflow.providers.google.cloud.operators.cloud_logging_sink"],
             },
+            {
+                "integration-name": "Google Cloud Generative AI",
+                "python-modules": ["airflow.providers.google.cloud.operators.gen_ai"],
+            },
         ],
         "sensors": [
             {
@@ -1057,6 +1067,10 @@ def get_provider_info():
                 "integration-name": "Google Cloud Logging",
                 "python-modules": ["airflow.providers.google.cloud.hooks.cloud_logging"],
             },
+            {
+                "integration-name": "Google Cloud Generative AI",
+                "python-modules": ["airflow.providers.google.cloud.hooks.gen_ai"],
+            },
         ],
         "triggers": [
             {
@@ -1408,6 +1422,7 @@ def get_provider_info():
             "airflow.providers.google.cloud.links.compute.ComputeInstanceDetailsLink",
             "airflow.providers.google.cloud.links.compute.ComputeInstanceTemplateDetailsLink",
             "airflow.providers.google.cloud.links.compute.ComputeInstanceGroupManagerDetailsLink",
+            "airflow.providers.google.cloud.links.cloud_run.CloudRunJobLoggingLink",
             "airflow.providers.google.cloud.links.cloud_tasks.CloudTasksQueueLink",
             "airflow.providers.google.cloud.links.cloud_tasks.CloudTasksLink",
             "airflow.providers.google.cloud.links.dataproc.DataprocLink",

airflow/providers/google/leveldb/hooks/leveldb.py

@@ -21,7 +21,7 @@ from __future__ import annotations
 from typing import Any
 
 from airflow.exceptions import AirflowException, AirflowOptionalProviderFeatureException
-from airflow.providers.google.version_compat import BaseHook
+from airflow.providers.common.compat.sdk import BaseHook
 
 try:
     import plyvel

airflow/providers/google/marketing_platform/links/analytics_admin.py

@@ -18,19 +18,13 @@ from __future__ import annotations
 
 from typing import TYPE_CHECKING, ClassVar
 
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS, BaseOperator
+from airflow.providers.common.compat.sdk import BaseOperatorLink, XCom
 
 if TYPE_CHECKING:
     from airflow.models.taskinstancekey import TaskInstanceKey
+    from airflow.providers.google.version_compat import BaseOperator
     from airflow.utils.context import Context
 
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseOperatorLink
-    from airflow.sdk.execution_time.xcom import XCom
-else:
-    from airflow.models import XCom
-    from airflow.models.baseoperatorlink import BaseOperatorLink  # type: ignore[no-redef]
-
 
 BASE_LINK = "https://analytics.google.com/analytics/web/"
 

airflow/providers/google/marketing_platform/sensors/campaign_manager.py

@@ -22,13 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.marketing_platform.hooks.campaign_manager import GoogleCampaignManagerHook
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseSensorOperator
-else:
-    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/marketing_platform/sensors/display_video.py

@@ -22,13 +22,8 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
 from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.marketing_platform.hooks.display_video import GoogleDisplayVideo360Hook
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseSensorOperator
-else:
-    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/suite/sensors/drive.py

@@ -22,13 +22,8 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.common.compat.sdk import BaseSensorOperator
 from airflow.providers.google.suite.hooks.drive import GoogleDriveHook
-from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import BaseSensorOperator
-else:
-    from airflow.sensors.base import BaseSensorOperator  # type: ignore[no-redef]
 
 if TYPE_CHECKING:
     from airflow.utils.context import Context

airflow/providers/google/version_compat.py

@@ -47,30 +47,10 @@ else:
     from airflow.hooks.base import BaseHook  # type: ignore[attr-defined,no-redef]
     from airflow.models import BaseOperator
 
-# Other SDK components: Available since 3.0+
-if AIRFLOW_V_3_0_PLUS:
-    from airflow.sdk import (
-        BaseOperatorLink,
-        BaseSensorOperator,
-        PokeReturnValue,
-    )
-else:
-    from airflow.models import BaseOperatorLink
-    from airflow.sensors.base import BaseSensorOperator, PokeReturnValue  # type: ignore[no-redef]
-
-try:
-    from airflow.sdk.execution_time.timeout import timeout
-except ImportError:
-    from airflow.utils.timeout import timeout  # type: ignore[assignment,attr-defined,no-redef]
-
 # Explicitly export these imports to protect them from being removed by linters
 __all__ = [
     "AIRFLOW_V_3_0_PLUS",
     "AIRFLOW_V_3_1_PLUS",
     "BaseHook",
     "BaseOperator",
-    "BaseSensorOperator",
-    "BaseOperatorLink",
-    "PokeReturnValue",
-    "timeout",
 ]

{apache_airflow_providers_google-18.0.0rc1.dist-info → apache_airflow_providers_google-18.1.0rc1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-google
-Version: 18.0.0rc1
+Version: 18.1.0rc1
 Summary: Provider package apache-airflow-providers-google for Apache Airflow
 Keywords: airflow-provider,google,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>
@@ -21,7 +21,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
 Classifier: Topic :: System :: Monitoring
 Requires-Dist: apache-airflow>=2.10.0rc1
-Requires-Dist: apache-airflow-providers-common-compat>=1.4.0rc1
+Requires-Dist: apache-airflow-providers-common-compat>=1.8.0rc1
 Requires-Dist: apache-airflow-providers-common-sql>=1.27.0rc1
 Requires-Dist: asgiref>=3.5.2
 Requires-Dist: dill>=0.2.3
@@ -113,8 +113,8 @@ Requires-Dist: apache-airflow-providers-sftp ; extra == "sftp"
 Requires-Dist: apache-airflow-providers-ssh ; extra == "ssh"
 Requires-Dist: apache-airflow-providers-trino ; extra == "trino"
 Project-URL: Bug Tracker, https://github.com/apache/airflow/issues
-Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-google/18.0.0/changelog.html
-Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-google/18.0.0
+Project-URL: Changelog, https://airflow.staged.apache.org/docs/apache-airflow-providers-google/18.1.0/changelog.html
+Project-URL: Documentation, https://airflow.staged.apache.org/docs/apache-airflow-providers-google/18.1.0
 Project-URL: Mastodon, https://fosstodon.org/@airflow
 Project-URL: Slack Chat, https://s.apache.org/airflow-slack
 Project-URL: Source Code, https://github.com/apache/airflow
@@ -164,7 +164,7 @@ Provides-Extra: trino
 
 Package ``apache-airflow-providers-google``
 
-Release: ``18.0.0``
+Release: ``18.1.0``
 
 
 Google services including:
@@ -184,7 +184,7 @@ This is a provider package for ``google`` provider. All classes for this provide
 are in ``airflow.providers.google`` python package.
 
 You can find package information and changelog for the provider
-in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/18.0.0/>`_.
+in the `documentation <https://airflow.apache.org/docs/apache-airflow-providers-google/18.1.0/>`_.
 
 Installation
 ------------
@@ -202,7 +202,7 @@ Requirements
 PIP package                                 Version required
 ==========================================  ======================================
 ``apache-airflow``                          ``>=2.10.0``
-``apache-airflow-providers-common-compat``  ``>=1.4.0``
+``apache-airflow-providers-common-compat``  ``>=1.8.0``
 ``apache-airflow-providers-common-sql``     ``>=1.27.0``
 ``asgiref``                                 ``>=3.5.2``
 ``dill``                                    ``>=0.2.3``
@@ -340,5 +340,5 @@ Extra                 Dependencies
 ====================  =========================================================================================================================================
 
 The changelog for the provider package can be found in the
-`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/18.0.0/changelog.html>`_.
+`changelog <https://airflow.apache.org/docs/apache-airflow-providers-google/18.1.0/changelog.html>`_.