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

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

@@ -65,13 +65,15 @@ from airflow.providers.google.cloud.links.cloud_storage_transfer import (
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.cloud_storage_transfer_service import (
+    CloudDataTransferServiceRunJobTrigger,
     CloudStorageTransferServiceCheckJobStatusTrigger,
 )
 from airflow.providers.google.cloud.utils.helpers import normalize_directory_path
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
+    from airflow.providers.openlineage.extractors import OperatorLineage
 
 
 class TransferJobPreprocessor:
@@ -96,8 +98,8 @@ class TransferJobPreprocessor:
 
         aws_hook = AwsBaseHook(self.aws_conn_id, resource_type="s3")
         aws_credentials = aws_hook.get_credentials()
-        aws_access_key_id = aws_credentials.access_key  # type: ignore[attr-defined]
-        aws_secret_access_key = aws_credentials.secret_key  # type: ignore[attr-defined]
+        aws_access_key_id = aws_credentials.access_key
+        aws_secret_access_key = aws_credentials.secret_key
         self.body[TRANSFER_SPEC][AWS_S3_DATA_SOURCE][AWS_ACCESS_KEY] = {
             ACCESS_KEY_ID: aws_access_key_id,
             SECRET_ACCESS_KEY: aws_secret_access_key,
@@ -280,7 +282,6 @@ class CloudDataTransferServiceCreateJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=result[NAME],
             )
@@ -370,7 +371,6 @@ class CloudDataTransferServiceUpdateJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=self.job_name,
             )
@@ -470,6 +470,8 @@ class CloudDataTransferServiceRunJobOperator(GoogleCloudBaseOperator):
         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).
