apache-airflow-providers-google 15.1.0rc1__py3-none-any.whl → 19.1.0rc1__py3-none-any.whl

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

@@ -788,3 +788,125 @@ class DataflowJobMessagesTrigger(BaseTrigger):
             poll_sleep=self.poll_sleep,
             impersonation_chain=self.impersonation_chain,
         )
+
+
+class DataflowJobStateCompleteTrigger(BaseTrigger):
+    """
+    Trigger that monitors if a Dataflow job has reached any of successful terminal state meant for that job.
+
+    :param job_id: Required. ID of the job.
+    :param project_id: Required. The Google Cloud project ID in which the job was started.
+    :param location: Optional. The location where the job is executed. If set to None then
+        the value of DEFAULT_DATAFLOW_LOCATION will be used.
+    :param wait_until_finished: Optional. Dataflow option to block pipeline until completion.
+    :param gcp_conn_id: The connection ID to use for connecting to Google Cloud.
+    :param poll_sleep: Time (seconds) to wait between two consecutive calls to check the job.
+    :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_id: str,
+        project_id: str | None,
+        location: str = DEFAULT_DATAFLOW_LOCATION,
+        wait_until_finished: bool | None = None,
+        gcp_conn_id: str = "google_cloud_default",
+        poll_sleep: int = 10,
+        impersonation_chain: str | Sequence[str] | None = None,
+    ):
+        super().__init__()
+        self.job_id = job_id
+        self.project_id = project_id
+        self.location = location
+        self.wait_until_finished = wait_until_finished
+        self.gcp_conn_id = gcp_conn_id
+        self.poll_sleep = poll_sleep
+        self.impersonation_chain = impersonation_chain
+
+    def serialize(self) -> tuple[str, dict[str, Any]]:
+        """Serialize class arguments and classpath."""
+        return (
+            "airflow.providers.google.cloud.triggers.dataflow.DataflowJobStateCompleteTrigger",
+            {
+                "job_id": self.job_id,
+                "project_id": self.project_id,
+                "location": self.location,
+                "wait_until_finished": self.wait_until_finished,
+                "gcp_conn_id": self.gcp_conn_id,
+                "poll_sleep": self.poll_sleep,
+                "impersonation_chain": self.impersonation_chain,
+            },
+        )
+
+    async def run(self):
+        """
+        Loop until the job reaches  successful final or error state.
+
+        Yields a TriggerEvent with success status, if the job reaches successful state for own type.
+
+        Yields a TriggerEvent with error status, if the client returns an unexpected terminal
+        job status or any exception is raised while looping.
+
+        In any other case the Trigger will wait for a specified amount of time
+        stored in self.poll_sleep variable.
+        """
+        try:
+            while True:
+                job = await self.async_hook.get_job(
+                    project_id=self.project_id,
+                    job_id=self.job_id,
+                    location=self.location,
+                )
+                job_state = job.current_state.name
+                job_type_name = job.type_.name
+
+                FAILED_STATES = DataflowJobStatus.FAILED_END_STATES | {DataflowJobStatus.JOB_STATE_DRAINED}
+                if job_state in FAILED_STATES:
+                    yield TriggerEvent(
+                        {
+                            "status": "error",
+                            "message": (
+                                f"Job with id '{self.job_id}' is in failed terminal state: {job_state}"
+                            ),
+                        }
+                    )
+                    return
+
+                if self.async_hook.job_reached_terminal_state(
+                    job={"id": self.job_id, "currentState": job_state, "type": job_type_name},
+                    wait_until_finished=self.wait_until_finished,
+                ):
+                    yield TriggerEvent(
+                        {
+                            "status": "success",
+                            "message": (
+                                f"Job with id '{self.job_id}' has reached successful final state: {job_state}"
+                            ),
+                        }
+                    )
+                    return
+                self.log.info("Sleeping for %s seconds.", self.poll_sleep)
+                await asyncio.sleep(self.poll_sleep)
+        except Exception as e:
+            self.log.error("Exception occurred while checking for job state!")
+            yield TriggerEvent(
+                {
+                    "status": "error",
+                    "message": str(e),
+                }
+            )
+
+    @cached_property
+    def async_hook(self) -> AsyncDataflowHook:
+        return AsyncDataflowHook(
+            gcp_conn_id=self.gcp_conn_id,
+            poll_sleep=self.poll_sleep,
+            impersonation_chain=self.impersonation_chain,
+        )

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

@@ -86,7 +86,7 @@ class DataFusionStartPipelineTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current pipeline status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:

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

@@ -103,7 +103,13 @@ class DataplexDataQualityJobTrigger(BaseTrigger):
                 self.polling_interval_seconds,
             )
             await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
+        yield TriggerEvent(
+            {
+                "job_id": self.job_id,
+                "job_state": DataScanJob.State(state).name,
+                "job": self._convert_to_dict(job),
+            }
+        )
 
     def _convert_to_dict(self, job: DataScanJob) -> dict:
         """Return a representation of a DataScanJob instance as a dict."""
@@ -185,7 +191,13 @@ class DataplexDataProfileJobTrigger(BaseTrigger):
                 self.polling_interval_seconds,
             )
             await asyncio.sleep(self.polling_interval_seconds)
-        yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": self._convert_to_dict(job)})
+        yield TriggerEvent(
+            {
+                "job_id": self.job_id,
+                "job_state": DataScanJob.State(state).name,
+                "job": self._convert_to_dict(job),
+            }
+        )
 
     def _convert_to_dict(self, job: DataScanJob) -> dict:
         """Return a representation of a DataScanJob instance as a dict."""

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

@@ -25,21 +25,25 @@ import time
 from collections.abc import AsyncIterator, Sequence
 from typing import TYPE_CHECKING, Any
 
+from asgiref.sync import sync_to_async
 from google.api_core.exceptions import NotFound
-from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
+from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, Job, JobStatus
 
 from airflow.exceptions import AirflowException
-from airflow.models.taskinstance import TaskInstance
 from airflow.providers.google.cloud.hooks.dataproc import DataprocAsyncHook, DataprocHook
 from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.triggers.base import BaseTrigger, TriggerEvent
-from airflow.utils.session import provide_session
 from airflow.utils.state import TaskInstanceState
 
 if TYPE_CHECKING:
     from sqlalchemy.orm.session import Session
 
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.models.taskinstance import TaskInstance
+    from airflow.utils.session import provide_session
+
 
 class DataprocBaseTrigger(BaseTrigger):
     """Base class for Dataproc triggers."""
@@ -117,40 +121,67 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
             },
         )
 
-    @provide_session
-    def get_task_instance(self, session: Session) -> TaskInstance:
-        """
-        Get the task instance for the current task.
+    if not AIRFLOW_V_3_0_PLUS:
 
