apache-airflow-providers-google 12.0.0rc2__py3-none-any.whl → 14.0.0rc1__py3-none-any.whl

airflow/providers/google/cloud/operators/bigquery_dts.py

@@ -24,14 +24,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING
 
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.bigquery_datatransfer_v1 import (
-    StartManualTransferRunsResponse,
-    TransferConfig,
-    TransferRun,
-    TransferState,
-)
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigquery_dts import BiqQueryDataTransferServiceHook, get_object_id
@@ -39,11 +31,17 @@ from airflow.providers.google.cloud.links.bigquery_dts import BigQueryDataTransf
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.bigquery_dts import BigQueryDataTransferRunTrigger
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.bigquery_datatransfer_v1 import (
+    StartManualTransferRunsResponse,
+    TransferConfig,
+    TransferRun,
+    TransferState,
+)
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 def _get_transfer_config_details(config_transfer_name: str):
@@ -299,6 +297,7 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
         self.gcp_conn_id = gcp_conn_id
         self.impersonation_chain = impersonation_chain
         self.deferrable = deferrable
+        self._transfer_run: dict = {}
 
     @cached_property
     def hook(self) -> BiqQueryDataTransferServiceHook:
@@ -339,12 +338,13 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
         self.xcom_push(context, key="run_id", value=run_id)
 
         if not self.deferrable:
