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

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

@@ -30,11 +30,6 @@ from enum import Enum
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from google.api_core.exceptions import AlreadyExists, NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.api_core.retry import Retry, exponential_sleep_generator
-from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
-
 from airflow.configuration import conf
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.dataproc import (
@@ -54,9 +49,6 @@ from airflow.providers.google.cloud.links.dataproc import (
     DataprocWorkflowLink,
     DataprocWorkflowTemplateLink,
 )
-from airflow.providers.google.cloud.openlineage.utils import (
-    inject_openlineage_properties_into_dataproc_job,
-)
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.cloud.triggers.dataproc import (
     DataprocBatchTrigger,
@@ -69,16 +61,19 @@ from airflow.providers.google.cloud.utils.dataproc import DataprocOperationType
 from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.utils import timezone
+from google.api_core.exceptions import AlreadyExists, NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.api_core.retry import Retry, exponential_sleep_generator
+from google.cloud.dataproc_v1 import Batch, Cluster, ClusterStatus, JobStatus
 
 if TYPE_CHECKING:
+    from airflow.utils.context import Context
     from google.api_core import operation
     from google.api_core.retry_async import AsyncRetry
     from google.protobuf.duration_pb2 import Duration
     from google.protobuf.field_mask_pb2 import FieldMask
     from google.type.interval_pb2 import Interval
 
-    from airflow.utils.context import Context
-
 
 class PreemptibilityType(Enum):
     """Contains possible Type values of Preemptibility applicable for every secondary worker of Cluster."""
@@ -1824,6 +1819,12 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         polling_interval_seconds: int = 10,
         cancel_on_kill: bool = True,
+        openlineage_inject_parent_job_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_parent_job_info", fallback=False
+        ),
+        openlineage_inject_transport_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_transport_info", fallback=False
+        ),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1843,11 +1844,17 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
         self.polling_interval_seconds = polling_interval_seconds
         self.cancel_on_kill = cancel_on_kill
         self.operation_name: str | None = None
+        self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info
+        self.openlineage_inject_transport_info = openlineage_inject_transport_info
 
     def execute(self, context: Context):
         self.log.info("Instantiating Inline Template")
         hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
         project_id = self.project_id or hook.project_id
+        if self.openlineage_inject_parent_job_info or self.openlineage_inject_transport_info:
+            self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.")
+            self._inject_openlineage_properties_into_dataproc_workflow_template(context)
+
         operation = hook.instantiate_inline_workflow_template(
             template=self.template,
             project_id=project_id,
@@ -1903,6 +1910,25 @@ class DataprocInstantiateInlineWorkflowTemplateOperator(GoogleCloudBaseOperator)
             hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
             hook.get_operations_client(region=self.region).cancel_operation(name=self.operation_name)
 
+    def _inject_openlineage_properties_into_dataproc_workflow_template(self, context: Context) -> None:
+        try:
+            from airflow.providers.google.cloud.openlineage.utils import (
+                inject_openlineage_properties_into_dataproc_workflow_template,
+            )
+
+            self.template = inject_openlineage_properties_into_dataproc_workflow_template(
+                template=self.template,
+                context=context,
+                inject_parent_job_info=self.openlineage_inject_parent_job_info,
+                inject_transport_info=self.openlineage_inject_transport_info,
+            )
+        except Exception as e:
+            self.log.warning(
+                "An error occurred while trying to inject OpenLineage information. "
+                "Dataproc template has not been modified by OpenLineage.",
+                exc_info=e,
+            )
+
 
 class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
     """
@@ -1968,6 +1994,9 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         openlineage_inject_parent_job_info: bool = conf.getboolean(
             "openlineage", "spark_inject_parent_job_info", fallback=False
         ),
+        openlineage_inject_transport_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_transport_info", fallback=False
+        ),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1990,15 +2019,15 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         self.job_id: str | None = None
         self.wait_timeout = wait_timeout
         self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info
+        self.openlineage_inject_transport_info = openlineage_inject_transport_info
 
     def execute(self, context: Context):
         self.log.info("Submitting job")
         self.hook = DataprocHook(gcp_conn_id=self.gcp_conn_id, impersonation_chain=self.impersonation_chain)
-        if self.openlineage_inject_parent_job_info:
+        if self.openlineage_inject_parent_job_info or self.openlineage_inject_transport_info:
             self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.")
-            self.job = inject_openlineage_properties_into_dataproc_job(
-                job=self.job, context=context, inject_parent_job_info=self.openlineage_inject_parent_job_info
-            )
+            self._inject_openlineage_properties_into_dataproc_job(context)
+
         job_object = self.hook.submit_job(
             project_id=self.project_id,
             region=self.region,
@@ -2072,6 +2101,25 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         if self.job_id and self.cancel_on_kill:
             self.hook.cancel_job(job_id=self.job_id, project_id=self.project_id, region=self.region)
 
+    def _inject_openlineage_properties_into_dataproc_job(self, context: Context) -> None:
+        try:
+            from airflow.providers.google.cloud.openlineage.utils import (
+                inject_openlineage_properties_into_dataproc_job,
+            )
+
+            self.job = inject_openlineage_properties_into_dataproc_job(
+                job=self.job,
+                context=context,
+                inject_parent_job_info=self.openlineage_inject_parent_job_info,
+                inject_transport_info=self.openlineage_inject_transport_info,
+            )
+        except Exception as e:
+            self.log.warning(
+                "An error occurred while trying to inject OpenLineage information. "
+                "Dataproc job has not been modified by OpenLineage.",
+                exc_info=e,
+            )
+
 
 class DataprocUpdateClusterOperator(GoogleCloudBaseOperator):
     """
@@ -2425,6 +2473,12 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         asynchronous: bool = False,
         deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False),
         polling_interval_seconds: int = 5,
