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

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

@@ -22,10 +22,11 @@ import time
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any
 
+from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
+from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.datafusion import SUCCESS_STATES, DataFusionHook, PipelineStates
 from airflow.providers.google.cloud.links.datafusion import (
     DataFusionInstanceLink,
@@ -36,32 +37,11 @@ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseO
 from airflow.providers.google.cloud.triggers.datafusion import DataFusionStartPipelineTrigger
 from airflow.providers.google.cloud.utils.datafusion import DataFusionPipelineType
 from airflow.providers.google.cloud.utils.helpers import resource_path_to_dict
-from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
-from google.api_core.retry import exponential_sleep_generator
 
 if TYPE_CHECKING:
-    from airflow.utils.context import Context
-
-
-class DataFusionPipelineLinkHelper:
-    """
-    Helper class for Pipeline links.
-
-    .. warning::
-        This class is deprecated. Consider using ``resource_path_to_dict()`` instead.
-    """
-
-    @staticmethod
-    @deprecated(
-        planned_removal_date="March 01, 2025",
-        use_instead="airflow.providers.google.cloud.utils.helpers.resource_path_to_dict",
-        category=AirflowProviderDeprecationWarning,
-    )
-    def get_project_id(instance):
-        instance = instance["name"]
-        project_id = next(x for x in instance.split("/") if x.startswith("airflow"))
-        return project_id
+    from airflow.providers.common.compat.sdk import Context
+    from airflow.providers.openlineage.extractors import OperatorLineage
 
 
 class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
@@ -132,10 +112,9 @@ class CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
 
 
@@ -290,10 +269,9 @@ class CloudDataFusionCreateInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
         return instance
 
@@ -379,10 +357,9 @@ class CloudDataFusionUpdateInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
 
 
@@ -450,10 +427,9 @@ class CloudDataFusionGetInstanceOperator(GoogleCloudBaseOperator):
         project_id = resource_path_to_dict(resource_name=instance["name"])["projects"]
         DataFusionInstanceLink.persist(
             context=context,
-            task_instance=self,
             project_id=project_id,
             instance_name=self.instance_name,
-            location=self.location,
+            region=self.location,
         )
         return instance
 
@@ -540,7 +516,6 @@ class CloudDataFusionCreatePipelineOperator(GoogleCloudBaseOperator):
         )
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,
@@ -714,7 +689,6 @@ class CloudDataFusionListPipelinesOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelinesLink.persist(
             context=context,
-            task_instance=self,
             uri=service_endpoint,
             namespace=self.namespace,
         )
@@ -804,6 +778,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         self.pipeline_timeout = pipeline_timeout
         self.deferrable = deferrable
         self.poll_interval = poll_interval
+        self.pipeline_id: str | None = None
 
         if success_states:
             self.success_states = success_states
@@ -823,18 +798,17 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
             project_id=self.project_id,
         )
         api_url = instance["apiEndpoint"]
-        pipeline_id = hook.start_pipeline(
+        self.pipeline_id = hook.start_pipeline(
             pipeline_name=self.pipeline_name,
             pipeline_type=self.pipeline_type,
             instance_url=api_url,
             namespace=self.namespace,
             runtime_args=self.runtime_args,
         )
-        self.log.info("Pipeline %s submitted successfully.", pipeline_id)
+        self.log.info("Pipeline %s submitted successfully.", self.pipeline_id)
 
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,
@@ -852,7 +826,7 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
                     namespace=self.namespace,
                     pipeline_name=self.pipeline_name,
                     pipeline_type=self.pipeline_type.value,
-                    pipeline_id=pipeline_id,
+                    pipeline_id=self.pipeline_id,
                     poll_interval=self.poll_interval,
                     gcp_conn_id=self.gcp_conn_id,
                     impersonation_chain=self.impersonation_chain,
@@ -862,19 +836,21 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         else:
             if not self.asynchronous:
                 # when NOT using asynchronous mode it will just wait for pipeline to finish and print message
-                self.log.info("Waiting when pipeline %s will be in one of the success states", pipeline_id)
+                self.log.info(
+                    "Waiting when pipeline %s will be in one of the success states", self.pipeline_id
+                )
                 hook.wait_for_pipeline_state(
                     success_states=self.success_states,
-                    pipeline_id=pipeline_id,
+                    pipeline_id=self.pipeline_id,
                     pipeline_name=self.pipeline_name,
                     pipeline_type=self.pipeline_type,
                     namespace=self.namespace,
                     instance_url=api_url,
                     timeout=self.pipeline_timeout,
                 )
-                self.log.info("Pipeline %s discovered success state.", pipeline_id)
+                self.log.info("Pipeline %s discovered success state.", self.pipeline_id)
             #  otherwise, return pipeline_id so that sensor can use it later to check the pipeline state
-        return pipeline_id
+        return self.pipeline_id
 
     def execute_complete(self, context: Context, event: dict[str, Any]):
         """
@@ -891,6 +867,31 @@ class CloudDataFusionStartPipelineOperator(GoogleCloudBaseOperator):
         )
         return event["pipeline_id"]
 
+    def get_openlineage_facets_on_complete(self, task_instance) -> OperatorLineage | None:
+        """Build and return OpenLineage facets and datasets for the completed pipeline start."""
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.facets import DataFusionRunFacet
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        pipeline_resource = f"{self.project_id}:{self.location}:{self.instance_name}:{self.pipeline_name}"
+
+        inputs = [Dataset(namespace="datafusion", name=pipeline_resource)]
+
+        if self.pipeline_id:
+            output_name = f"{pipeline_resource}:{self.pipeline_id}"
+        else:
+            output_name = f"{pipeline_resource}:unknown"
+        outputs = [Dataset(namespace="datafusion", name=output_name)]
+
+        run_facets = {
+            "dataFusionRun": DataFusionRunFacet(
+                runId=self.pipeline_id,
+                runtimeArgs=self.runtime_args,
+            )
+        }
+
+        return OperatorLineage(inputs=inputs, outputs=outputs, run_facets=run_facets, job_facets={})
+
 
 class CloudDataFusionStopPipelineOperator(GoogleCloudBaseOperator):
     """
@@ -964,7 +965,6 @@ class CloudDataFusionStopPipelineOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,