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

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

@@ -28,7 +28,7 @@ from typing import TYPE_CHECKING, Any
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.dataflow import (
     DEFAULT_DATAFLOW_LOCATION,
     DataflowHook,
@@ -43,7 +43,7 @@ from airflow.providers.google.common.consts import GOOGLE_DEFAULT_DEFERRABLE_MET
 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
 
 
 class CheckJobRunning(Enum):
@@ -383,7 +383,12 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
     def execute(self, context: Context):
         def set_current_job(current_job):
             self.job = current_job
-            DataflowJobLink.persist(self, context, self.project_id, self.location, self.job.get("id"))
+            DataflowJobLink.persist(
+                context=context,
+                project_id=self.project_id,
+                region=self.location,
+                job_id=self.job.get("id"),
+            )
 
         options = self.dataflow_default_options
         options.update(self.options)
@@ -404,7 +409,7 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
                 append_job_name=self.append_job_name,
             )
             job_id = self.hook.extract_job_id(self.job)
-            self.xcom_push(context, key="job_id", value=job_id)
+            context["task_instance"].xcom_push(key="job_id", value=job_id)
             return job_id
 
         self.job = self.hook.launch_job_with_template(
@@ -418,7 +423,9 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
             environment=self.environment,
         )
         job_id = self.hook.extract_job_id(self.job)