-        :param session: Sqlalchemy session
-        """
-        query = session.query(TaskInstance).filter(
-            TaskInstance.dag_id == self.task_instance.dag_id,
-            TaskInstance.task_id == self.task_instance.task_id,
-            TaskInstance.run_id == self.task_instance.run_id,
-            TaskInstance.map_index == self.task_instance.map_index,
+        @provide_session
+        def get_task_instance(self, session: Session) -> TaskInstance:
+            """
+            Get the task instance for the current task.
+
+            :param session: Sqlalchemy session
+            """
+            query = session.query(TaskInstance).filter(
+                TaskInstance.dag_id == self.task_instance.dag_id,
+                TaskInstance.task_id == self.task_instance.task_id,
+                TaskInstance.run_id == self.task_instance.run_id,
+                TaskInstance.map_index == self.task_instance.map_index,
+            )
+            task_instance = query.one_or_none()
+            if task_instance is None:
+                raise AirflowException(
+                    "TaskInstance with dag_id: %s,task_id: %s, run_id: %s and map_index: %s is not found",
+                    self.task_instance.dag_id,
+                    self.task_instance.task_id,
+                    self.task_instance.run_id,
+                    self.task_instance.map_index,
+                )
+            return task_instance
+
+    async def get_task_state(self):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+            dag_id=self.task_instance.dag_id,
+            task_ids=[self.task_instance.task_id],
+            run_ids=[self.task_instance.run_id],
+            map_index=self.task_instance.map_index,
         )
-        task_instance = query.one_or_none()
-        if task_instance is None:
+        try:
+            task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
+        except Exception:
             raise AirflowException(
-                "TaskInstance with dag_id: %s,task_id: %s, run_id: %s and map_index: %s is not found",
+                "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
                 self.task_instance.dag_id,
                 self.task_instance.task_id,
                 self.task_instance.run_id,
                 self.task_instance.map_index,
             )
-        return task_instance
+        return task_state
 
-    def safe_to_cancel(self) -> bool:
+    async def safe_to_cancel(self) -> bool:
         """
         Whether it is safe to cancel the external job which is being executed by this trigger.
 
         This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
         Because in those cases, we should NOT cancel the external job.
         """
-        # Database query is needed to get the latest state of the task instance.
-        task_instance = self.get_task_instance()  # type: ignore[call-arg]
-        return task_instance.state != TaskInstanceState.DEFERRED
+        if AIRFLOW_V_3_0_PLUS:
+            task_state = await self.get_task_state()
+        else:
+            # Database query is needed to get the latest state of the task instance.
+            task_instance = self.get_task_instance()  # type: ignore[call-arg]
+            task_state = task_instance.state
+        return task_state != TaskInstanceState.DEFERRED
 
     async def run(self):
         try:
@@ -163,11 +194,13 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
                 if state in (JobStatus.State.DONE, JobStatus.State.CANCELLED, JobStatus.State.ERROR):
                     break
                 await asyncio.sleep(self.polling_interval_seconds)
-            yield TriggerEvent({"job_id": self.job_id, "job_state": state, "job": job})
+            yield TriggerEvent(
+                {"job_id": self.job_id, "job_state": JobStatus.State(state).name, "job": Job.to_dict(job)}
+            )
         except asyncio.CancelledError:
             self.log.info("Task got cancelled.")
             try:
-                if self.job_id and self.cancel_on_kill and self.safe_to_cancel():
+                if self.job_id and self.cancel_on_kill and await self.safe_to_cancel():
                     self.log.info(
                         "Cancelling the job as it is safe to do so. Note that the airflow TaskInstance is not"
                         " in deferred state."
@@ -181,7 +214,12 @@ class DataprocSubmitTrigger(DataprocBaseTrigger):
                         job_id=self.job_id, project_id=self.project_id, region=self.region
                     )
                     self.log.info("Job: %s is cancelled", self.job_id)
-                    yield TriggerEvent({"job_id": self.job_id, "job_state": ClusterStatus.State.DELETING})
+                    yield TriggerEvent(
+                        {
+                            "job_id": self.job_id,
+                            "job_state": ClusterStatus.State.DELETING.name,
+                        }
+                    )
             except Exception as e:
                 self.log.error("Failed to cancel the job: %s with error : %s", self.job_id, str(e))
                 raise e
@@ -224,35 +262,62 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
             },
         )
 
-    @provide_session
-    def get_task_instance(self, session: Session) -> TaskInstance:
-        query = session.query(TaskInstance).filter(
-            TaskInstance.dag_id == self.task_instance.dag_id,
-            TaskInstance.task_id == self.task_instance.task_id,
-            TaskInstance.run_id == self.task_instance.run_id,
-            TaskInstance.map_index == self.task_instance.map_index,
+    if not AIRFLOW_V_3_0_PLUS:
+
+        @provide_session
+        def get_task_instance(self, session: Session) -> TaskInstance:
+            query = session.query(TaskInstance).filter(
+                TaskInstance.dag_id == self.task_instance.dag_id,
+                TaskInstance.task_id == self.task_instance.task_id,
+                TaskInstance.run_id == self.task_instance.run_id,
+                TaskInstance.map_index == self.task_instance.map_index,
+            )
+            task_instance = query.one_or_none()
+            if task_instance is None:
+                raise AirflowException(
+                    "TaskInstance with dag_id: %s,task_id: %s, run_id: %s and map_index: %s is not found.",
+                    self.task_instance.dag_id,
+                    self.task_instance.task_id,
+                    self.task_instance.run_id,
+                    self.task_instance.map_index,
+                )
+            return task_instance
+
+    async def get_task_state(self):
+        from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance
+
+        task_states_response = await sync_to_async(RuntimeTaskInstance.get_task_states)(
+            dag_id=self.task_instance.dag_id,
+            task_ids=[self.task_instance.task_id],
+            run_ids=[self.task_instance.run_id],
+            map_index=self.task_instance.map_index,
         )
-        task_instance = query.one_or_none()
-        if task_instance is None:
+        try:
+            task_state = task_states_response[self.task_instance.run_id][self.task_instance.task_id]
+        except Exception:
             raise AirflowException(
-                "TaskInstance with dag_id: %s,task_id: %s, run_id: %s and map_index: %s is not found.",
+                "TaskInstance with dag_id: %s, task_id: %s, run_id: %s and map_index: %s is not found",
                 self.task_instance.dag_id,
                 self.task_instance.task_id,
                 self.task_instance.run_id,
                 self.task_instance.map_index,
             )
-        return task_instance
+        return task_state
 
-    def safe_to_cancel(self) -> bool:
+    async def safe_to_cancel(self) -> bool:
         """
         Whether it is safe to cancel the external job which is being executed by this trigger.
 
         This is to avoid the case that `asyncio.CancelledError` is called because the trigger itself is stopped.
         Because in those cases, we should NOT cancel the external job.
         """
-        # Database query is needed to get the latest state of the task instance.
-        task_instance = self.get_task_instance()  # type: ignore[call-arg]
-        return task_instance.state != TaskInstanceState.DEFERRED
+        if AIRFLOW_V_3_0_PLUS:
+            task_state = await self.get_task_state()
+        else:
+            # Database query is needed to get the latest state of the task instance.
+            task_instance = self.get_task_instance()  # type: ignore[call-arg]
+            task_state = task_instance.state
+        return task_state != TaskInstanceState.DEFERRED
 
     async def run(self) -> AsyncIterator[TriggerEvent]:
         try:
@@ -264,8 +329,8 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
                     yield TriggerEvent(
                         {
                             "cluster_name": self.cluster_name,
-                            "cluster_state": ClusterStatus.State.DELETING,
-                            "cluster": cluster,
+                            "cluster_state": ClusterStatus.State.DELETING.name,  # type: ignore
+                            "cluster": Cluster.to_dict(cluster),
                         }
                     )
                     return
@@ -273,17 +338,18 @@ class DataprocClusterTrigger(DataprocBaseTrigger):
                     yield TriggerEvent(
                         {
                             "cluster_name": self.cluster_name,
-                            "cluster_state": state,
-                            "cluster": cluster,
+                            "cluster_state": ClusterStatus.State(state).name,
+                            "cluster": Cluster.to_dict(cluster),
                         }
                     )
                     return
-                self.log.info("Current state is %s", state)
-                self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
-                await asyncio.sleep(self.polling_interval_seconds)
+                else:
+                    self.log.info("Current state is %s", state)
+                    self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
+                    await asyncio.sleep(self.polling_interval_seconds)
         except asyncio.CancelledError:
             try:
-                if self.delete_on_error and self.safe_to_cancel():
+                if self.delete_on_error and await self.safe_to_cancel():
                     self.log.info(
                         "Deleting the cluster as it is safe to delete as the airflow TaskInstance is not in "
                         "deferred state."
@@ -369,12 +435,16 @@ class DataprocBatchTrigger(DataprocBaseTrigger):
 
             if state in (Batch.State.FAILED, Batch.State.SUCCEEDED, Batch.State.CANCELLED):
                 break
-            self.log.info("Current state is %s", state)
+            self.log.info("Current state is %s", Batch.State(state).name)
             self.log.info("Sleeping for %s seconds.", self.polling_interval_seconds)
             await asyncio.sleep(self.polling_interval_seconds)
 
         yield TriggerEvent(
-            {"batch_id": self.batch_id, "batch_state": state, "batch_state_message": batch.state_message}
+            {
+                "batch_id": self.batch_id,
+                "batch_state": Batch.State(state).name,
+                "batch_state_message": batch.state_message,
+            }
         )
 
 
@@ -432,9 +502,9 @@ class DataprocDeleteClusterTrigger(DataprocBaseTrigger):
         try:
             while self.end_time > time.time():
                 cluster = await self.get_async_hook().get_cluster(
-                    region=self.region,  # type: ignore[arg-type]
+                    region=self.region,
                     cluster_name=self.cluster_name,
-                    project_id=self.project_id,  # type: ignore[arg-type]
+                    project_id=self.project_id,
                     metadata=self.metadata,
                 )
                 self.log.info(

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

@@ -153,7 +153,7 @@ class GKEStartPodTrigger(KubernetesPodTrigger):
         )
 
     @cached_property
-    def hook(self) -> GKEKubernetesAsyncHook:  # type: ignore[override]
+    def hook(self) -> GKEKubernetesAsyncHook:
         return GKEKubernetesAsyncHook(
             cluster_url=self._cluster_url,
             ssl_ca_cert=self._ssl_ca_cert,
@@ -200,7 +200,7 @@ class GKEOperationTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get operation status and yields corresponding event."""
         hook = self._get_hook()
         try:
@@ -260,9 +260,10 @@ class GKEJobTrigger(BaseTrigger):
         ssl_ca_cert: str,
         job_name: str,
         job_namespace: str,
-        pod_name: str,
+        pod_names: list[str],
         pod_namespace: str,
         base_container_name: str,
+        pod_name: str | None = None,
         gcp_conn_id: str = "google_cloud_default",
         poll_interval: float = 2,
         impersonation_chain: str | Sequence[str] | None = None,
@@ -274,7 +275,13 @@ class GKEJobTrigger(BaseTrigger):
         self.ssl_ca_cert = ssl_ca_cert
         self.job_name = job_name
         self.job_namespace = job_namespace
-        self.pod_name = pod_name
+        if pod_name is not None:
+            self._pod_name = pod_name
+            self.pod_names = [
+                self.pod_name,
+            ]
+        else:
+            self.pod_names = pod_names
         self.pod_namespace = pod_namespace
         self.base_container_name = base_container_name
         self.gcp_conn_id = gcp_conn_id
@@ -283,6 +290,15 @@ class GKEJobTrigger(BaseTrigger):
         self.get_logs = get_logs
         self.do_xcom_push = do_xcom_push
 
+    @property
+    def pod_name(self):
+        warnings.warn(
+            "`pod_name` parameter is deprecated, please use `pod_names`",
+            AirflowProviderDeprecationWarning,
+            stacklevel=2,
+        )
+        return self._pod_name
+
     def serialize(self) -> tuple[str, dict[str, Any]]:
         """Serialize KubernetesCreateJobTrigger arguments and classpath."""
         return (
@@ -292,7 +308,7 @@ class GKEJobTrigger(BaseTrigger):
                 "ssl_ca_cert": self.ssl_ca_cert,
                 "job_name": self.job_name,
                 "job_namespace": self.job_namespace,
-                "pod_name": self.pod_name,
+                "pod_names": self.pod_names,
                 "pod_namespace": self.pod_namespace,
                 "base_container_name": self.base_container_name,
                 "gcp_conn_id": self.gcp_conn_id,
@@ -303,10 +319,8 @@ class GKEJobTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job status and yield a TriggerEvent."""
-        if self.get_logs or self.do_xcom_push:
-            pod = await self.hook.get_pod(name=self.pod_name, namespace=self.pod_namespace)
         if self.do_xcom_push:
             kubernetes_provider = ProvidersManager().providers["apache-airflow-providers-cncf-kubernetes"]
             kubernetes_provider_name = kubernetes_provider.data["package-name"]
@@ -318,22 +332,26 @@ class GKEJobTrigger(BaseTrigger):
                     f"package {kubernetes_provider_name}=={kubernetes_provider_version} which doesn't "
                     f"support this feature. Please upgrade it to version higher than or equal to {min_version}."
                 )
-            await self.hook.wait_until_container_complete(
-                name=self.pod_name,
-                namespace=self.pod_namespace,
-                container_name=self.base_container_name,
-                poll_interval=self.poll_interval,
-            )
-            self.log.info("Checking if xcom sidecar container is started.")
-            await self.hook.wait_until_container_started(
-                name=self.pod_name,
-                namespace=self.pod_namespace,
-                container_name=PodDefaults.SIDECAR_CONTAINER_NAME,
-                poll_interval=self.poll_interval,
-            )
-            self.log.info("Extracting result from xcom sidecar container.")
-            loop = asyncio.get_running_loop()
-            xcom_result = await loop.run_in_executor(None, self.pod_manager.extract_xcom, pod)
+            xcom_results = []
+            for pod_name in self.pod_names:
+                pod = await self.hook.get_pod(name=pod_name, namespace=self.pod_namespace)
+                await self.hook.wait_until_container_complete(
+                    name=pod_name,
+                    namespace=self.pod_namespace,
+                    container_name=self.base_container_name,
+                    poll_interval=self.poll_interval,
+                )
+                self.log.info("Checking if xcom sidecar container is started.")
+                await self.hook.wait_until_container_started(
+                    name=pod_name,
+                    namespace=self.pod_namespace,
+                    container_name=PodDefaults.SIDECAR_CONTAINER_NAME,
+                    poll_interval=self.poll_interval,
+                )
+                self.log.info("Extracting result from xcom sidecar container.")
+                loop = asyncio.get_running_loop()
+                xcom_result = await loop.run_in_executor(None, self.pod_manager.extract_xcom, pod)
+                xcom_results.append(xcom_result)
         job: V1Job = await self.hook.wait_until_job_complete(
             name=self.job_name, namespace=self.job_namespace, poll_interval=self.poll_interval
         )
@@ -345,12 +363,12 @@ class GKEJobTrigger(BaseTrigger):
             {
                 "name": job.metadata.name,
                 "namespace": job.metadata.namespace,
-                "pod_name": pod.metadata.name if self.get_logs else None,
-                "pod_namespace": pod.metadata.namespace if self.get_logs else None,
+                "pod_names": [pod_name for pod_name in self.pod_names] if self.get_logs else None,
+                "pod_namespace": self.pod_namespace if self.get_logs else None,
                 "status": status,
                 "message": message,
                 "job": job_dict,
-                "xcom_result": xcom_result if self.do_xcom_push else None,
+                "xcom_result": xcom_results if self.do_xcom_push else None,
             }
         )
 

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

@@ -90,7 +90,7 @@ class MLEngineStartTrainingJobTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
+    async def run(self) -> AsyncIterator[TriggerEvent]:
         """Get current job execution status and yields a TriggerEvent."""
         hook = self._get_async_hook()
         try:

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

@@ -85,27 +85,23 @@ class PubsubPullTrigger(BaseTrigger):
             },
         )
 
-    async def run(self) -> AsyncIterator[TriggerEvent]:  # type: ignore[override]
-        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)
+    async def run(self) -> AsyncIterator[TriggerEvent]:
+        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/cloud/utils/bigquery_get_data.py

@@ -23,9 +23,9 @@ from google.cloud.bigquery.table import Row, RowIterator
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
-    from logging import Logger
 
     from airflow.providers.google.cloud.hooks.bigquery import BigQueryHook
+    from airflow.sdk.types import Logger
 
 
 def bigquery_get_data(