+    :param timeout: Time to wait for the operation to end in seconds. Defaults to 60 seconds if not specified.
+    :param deferrable: Run operator in the deferrable mode.
     """
 
     # [START gcp_transfer_job_run_template_fields]
@@ -491,6 +493,8 @@ class CloudDataTransferServiceRunJobOperator(GoogleCloudBaseOperator):
         api_version: str = "v1",
         project_id: str = PROVIDE_PROJECT_ID,
         google_impersonation_chain: str | Sequence[str] | None = None,
+        timeout: float | None = None,
+        deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -499,6 +503,8 @@ class CloudDataTransferServiceRunJobOperator(GoogleCloudBaseOperator):
         self.gcp_conn_id = gcp_conn_id
         self.api_version = api_version
         self.google_impersonation_chain = google_impersonation_chain
+        self.timeout = timeout
+        self.deferrable = deferrable
 
     def _validate_inputs(self) -> None:
         if not self.job_name:
@@ -516,13 +522,36 @@ class CloudDataTransferServiceRunJobOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferJobLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 job_name=self.job_name,
             )
 
+        if self.deferrable:
+            self.defer(
+                timeout=timedelta(seconds=self.timeout or 60),
+                trigger=CloudDataTransferServiceRunJobTrigger(
+                    job_name=self.job_name,
+                    project_id=project_id,
+                    gcp_conn_id=self.gcp_conn_id,
+                    impersonation_chain=self.google_impersonation_chain,
+                ),
+                method_name="execute_complete",
+            )
+
         return hook.run_transfer_job(job_name=self.job_name, project_id=project_id)
 
+    def execute_complete(self, context: Context, event: dict[str, Any]) -> Any:
+        """
+        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"])
+
+        return event["job_result"]
+
 
 class CloudDataTransferServiceGetOperationOperator(GoogleCloudBaseOperator):
     """
@@ -589,7 +618,6 @@ class CloudDataTransferServiceGetOperationOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
                 operation_name=self.operation_name,
             )
@@ -666,7 +694,6 @@ class CloudDataTransferServiceListOperationsOperator(GoogleCloudBaseOperator):
         if project_id:
             CloudStorageTransferListLink.persist(
                 context=context,
-                task_instance=self,
                 project_id=project_id,
             )
 
@@ -969,6 +996,7 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
         self.aws_role_arn = aws_role_arn
         self.deferrable = deferrable
         self._validate_inputs()
+        self._transfer_job: dict[str, Any] | None = None
 
     def _validate_inputs(self) -> None:
         if self.delete_job_after_completion and not self.wait:
@@ -983,19 +1011,18 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
 
         TransferJobPreprocessor(body=body, aws_conn_id=self.aws_conn_id, default_schedule=True).process_body()
 
-        job = hook.create_transfer_job(body=body)
-
+        self._transfer_job = hook.create_transfer_job(body=body)
         if self.wait:
             if not self.deferrable:
-                hook.wait_for_transfer_job(job, timeout=self.timeout)
+                hook.wait_for_transfer_job(self._transfer_job, timeout=self.timeout)
                 if self.delete_job_after_completion:
-                    hook.delete_transfer_job(job_name=job[NAME], project_id=self.project_id)
+                    hook.delete_transfer_job(job_name=self._transfer_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],
+                        job_name=self._transfer_job[NAME],
+                        project_id=self._transfer_job[PROJECT_ID],
                         gcp_conn_id=self.gcp_conn_id,
                         impersonation_chain=self.google_impersonation_chain,
                     ),
@@ -1045,6 +1072,57 @@ class CloudDataTransferServiceS3ToGCSOperator(GoogleCloudBaseOperator):
 
         return body
 
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """Provide OpenLineage OperatorLineage for the S3->GCS transfer."""
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.facets import (
+            CloudStorageTransferJobFacet,
+            CloudStorageTransferRunFacet,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        input_ds = Dataset(
+            namespace=f"s3://{self.s3_bucket}",
+            name=normalize_directory_path(self.s3_path) or "",
+        )
+
+        output_ds = Dataset(
+            namespace=f"gs://{self.gcs_bucket}",
+            name=normalize_directory_path(self.gcs_path) or "",
+        )
+
+        job = self._transfer_job or {}
+        job_facet = CloudStorageTransferJobFacet(
+            jobName=job.get(NAME),
+            projectId=job.get(PROJECT_ID, self.project_id),
+            description=job.get(DESCRIPTION, self.description),
+            status=job.get(STATUS),
+            sourceBucket=job.get(TRANSFER_SPEC, {})
+            .get(AWS_S3_DATA_SOURCE, {})
+            .get(BUCKET_NAME, self.s3_bucket),
+            sourcePath=job.get(TRANSFER_SPEC, {}).get(AWS_S3_DATA_SOURCE, {}).get(PATH, self.s3_path),
+            targetBucket=job.get(TRANSFER_SPEC, {}).get(GCS_DATA_SINK, {}).get(BUCKET_NAME, self.gcs_bucket),
+            targetPath=job.get(TRANSFER_SPEC, {}).get(GCS_DATA_SINK, {}).get(PATH, self.gcs_path),
+            objectConditions=job.get(TRANSFER_SPEC, {}).get("objectConditions", self.object_conditions),
+            transferOptions=job.get(TRANSFER_SPEC, {}).get("transferOptions", self.transfer_options),
+            schedule=job.get(SCHEDULE, self.schedule),
+        )
+
+        run_facet = CloudStorageTransferRunFacet(
+            jobName=job.get(NAME),
+            wait=self.wait,
+            timeout=self.timeout,
+            deferrable=self.deferrable,
+            deleteJobAfterCompletion=self.delete_job_after_completion,
+        )
+
+        return OperatorLineage(
+            inputs=[input_ds],
+            outputs=[output_ds],
+            job_facets={"cloudStorageTransferJob": job_facet},
+            run_facets={"cloudStorageTransferRun": run_facet},
+        )
+
 
 class CloudDataTransferServiceGCSToGCSOperator(GoogleCloudBaseOperator):
     """

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

@@ -23,9 +23,11 @@ from collections.abc import Sequence
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any
 
-from json_merge_patch import merge
+from google.api_core import exceptions
+from google.cloud.compute_v1.types import Instance, InstanceGroupManager, InstanceTemplate
 
 from airflow.exceptions import AirflowException
+from airflow.providers.google._vendor.json_merge_patch import merge
 from airflow.providers.google.cloud.hooks.compute import ComputeEngineHook
 from airflow.providers.google.cloud.links.compute import (
     ComputeInstanceDetailsLink,
@@ -36,13 +38,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 airflow.utils.context import Context
     from google.api_core.retry import Retry
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class ComputeEngineBaseOperator(GoogleCloudBaseOperator):
     """Abstract base operator for Google Compute Engine operators to inherit from."""
@@ -73,6 +74,13 @@ class ComputeEngineBaseOperator(GoogleCloudBaseOperator):
         if not self.zone:
             raise AirflowException("The required parameter 'zone' is missing")
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.zone,
+            "resource_id": self.resource_id,
+        }
+
     def execute(self, context: Context):
         pass
 