-        DataflowJobLink.persist(self, context, self.project_id, self.location, job_id)
+        DataflowJobLink.persist(
+            context=context, project_id=self.project_id, region=self.location, job_id=job_id
+        )
         self.defer(
             trigger=TemplateJobStartTrigger(
                 project_id=self.project_id,
@@ -439,7 +446,7 @@ class DataflowTemplatedJobStartOperator(GoogleCloudBaseOperator):
             raise AirflowException(event["message"])
 
         job_id = event["job_id"]
-        self.xcom_push(context, key="job_id", value=job_id)
+        context["task_instance"].xcom_push(key="job_id", value=job_id)
         self.log.info("Task %s completed with response %s", self.task_id, event["message"])
         return job_id
 
@@ -576,6 +583,7 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
     def hook(self) -> DataflowHook:
         hook = DataflowHook(
             gcp_conn_id=self.gcp_conn_id,
+            poll_sleep=self.poll_sleep,
             drain_pipeline=self.drain_pipeline,
             cancel_timeout=self.cancel_timeout,
             wait_until_finished=self.wait_until_finished,
@@ -590,7 +598,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
 
         def set_current_job(current_job):
             self.job = current_job
-            DataflowJobLink.persist(self, context, self.project_id, self.location, self.job.get("id"))
+            DataflowJobLink.persist(
+                context=context, project_id=self.project_id, region=self.location, job_id=self.job.get("id")
+            )
 
         if not self.deferrable:
             self.job = self.hook.start_flex_template(
@@ -600,7 +610,7 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
                 on_new_job_callback=set_current_job,
             )
             job_id = self.hook.extract_job_id(self.job)
-            self.xcom_push(context, key="job_id", value=job_id)
+            context["task_instance"].xcom_push(key="job_id", value=job_id)
             return self.job
 
         self.job = self.hook.launch_job_with_flex_template(
@@ -609,7 +619,9 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
             project_id=self.project_id,
         )
         job_id = self.hook.extract_job_id(self.job)
-        DataflowJobLink.persist(self, context, self.project_id, self.location, job_id)
+        DataflowJobLink.persist(
+            context=context, project_id=self.project_id, region=self.location, job_id=job_id
+        )
         self.defer(
             trigger=TemplateJobStartTrigger(
                 project_id=self.project_id,
@@ -639,7 +651,7 @@ class DataflowStartFlexTemplateOperator(GoogleCloudBaseOperator):
 
         job_id = event["job_id"]
         self.log.info("Task %s completed with response %s", job_id, event["message"])
-        self.xcom_push(context, key="job_id", value=job_id)
+        context["task_instance"].xcom_push(key="job_id", value=job_id)
         job = self.hook.get_job(job_id=job_id, project_id=self.project_id, location=self.location)
         return job
 
@@ -764,7 +776,9 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
             location=self.region,
         )
 
-        DataflowJobLink.persist(self, context, self.project_id, self.region, self.job_id)
+        DataflowJobLink.persist(
+            context=context, project_id=self.project_id, region=self.region, job_id=self.job_id
+        )
 
         if self.deferrable:
             self.defer(
@@ -794,7 +808,7 @@ class DataflowStartYamlJobOperator(GoogleCloudBaseOperator):
             raise AirflowException(event["message"])
         job = event["job"]
         self.log.info("Job %s completed with response %s", job["id"], event["message"])
-        self.xcom_push(context, key="job_id", value=job["id"])
+        context["task_instance"].xcom_push(key="job_id", value=job["id"])
 
         return job
 
@@ -971,6 +985,14 @@ class DataflowCreatePipelineOperator(GoogleCloudBaseOperator):
 
         self.pipeline_name = self.body["name"].split("/")[-1] if self.body else None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+            "location": self.location,
+            "pipeline_name": self.pipeline_name,
+        }
+
     def execute(self, context: Context):
         if self.body is None:
             raise AirflowException(
@@ -1003,8 +1025,8 @@ class DataflowCreatePipelineOperator(GoogleCloudBaseOperator):
                     pipeline_name=self.pipeline_name,
                     location=self.location,
                 )
-        DataflowPipelineLink.persist(self, context, self.project_id, self.location, self.pipeline_name)
-        self.xcom_push(context, key="pipeline_name", value=self.pipeline_name)
+        DataflowPipelineLink.persist(context=context)
+        context["task_instance"].xcom_push(key="pipeline_name", value=self.pipeline_name)
         if self.pipeline:
             if "error" in self.pipeline:
                 raise AirflowException(self.pipeline.get("error").get("message"))
@@ -1075,8 +1097,10 @@ class DataflowRunPipelineOperator(GoogleCloudBaseOperator):
                 location=self.location,
             )["job"]
             job_id = self.dataflow_hook.extract_job_id(self.job)
-            self.xcom_push(context, key="job_id", value=job_id)
-            DataflowJobLink.persist(self, context, self.project_id, self.location, job_id)
+            context["task_instance"].xcom_push(key="job_id", value=job_id)
+            DataflowJobLink.persist(
+                context=context, project_id=self.project_id, region=self.location, job_id=job_id
+            )
         except HttpError as e:
             if e.resp.status == 404:
                 raise AirflowException("Pipeline with given name was not found.")

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

@@ -28,7 +28,7 @@ from airflow.providers.google.cloud.links.dataform import (
 if TYPE_CHECKING:
     from google.api_core.retry import Retry
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dataform_v1beta1.types import (
@@ -258,7 +258,6 @@ class DataformCreateWorkflowInvocationOperator(GoogleCloudBaseOperator):
         )
         workflow_invocation_id = result.name.split("/")[-1]
         DataformWorkflowInvocationLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id,
             region=self.region,
@@ -347,6 +346,13 @@ class DataformGetWorkflowInvocationOperator(GoogleCloudBaseOperator):
             timeout=self.timeout,
             metadata=self.metadata,
         )
+        DataformWorkflowInvocationLink.persist(
+            context=context,
+            project_id=self.project_id,
+            region=self.region,
+            repository_id=self.repository_id,
+            workflow_invocation_id=self.workflow_invocation_id,
+        )
         return WorkflowInvocation.to_dict(result)
 
 
@@ -412,7 +418,6 @@ class DataformQueryWorkflowInvocationActionsOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         DataformWorkflowInvocationLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id,
             region=self.region,
@@ -494,6 +499,13 @@ class DataformCancelWorkflowInvocationOperator(GoogleCloudBaseOperator):
             gcp_conn_id=self.gcp_conn_id,
             impersonation_chain=self.impersonation_chain,
         )
+        DataformWorkflowInvocationLink.persist(
+            context=context,
+            project_id=self.project_id,
+            region=self.region,
+            repository_id=self.repository_id,
+            workflow_invocation_id=self.workflow_invocation_id,
+        )
         hook.cancel_workflow_invocation(
             project_id=self.project_id,
             region=self.region,
@@ -576,7 +588,6 @@ class DataformCreateRepositoryOperator(GoogleCloudBaseOperator):
         )
 
         DataformRepositoryLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id,
             region=self.region,
@@ -735,7 +746,6 @@ class DataformCreateWorkspaceOperator(GoogleCloudBaseOperator):
         )
 
         DataformWorkspaceLink.persist(
-            operator_instance=self,
             context=context,
             project_id=self.project_id,
             region=self.region,

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

@@ -26,7 +26,7 @@ from google.api_core.retry import exponential_sleep_generator
 from googleapiclient.errors import HttpError
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowException
+from airflow.providers.common.compat.sdk import AirflowException
 from airflow.providers.google.cloud.hooks.datafusion import SUCCESS_STATES, DataFusionHook, PipelineStates
 from airflow.providers.google.cloud.links.datafusion import (
     DataFusionInstanceLink,
@@ -40,7 +40,8 @@ from airflow.providers.google.cloud.utils.helpers import resource_path_to_dict
 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 CloudDataFusionRestartInstanceOperator(GoogleCloudBaseOperator):
@@ -111,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,
         )
 
 
@@ -269,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
 
@@ -358,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,
         )
 
 
@@ -429,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
 
@@ -519,7 +516,6 @@ class CloudDataFusionCreatePipelineOperator(GoogleCloudBaseOperator):
         )
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,
@@ -693,7 +689,6 @@ class CloudDataFusionListPipelinesOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelinesLink.persist(
             context=context,
-            task_instance=self,
             uri=service_endpoint,
             namespace=self.namespace,
         )
@@ -783,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
@@ -802,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,
@@ -831,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,
@@ -841,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]):
         """
@@ -870,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):
     """
@@ -943,7 +965,6 @@ class CloudDataFusionStopPipelineOperator(GoogleCloudBaseOperator):
 
         DataFusionPipelineLink.persist(
             context=context,
-            task_instance=self,
             uri=instance["serviceEndpoint"],
             pipeline_name=self.pipeline_name,
             namespace=self.namespace,