-            result = self._wait_for_transfer_to_be_done(
+            # Save as attribute for further use by OpenLineage
+            self._transfer_run = self._wait_for_transfer_to_be_done(
                 run_id=run_id,
                 transfer_config_id=transfer_config["config_id"],
             )
             self.log.info("Transfer run %s submitted successfully.", run_id)
-            return result
+            return self._transfer_run
 
         self.defer(
             trigger=BigQueryDataTransferRunTrigger(
@@ -412,4 +412,117 @@ class BigQueryDataTransferServiceStartTransferRunsOperator(GoogleCloudBaseOperat
             event["message"],
         )
 
-        return TransferRun.to_dict(transfer_run)
+        # Save as attribute for further use by OpenLineage
+        self._transfer_run = TransferRun.to_dict(transfer_run)
+        return self._transfer_run
+
+    def get_openlineage_facets_on_complete(self, _):
+        """Implement _on_complete as we need a run config to extract information."""
+        from urllib.parse import urlsplit
+
+        from airflow.providers.common.compat.openlineage.facet import Dataset, ErrorMessageRunFacet
+        from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            extract_ds_name_from_gcs_path,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+        from airflow.providers.openlineage.sqlparser import DatabaseInfo, SQLParser
+
+        if not self._transfer_run:
+            self.log.debug("No BigQuery Data Transfer configuration was found by OpenLineage.")
+            return OperatorLineage()
+
+        data_source_id = self._transfer_run["data_source_id"]
+        dest_dataset_id = self._transfer_run["destination_dataset_id"]
+        params = self._transfer_run["params"]
+
+        input_datasets, output_datasets = [], []
+        run_facets, job_facets = {}, {}
+        if data_source_id in ("google_cloud_storage", "amazon_s3", "azure_blob_storage"):
+            if data_source_id == "google_cloud_storage":
+                bucket, path = _parse_gcs_url(params["data_path_template"])  # gs://bucket...
+                namespace = f"gs://{bucket}"
+                name = extract_ds_name_from_gcs_path(path)
+            elif data_source_id == "amazon_s3":
+                parsed_url = urlsplit(params["data_path"])  # s3://bucket...
+                namespace = f"s3://{parsed_url.netloc}"
+                name = extract_ds_name_from_gcs_path(parsed_url.path)
+            else:  # azure_blob_storage
+                storage_account = params["storage_account"]
+                container = params["container"]
+                namespace = f"abfss://{container}@{storage_account}.dfs.core.windows.net"
+                name = extract_ds_name_from_gcs_path(params["data_path"])
+
+            input_datasets.append(Dataset(namespace=namespace, name=name))
+            dest_table_name = params["destination_table_name_template"]
+            output_datasets.append(
+                Dataset(
+                    namespace=BIGQUERY_NAMESPACE,
+                    name=f"{self.project_id}.{dest_dataset_id}.{dest_table_name}",
+                )
+            )
+        elif data_source_id in ("postgresql", "oracle", "mysql"):
+            scheme = data_source_id if data_source_id != "postgresql" else "postgres"
+            host = params["connector.endpoint.host"]
+            port = params["connector.endpoint.port"]
+
+            for asset in params["assets"]:
+                # MySQL: db/table; Other: db/schema/table;
+                table_name = asset.split("/")[-1]
+
+                input_datasets.append(
+                    Dataset(namespace=f"{scheme}://{host}:{int(port)}", name=asset.replace("/", "."))
+                )
+                output_datasets.append(
+                    Dataset(
+                        namespace=BIGQUERY_NAMESPACE, name=f"{self.project_id}.{dest_dataset_id}.{table_name}"
+                    )
+                )
+        elif data_source_id == "scheduled_query":
+            bq_db_info = DatabaseInfo(
+                scheme="bigquery",
+                authority=None,
+                database=self.project_id,
+            )
+            parser_result = SQLParser("bigquery").generate_openlineage_metadata_from_sql(
+                sql=params["query"],
+                database_info=bq_db_info,
+                database=self.project_id,
+                use_connection=False,
+                hook=None,  # Hook is not used when use_connection=False
+                sqlalchemy_engine=None,
+            )
+            if parser_result.inputs:
+                input_datasets.extend(parser_result.inputs)
+            if parser_result.outputs:
+                output_datasets.extend(parser_result.outputs)
+            if parser_result.job_facets:
+                job_facets = {**job_facets, **parser_result.job_facets}
+            if parser_result.run_facets:
+                run_facets = {**run_facets, **parser_result.run_facets}
+            dest_table_name = params.get("destination_table_name_template")
+            if dest_table_name:
+                output_datasets.append(
+                    Dataset(
+                        namespace=BIGQUERY_NAMESPACE,
+                        name=f"{self.project_id}.{dest_dataset_id}.{dest_table_name}",
+                    )
+                )
+        else:
+            self.log.debug(
+                "BigQuery Data Transfer data_source_id `%s` is not supported by OpenLineage.", data_source_id
+            )
+            return OperatorLineage()
+
+        error_status = self._transfer_run.get("error_status")
+        if error_status and str(error_status["code"]) != "0":
+            run_facets["errorMessage"] = ErrorMessageRunFacet(
+                message=error_status["message"],
+                programmingLanguage="python",
+                stackTrace=str(error_status["details"]),
+            )
+
+        return OperatorLineage(
+            inputs=input_datasets, outputs=output_datasets, job_facets=job_facets, run_facets=run_facets
+        )

airflow/providers/google/cloud/operators/bigtable.py

@@ -23,7 +23,6 @@ from collections.abc import Iterable, Sequence
 from typing import TYPE_CHECKING
 
 import google.api_core.exceptions
-
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.bigtable import BigtableHook
 from airflow.providers.google.cloud.links.bigtable import (
@@ -37,11 +36,10 @@ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 if TYPE_CHECKING:
     import enum
 
+    from airflow.utils.context import Context
     from google.cloud.bigtable import enums
     from google.cloud.bigtable.column_family import GarbageCollectionRule
 
-    from airflow.utils.context import Context
-
 
 class BigtableValidationMixin:
     """Common class for Cloud Bigtable operators for validating required fields."""

airflow/providers/google/cloud/operators/cloud_base.py

@@ -19,9 +19,8 @@
 
 from __future__ import annotations
 
-from google.api_core.gapic_v1.method import DEFAULT
-
 from airflow.models import BaseOperator
+from google.api_core.gapic_v1.method import DEFAULT
 
 
 class GoogleCloudBaseOperator(BaseOperator):

airflow/providers/google/cloud/operators/cloud_batch.py

@@ -20,18 +20,16 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.cloud.batch_v1 import Job, Task
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_batch import CloudBatchHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_batch import CloudBatchJobFinishedTrigger
+from google.cloud.batch_v1 import Job, Task
 
 if TYPE_CHECKING:
-    from google.api_core import operation
-
     from airflow.utils.context import Context
+    from google.api_core import operation
 
 
 class CloudBatchSubmitJobOperator(GoogleCloudBaseOperator):

airflow/providers/google/cloud/operators/cloud_build.py

@@ -26,9 +26,6 @@ from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 from urllib.parse import unquote, urlsplit
 
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_build import CloudBuildHook
@@ -44,11 +41,12 @@ from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_MET
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils import yaml
 from airflow.utils.helpers import exactly_one
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.devtools.cloudbuild_v1.types import Build, BuildTrigger, RepoSource
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 REGEX_REPO_PATH = re.compile(r"^/(?P<project_id>[^/]+)/(?P<repo_name>[^/]+)[\+/]*(?P<branch_name>[^:]+)?")

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

@@ -21,11 +21,6 @@ import shlex
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.api_core.exceptions import AlreadyExists
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.orchestration.airflow.service_v1 import ImageVersion
-from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_composer import CloudComposerHook
@@ -36,13 +31,16 @@ from airflow.providers.google.cloud.triggers.cloud_composer import (
     CloudComposerExecutionTrigger,
 )
 from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_METHOD_NAME
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.orchestration.airflow.service_v1 import ImageVersion
+from google.cloud.orchestration.airflow.service_v1.types import Environment, ExecuteAirflowCommandResponse
 
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
-
 CLOUD_COMPOSER_BASE_LINK = "https://console.cloud.google.com/composer/environments"
 CLOUD_COMPOSER_DETAILS_LINK = (
     CLOUD_COMPOSER_BASE_LINK + "/detail/{region}/{environment_id}/monitoring?project={project_id}"

airflow/providers/google/cloud/operators/cloud_memorystore.py

@@ -29,10 +29,6 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud.memcache_v1beta2.types import cloud_memcache
-from google.cloud.redis_v1 import FailoverInstanceRequest, InputConfig, Instance, OutputConfig
-
 from airflow.providers.google.cloud.hooks.cloud_memorystore import (
     CloudMemorystoreHook,
     CloudMemorystoreMemcachedHook,
@@ -45,13 +41,15 @@ from airflow.providers.google.cloud.links.cloud_memorystore import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud.memcache_v1beta2.types import cloud_memcache
+from google.cloud.redis_v1 import FailoverInstanceRequest, InputConfig, Instance, OutputConfig
 
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
-
 
 class CloudMemorystoreCreateInstanceOperator(GoogleCloudBaseOperator):
     """

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

@@ -21,21 +21,19 @@ from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
 import google.cloud.exceptions
-from google.api_core.exceptions import AlreadyExists
-from google.cloud.run_v2 import Job, Service
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.cloud_run import CloudRunHook, CloudRunServiceHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_run import CloudRunJobFinishedTrigger, RunJobStatus
+from google.api_core.exceptions import AlreadyExists
+from google.cloud.run_v2 import Job, Service
 
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core import operation
     from google.cloud.run_v2.types import Execution
 
-    from airflow.utils.context import Context
-
 
 class CloudRunCreateJobOperator(GoogleCloudBaseOperator):
     """
@@ -300,6 +298,9 @@ class CloudRunExecuteJobOperator(GoogleCloudBaseOperator):
             region=self.region, project_id=self.project_id, job_name=self.job_name, overrides=self.overrides
         )
 
+        if self.operation is None:
+            raise AirflowException("Operation is None")
+
         if not self.deferrable:
             result: Execution = self._wait_for_operation(self.operation)
             self._fail_if_execution_failed(result)

airflow/providers/google/cloud/operators/cloud_sql.py

@@ -20,6 +20,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Mapping, Sequence
+from contextlib import contextmanager
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
@@ -38,8 +39,7 @@ from airflow.providers.google.common.links.storage import FileDetailsLink
 
 if TYPE_CHECKING:
     from airflow.models import Connection
-    from airflow.providers.mysql.hooks.mysql import MySqlHook
-    from airflow.providers.postgres.hooks.postgres import PostgresHook
+    from airflow.providers.openlineage.extractors import OperatorLineage
     from airflow.utils.context import Context
 
 
@@ -1256,7 +1256,8 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
         self.ssl_client_key = ssl_client_key
         self.ssl_secret_id = ssl_secret_id
 
-    def _execute_query(self, hook: CloudSQLDatabaseHook, database_hook: PostgresHook | MySqlHook) -> None:
+    @contextmanager
+    def cloud_sql_proxy_context(self, hook: CloudSQLDatabaseHook):
         cloud_sql_proxy_runner = None
         try:
             if hook.use_proxy:
@@ -1266,27 +1267,27 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
                 # be taken over here by another bind(0).
                 # It's quite unlikely to happen though!
                 cloud_sql_proxy_runner.start_proxy()
-            self.log.info('Executing: "%s"', self.sql)
-            database_hook.run(self.sql, self.autocommit, parameters=self.parameters)
+            yield
         finally:
             if cloud_sql_proxy_runner:
                 cloud_sql_proxy_runner.stop_proxy()
 
     def execute(self, context: Context):
-        self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
-
         hook = self.hook
         hook.validate_ssl_certs()
         connection = hook.create_connection()
         hook.validate_socket_path_length()
         database_hook = hook.get_database_hook(connection=connection)
         try:
-            self._execute_query(hook, database_hook)
+            with self.cloud_sql_proxy_context(hook):
+                self.log.info('Executing: "%s"', self.sql)
+                database_hook.run(self.sql, self.autocommit, parameters=self.parameters)
         finally:
             hook.cleanup_database_hook()
 
     @cached_property
     def hook(self):
+        self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
         return CloudSQLDatabaseHook(
             gcp_cloudsql_conn_id=self.gcp_cloudsql_conn_id,
             gcp_conn_id=self.gcp_conn_id,
@@ -1297,3 +1298,14 @@ class CloudSQLExecuteQueryOperator(GoogleCloudBaseOperator):
             ssl_key=self.ssl_client_key,
             ssl_secret_id=self.ssl_secret_id,
         )
+
+    def get_openlineage_facets_on_complete(self, _) -> OperatorLineage | None:
+        from airflow.providers.common.compat.openlineage.utils.sql import get_openlineage_facets_with_sql
+
+        with self.cloud_sql_proxy_context(self.hook):
+            return get_openlineage_facets_with_sql(
+                hook=self.hook.db_hook,
+                sql=self.sql,  # type:ignore[arg-type]  # Iterable[str] instead of list[str]
+                conn_id=self.gcp_cloudsql_conn_id,
+                database=self.hook.database,
+            )

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

@@ -21,9 +21,10 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from copy import deepcopy
-from datetime import date, time
-from typing import TYPE_CHECKING
+from datetime import date, time, timedelta
+from typing import TYPE_CHECKING, Any
 
+from airflow.configuration import conf
 from airflow.exceptions import AirflowException
 from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
 from airflow.providers.google.cloud.hooks.cloud_storage_transfer_service import (
@@ -63,6 +64,9 @@ from airflow.providers.google.cloud.links.cloud_storage_transfer import (
     CloudStorageTransferListLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.cloud.triggers.cloud_storage_transfer_service import (
+    CloudStorageTransferServiceCheckJobStatusTrigger,
+)
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
@@ -908,6 +912,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
     :param aws_role_arn: Optional AWS role ARN for workload identity federation. This will
         override the `aws_conn_id` for authentication between GCP and AWS; see
         https://cloud.google.com/storage-transfer/docs/reference/rest/v1/TransferSpec#AwsS3Data
+    :param deferrable: Run operator in the deferrable mode.
     """
 
     template_fields: Sequence[str] = (
@@ -942,6 +947,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         google_impersonation_chain: str | Sequence[str] | None = None,
         delete_job_after_completion: bool = False,
         aws_role_arn: str | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -961,6 +967,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         self.google_impersonation_chain = google_impersonation_chain
         self.delete_job_after_completion = delete_job_after_completion
         self.aws_role_arn = aws_role_arn
+        self.deferrable = deferrable
         self._validate_inputs()
 
     def _validate_inputs(self) -> None:
@@ -979,9 +986,31 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         job = hook.create_transfer_job(body=body)
 
         if self.wait:
-            hook.wait_for_transfer_job(job, timeout=self.timeout)
-            if self.delete_job_after_completion:
-                hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
+            if not self.deferrable:
+                hook.wait_for_transfer_job(job, timeout=self.timeout)
+                if self.delete_job_after_completion:
+                    hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
+            else:
+                self.defer(
+                    timeout=timedelta(seconds=self.timeout or 60),
+                    trigger=CloudStorageTransferServiceCheckJobStatusTrigger(
+                        job_name=job[NAME],
+                        project_id=job[PROJECT_ID],
+                        gcp_conn_id=self.gcp_conn_id,
+                        impersonation_chain=self.google_impersonation_chain,
+                    ),
+                    method_name="execute_complete",
+                )
+
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
+        """
+        Act as a callback for when the trigger fires.
+
+        This returns immediately. It relies on trigger to throw an exception,
+        otherwise it assumes execution was successful.
+        """
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
 
     def _create_body(self) -> dict:
         body = {
@@ -1079,6 +1108,7 @@ class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
         account from the list granting this role to the originating account (templated).
     :param delete_job_after_completion: If True, delete the job after complete.
         If set to True, 'wait' must be set to True.
+    :param deferrable: Run operator in the deferrable mode.
     """
 
     # [START gcp_transfer_gcs_to_gcs_template_fields]
@@ -1113,6 +1143,7 @@ class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
         timeout: float | None = None,
         google_impersonation_chain: str | Sequence[str] | None = None,
         delete_job_after_completion: bool = False,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1130,6 +1161,7 @@ class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
         self.timeout = timeout
         self.google_impersonation_chain = google_impersonation_chain
         self.delete_job_after_completion = delete_job_after_completion
+        self.deferrable = deferrable
         self._validate_inputs()
 
     def _validate_inputs(self) -> None:
@@ -1149,9 +1181,31 @@ class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
         job = hook.create_transfer_job(body=body)
 
         if self.wait:
-            hook.wait_for_transfer_job(job, timeout=self.timeout)
-            if self.delete_job_after_completion:
-                hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
+            if not self.deferrable:
+                hook.wait_for_transfer_job(job, timeout=self.timeout)
+                if self.delete_job_after_completion:
+                    hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
+            else:
+                self.defer(
+                    timeout=timedelta(seconds=self.timeout or 60),
+                    trigger=CloudStorageTransferServiceCheckJobStatusTrigger(
+                        job_name=job[NAME],
+                        project_id=job[PROJECT_ID],
+                        gcp_conn_id=self.gcp_conn_id,
+                        impersonation_chain=self.google_impersonation_chain,
+                    ),
+                    method_name="execute_complete",
+                )
+
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> None:
+        """
+        Act as a callback for when the trigger fires.
+
+        This returns immediately. It relies on trigger to throw an exception,
+        otherwise it assumes execution was successful.
+        """
+        if event["status"] == "error":
+            raise AirflowException(event["message"])
 
     def _create_body(self) -> dict:
         body = {

airflow/providers/google/cloud/operators/compute.py

@@ -23,8 +23,6 @@ from collections.abc import Sequence
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
-from google.api_core import exceptions
-from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 from json_merge_patch import merge
 
 from airflow.exceptions import AirflowException
@@ -38,11 +36,12 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.cloud.utils.field_sanitizer import GcpBodyFieldSanitizer
 from airflow.providers.google.cloud.utils.field_validator import GcpBodyFieldValidator
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
+from google.api_core import exceptions
+from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 
 if TYPE_CHECKING:
-    from google.api_core.retry import Retry
-
     from airflow.utils.context import Context
+    from google.api_core.retry import Retry
 
 
 class ComputeEngineBaseOperator(GoogleCloudBaseOperator):

airflow/providers/google/cloud/operators/datacatalog.py

@@ -19,6 +19,14 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.cloud.hooks.datacatalog import CloudDataCatalogHook
+from airflow.providers.google.cloud.links.datacatalog import (
+    DataCatalogEntryGroupLink,
+    DataCatalogEntryLink,
+    DataCatalogTagTemplateLink,
+)
+from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
+from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from google.api_core.exceptions import AlreadyExists, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.datacatalog import (
@@ -32,21 +40,11 @@ from google.cloud.datacatalog import (
     TagTemplateField,
 )
 
-from airflow.providers.google.cloud.hooks.datacatalog import CloudDataCatalogHook
-from airflow.providers.google.cloud.links.datacatalog import (
-    DataCatalogEntryGroupLink,
-    DataCatalogEntryLink,
-    DataCatalogTagTemplateLink,
-)
-from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
-
 
 class CloudDataCatalogCreateEntryOperator(GoogleCloudBaseOperator):
     """