+        openlineage_inject_parent_job_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_parent_job_info", fallback=False
+        ),
+        openlineage_inject_transport_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_transport_info", fallback=False
+        ),
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -2446,6 +2500,8 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         self.asynchronous = asynchronous
         self.deferrable = deferrable
         self.polling_interval_seconds = polling_interval_seconds
+        self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_info
+        self.openlineage_inject_transport_info = openlineage_inject_transport_info
 
     def execute(self, context: Context):
         if self.asynchronous and self.deferrable:
@@ -2468,6 +2524,10 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         else:
             self.log.info("Starting batch. The batch ID will be generated since it was not provided.")
 
+        if self.openlineage_inject_parent_job_info or self.openlineage_inject_transport_info:
+            self.log.info("Automatic injection of OpenLineage information into Spark properties is enabled.")
+            self._inject_openlineage_properties_into_dataproc_batch(context)
+
         try:
             self.operation = self.hook.create_batch(
                 region=self.region,
@@ -2483,7 +2543,10 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             self.log.info("Batch with given id already exists.")
             self.log.info("Attaching to the job %s if it is still running.", batch_id)
         else:
-            batch_id = self.operation.metadata.batch.split("/")[-1]
+            if self.operation and self.operation.metadata:
+                batch_id = self.operation.metadata.batch.split("/")[-1]
+            else:
+                raise AirflowException("Operation metadata is not available.")
             self.log.info("The batch %s was created.", batch_id)
 
         DataprocBatchLink.persist(
@@ -2609,8 +2672,11 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
             self.log.info("Batch with given id already exists.")
             self.log.info("Attaching to the job %s if it is still running.", self.batch_id)
         else:
-            batch_id = self.operation.metadata.batch.split("/")[-1]
-            self.log.info("The batch %s was created.", batch_id)
+            if self.operation and self.operation.metadata:
+                batch_id = self.operation.metadata.batch.split("/")[-1]
+                self.log.info("The batch %s was created.", batch_id)
+            else:
+                raise AirflowException("Operation metadata is not available.")
 
         self.log.info("Waiting for the completion of batch job %s", batch_id)
         batch = self.hook.wait_for_batch(
@@ -2623,6 +2689,25 @@ class DataprocCreateBatchOperator(GoogleCloudBaseOperator):
         )
         return batch, batch_id
 
+    def _inject_openlineage_properties_into_dataproc_batch(self, context: Context) -> None:
+        try:
+            from airflow.providers.google.cloud.openlineage.utils import (
+                inject_openlineage_properties_into_dataproc_batch,
+            )
+
+            self.batch = inject_openlineage_properties_into_dataproc_batch(
+                batch=self.batch,
+                context=context,
+                inject_parent_job_info=self.openlineage_inject_parent_job_info,
+                inject_transport_info=self.openlineage_inject_transport_info,
+            )
+        except Exception as e:
+            self.log.warning(
+                "An error occurred while trying to inject OpenLineage information. "
+                "Dataproc batch has not been modified by OpenLineage.",
+                exc_info=e,
+            )
+
 
 class DataprocDeleteBatchOperator(GoogleCloudBaseOperator):
     """

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

@@ -23,25 +23,23 @@ import time
 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.api_core.retry import Retry, exponential_sleep_generator
-from google.cloud.metastore_v1 import MetadataExport, MetadataManagementActivity
-from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
-from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
-
 from airflow.exceptions import AirflowException
 from airflow.models import BaseOperator, BaseOperatorLink
 from airflow.models.xcom import XCom
 from airflow.providers.google.cloud.hooks.dataproc_metastore import DataprocMetastoreHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.links.storage import StorageLink
+from google.api_core.exceptions import AlreadyExists
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.api_core.retry import Retry, exponential_sleep_generator
+from google.cloud.metastore_v1 import MetadataExport, MetadataManagementActivity
+from google.cloud.metastore_v1.types import Backup, MetadataImport, Service
+from google.cloud.metastore_v1.types.metastore import DatabaseDumpSpec, Restore
 
 if TYPE_CHECKING:
-    from google.protobuf.field_mask_pb2 import FieldMask
-
     from airflow.models.taskinstancekey import TaskInstanceKey
     from airflow.utils.context import Context
+    from google.protobuf.field_mask_pb2 import FieldMask
 
 
 BASE_LINK = "https://console.cloud.google.com"

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

@@ -22,6 +22,22 @@ from __future__ import annotations
 from collections.abc import Sequence
 from typing import TYPE_CHECKING
 
+from airflow.providers.google.cloud.hooks.dlp import CloudDLPHook
+from airflow.providers.google.cloud.links.data_loss_prevention import (
+    CloudDLPDeidentifyTemplateDetailsLink,
+    CloudDLPDeidentifyTemplatesListLink,
+    CloudDLPInfoTypeDetailsLink,
+    CloudDLPInfoTypesListLink,
+    CloudDLPInspectTemplateDetailsLink,
+    CloudDLPInspectTemplatesListLink,
+    CloudDLPJobDetailsLink,
+    CloudDLPJobsListLink,
+    CloudDLPJobTriggerDetailsLink,
+    CloudDLPJobTriggersListLink,
+    CloudDLPPossibleInfoTypesListLink,
+)
+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, InvalidArgument, NotFound
 from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
 from google.cloud.dlp_v2.types import (
@@ -45,29 +61,11 @@ from google.cloud.dlp_v2.types import (
     StoredInfoTypeConfig,
 )
 
-from airflow.providers.google.cloud.hooks.dlp import CloudDLPHook
-from airflow.providers.google.cloud.links.data_loss_prevention import (
-    CloudDLPDeidentifyTemplateDetailsLink,
-    CloudDLPDeidentifyTemplatesListLink,
-    CloudDLPInfoTypeDetailsLink,
-    CloudDLPInfoTypesListLink,
-    CloudDLPInspectTemplateDetailsLink,
-    CloudDLPInspectTemplatesListLink,
-    CloudDLPJobDetailsLink,
-    CloudDLPJobsListLink,
-    CloudDLPJobTriggerDetailsLink,
-    CloudDLPJobTriggersListLink,
-    CloudDLPPossibleInfoTypesListLink,
-)
-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 CloudDLPCancelDLPJobOperator(GoogleCloudBaseOperator):
     """

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

@@ -35,14 +35,13 @@ 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.exceptions import Conflict
-from google.cloud.exceptions import GoogleCloudError
-
 from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.hooks.gcs import GCSHook
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
 from airflow.providers.google.common.links.storage import FileDetailsLink, StorageLink
 from airflow.utils import timezone
+from google.api_core.exceptions import Conflict
+from google.cloud.exceptions import GoogleCloudError
 
 
 class GCSCreateBucketOperator(GoogleCloudBaseOperator):
@@ -787,22 +786,20 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
 
     def execute(self, context: Context) -> list[str]:
         # Define intervals and prefixes.
-        try:
-            orig_start = context["data_interval_start"]
-            orig_end = context["data_interval_end"]
-        except KeyError:
-            orig_start = pendulum.instance(context["logical_date"])
-            next_dagrun = context["dag"].next_dagrun_info(last_automated_dagrun=None, restricted=False)
-            if next_dagrun and next_dagrun.data_interval and next_dagrun.data_interval.end:
-                orig_end = next_dagrun.data_interval.end
-            else:
-                orig_end = None
+        orig_start = context["data_interval_start"]
+        orig_end = context["data_interval_end"]
+
+        if orig_start is None or orig_end is None:
+            raise RuntimeError("`data_interval_start` & `data_interval_end` must not be None")
+
+        if not isinstance(orig_start, pendulum.DateTime):
+            orig_start = pendulum.instance(orig_start)
+
+        if not isinstance(orig_end, pendulum.DateTime):
+            orig_end = pendulum.instance(orig_end)
 
         timespan_start = orig_start
-        if orig_end is None:  # Only possible in Airflow before 2.2.
-            self.log.warning("No following schedule found, setting timespan end to max %s", orig_end)
-            timespan_end = pendulum.instance(datetime.datetime.max)
-        elif orig_start >= orig_end:  # Airflow 2.2 sets start == end for non-perodic schedules.
+        if orig_start >= orig_end:  # Airflow 2.2 sets start == end for non-perodic schedules.
             self.log.warning("DAG schedule not periodic, setting timespan end to max %s", orig_end)
             timespan_end = pendulum.instance(datetime.datetime.max)
         else:

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

@@ -24,7 +24,6 @@ from collections.abc import Sequence
 from functools import cached_property
 from typing import TYPE_CHECKING, Any
 
-from google.api_core.exceptions import AlreadyExists
 from kubernetes.client import V1JobList, models as k8s
 from packaging.version import parse as parse_version
 
@@ -61,6 +60,7 @@ from airflow.providers.google.common.deprecated import deprecated
 from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
 from airflow.providers_manager import ProvidersManager
 from airflow.utils.timezone import utcnow
+from google.api_core.exceptions import AlreadyExists
 
 try:
     from airflow.providers.cncf.kubernetes.operators.job import KubernetesDeleteJobOperator
@@ -73,11 +73,11 @@ except ImportError:
     )
 
 if TYPE_CHECKING:
-    from google.cloud.container_v1.types import Cluster
     from kubernetes.client.models import V1Job
     from pendulum import DateTime
 
     from airflow.utils.context import Context
+    from google.cloud.container_v1.types import Cluster
 
 KUBE_CONFIG_ENV_VAR = "KUBECONFIG"