@@ -224,9 +232,6 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
             self.log.info("The %s Instance already exists", self.resource_id)
             ComputeInstanceDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return Instance.to_dict(existing_instance)
@@ -246,9 +251,6 @@ class ComputeEngineInsertInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(new_instance)
@@ -396,9 +398,6 @@ class ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
             self.log.info("The %s Instance already exists", self.resource_id)
             ComputeInstanceDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return Instance.to_dict(existing_instance)
@@ -419,9 +418,6 @@ class ComputeEngineInsertInstanceFromTemplateOperator(ComputeEngineBaseOperator)
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return Instance.to_dict(new_instance_from_template)
@@ -597,9 +593,6 @@ class ComputeEngineStartInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.start_instance(zone=self.zone, resource_id=self.resource_id, project_id=self.project_id)
@@ -658,9 +651,6 @@ class ComputeEngineStopInstanceOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.stop_instance(zone=self.zone, resource_id=self.resource_id, project_id=self.project_id)
@@ -763,9 +753,6 @@ class ComputeEngineSetMachineTypeOperator(ComputeEngineBaseOperator):
         self._validate_all_body_fields()
         ComputeInstanceDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         hook.set_machine_type(
@@ -971,8 +958,6 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
             self.log.info("The %s Template already exists.", existing_template)
             ComputeInstanceTemplateDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return InstanceTemplate.to_dict(existing_template)
@@ -990,8 +975,6 @@ class ComputeEngineInsertInstanceTemplateOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceTemplateDetailsLink.persist(
             context=context,
-            task_instance=self,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return InstanceTemplate.to_dict(new_template)
@@ -1237,7 +1220,6 @@ class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
             )
             ComputeInstanceTemplateDetailsLink.persist(
                 context=context,
-                task_instance=self,
                 resource_id=self.body_patch["name"],
                 project_id=self.project_id or hook.project_id,
             )
@@ -1258,7 +1240,6 @@ class ComputeEngineCopyInstanceTemplateOperator(ComputeEngineBaseOperator):
         )
         ComputeInstanceTemplateDetailsLink.persist(
             context=context,
-            task_instance=self,
             resource_id=self.body_patch["name"],
             project_id=self.project_id or hook.project_id,
         )
@@ -1389,9 +1370,6 @@ class ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
             self.log.info("Calling patch instance template with updated body: %s", patch_body)
             ComputeInstanceGroupManagerDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
             )
             return hook.patch_instance_group_manager(
@@ -1401,16 +1379,12 @@ class ComputeEngineInstanceGroupUpdateManagerTemplateOperator(ComputeEngineBaseO
                 request_id=self.request_id,
                 project_id=self.project_id,
             )
-        else:
-            # Idempotence achieved
-            ComputeInstanceGroupManagerDetailsLink.persist(
-                context=context,
-                task_instance=self,
-                location_id=self.zone,
-                resource_id=self.resource_id,
-                project_id=self.project_id or hook.project_id,
-            )
-            return True
+        # Idempotence achieved
+        ComputeInstanceGroupManagerDetailsLink.persist(
+            context=context,
+            project_id=self.project_id or hook.project_id,
+        )
+        return True
 
 
 class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator):
@@ -1552,10 +1526,7 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
             self.log.info("The %s Instance Group Manager already exists", existing_instance_group_manager)
             ComputeInstanceGroupManagerDetailsLink.persist(
                 context=context,
-                task_instance=self,
-                resource_id=self.resource_id,
                 project_id=self.project_id or hook.project_id,
-                location_id=self.zone,
             )
             return InstanceGroupManager.to_dict(existing_instance_group_manager)
         self._field_sanitizer.sanitize(self.body)
@@ -1574,9 +1545,6 @@ class ComputeEngineInsertInstanceGroupManagerOperator(ComputeEngineBaseOperator)
         )
         ComputeInstanceGroupManagerDetailsLink.persist(
             context=context,
-            task_instance=self,
-            location_id=self.zone,
-            resource_id=self.resource_id,
             project_id=self.project_id or hook.project_id,
         )
         return InstanceGroupManager.to_dict(new_instance_group_manager)