apache-airflow-providers-google 11.0.0rc1__py3-none-any.whl → 12.0.0__py3-none-any.whl

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

@@ -1365,7 +1365,7 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
 
         try:
             self.log.info("Creating table")
-            table = bq_hook.create_empty_table(
+            self._table = bq_hook.create_empty_table(
                 project_id=self.project_id,
                 dataset_id=self.dataset_id,
                 table_id=self.table_id,
@@ -1382,12 +1382,15 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
             persist_kwargs = {
                 "context": context,
                 "task_instance": self,
-                "project_id": table.to_api_repr()["tableReference"]["projectId"],
-                "dataset_id": table.to_api_repr()["tableReference"]["datasetId"],
-                "table_id": table.to_api_repr()["tableReference"]["tableId"],
+                "project_id": self._table.to_api_repr()["tableReference"]["projectId"],
+                "dataset_id": self._table.to_api_repr()["tableReference"]["datasetId"],
+                "table_id": self._table.to_api_repr()["tableReference"]["tableId"],
             }
             self.log.info(
-                "Table %s.%s.%s created successfully", table.project, table.dataset_id, table.table_id
+                "Table %s.%s.%s created successfully",
+                self._table.project,
+                self._table.dataset_id,
+                self._table.table_id,
             )
         except Conflict:
             error_msg = f"Table {self.dataset_id}.{self.table_id} already exists."
@@ -1407,6 +1410,24 @@ class BigQueryCreateEmptyTableOperator(GoogleCloudBaseOperator):
 
         BigQueryTableLink.persist(**persist_kwargs)
 
+    def get_openlineage_facets_on_complete(self, task_instance):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            get_facets_from_bq_table,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        table_info = self._table.to_api_repr()["tableReference"]
+        table_id = ".".join((table_info["projectId"], table_info["datasetId"], table_info["tableId"]))
+        output_dataset = Dataset(
+            namespace=BIGQUERY_NAMESPACE,
+            name=table_id,
+            facets=get_facets_from_bq_table(self._table),
+        )
+
+        return OperatorLineage(outputs=[output_dataset])
+
 
 class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
     """
@@ -1632,15 +1653,15 @@ class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
             impersonation_chain=self.impersonation_chain,
         )
         if self.table_resource:
-            table = bq_hook.create_empty_table(
+            self._table = bq_hook.create_empty_table(
                 table_resource=self.table_resource,
             )
             BigQueryTableLink.persist(
                 context=context,
                 task_instance=self,
-                dataset_id=table.to_api_repr()["tableReference"]["datasetId"],
-                project_id=table.to_api_repr()["tableReference"]["projectId"],
-                table_id=table.to_api_repr()["tableReference"]["tableId"],
+                dataset_id=self._table.to_api_repr()["tableReference"]["datasetId"],
+                project_id=self._table.to_api_repr()["tableReference"]["projectId"],
+                table_id=self._table.to_api_repr()["tableReference"]["tableId"],
             )
             return
 
@@ -1691,18 +1712,36 @@ class BigQueryCreateExternalTableOperator(GoogleCloudBaseOperator):
             "encryptionConfiguration": self.encryption_configuration,
         }
 
-        table = bq_hook.create_empty_table(
+        self._table = bq_hook.create_empty_table(
             table_resource=table_resource,
         )
 
         BigQueryTableLink.persist(
             context=context,
             task_instance=self,
-            dataset_id=table.to_api_repr()["tableReference"]["datasetId"],
-            project_id=table.to_api_repr()["tableReference"]["projectId"],
-            table_id=table.to_api_repr()["tableReference"]["tableId"],
+            dataset_id=self._table.to_api_repr()["tableReference"]["datasetId"],
+            project_id=self._table.to_api_repr()["tableReference"]["projectId"],
+            table_id=self._table.to_api_repr()["tableReference"]["tableId"],
+        )
+
+    def get_openlineage_facets_on_complete(self, task_instance):
+        from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import (
+            BIGQUERY_NAMESPACE,
+            get_facets_from_bq_table,
+        )
+        from airflow.providers.openlineage.extractors import OperatorLineage
+
+        table_info = self._table.to_api_repr()["tableReference"]
+        table_id = ".".join((table_info["projectId"], table_info["datasetId"], table_info["tableId"]))
+        output_dataset = Dataset(
+            namespace=BIGQUERY_NAMESPACE,
+            name=table_id,
+            facets=get_facets_from_bq_table(self._table),
         )
 
+        return OperatorLineage(outputs=[output_dataset])
+
 
 class BigQueryDeleteDatasetOperator(GoogleCloudBaseOperator):
     """
@@ -2593,10 +2632,16 @@ class BigQueryInsertJobOperator(GoogleCloudBaseOperator, _BigQueryOpenLineageMix
             nowait=True,
         )
 
-    @staticmethod
-    def _handle_job_error(job: BigQueryJob | UnknownJob) -> None:
+    def _handle_job_error(self, job: BigQueryJob | UnknownJob) -> None:
+        self.log.info("Job %s is completed. Checking the job status", self.job_id)
+        # Log any transient errors encountered during the job execution
+        for error in job.errors or []:
+            self.log.error("BigQuery Job Error: %s", error)
         if job.error_result:
             raise AirflowException(f"BigQuery job {job.job_id} failed: {job.error_result}")
+        # Check the final state.
+        if job.state != "DONE":
+            raise AirflowException(f"Job failed with state: {job.state}")
 
     def execute(self, context: Any):
         hook = BigQueryHook(

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

@@ -54,6 +54,9 @@ 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,
@@ -1962,6 +1965,9 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         polling_interval_seconds: int = 10,
         cancel_on_kill: bool = True,
         wait_timeout: int | None = None,
+        openlineage_inject_parent_job_info: bool = conf.getboolean(
+            "openlineage", "spark_inject_parent_job_info", fallback=False
+        ),
         **kwargs,
     ) -> None:
         super().__init__(**kwargs)
@@ -1983,10 +1989,16 @@ class DataprocSubmitJobOperator(GoogleCloudBaseOperator):
         self.hook: DataprocHook | None = None
         self.job_id: str | None = None
         self.wait_timeout = wait_timeout
+        self.openlineage_inject_parent_job_info = openlineage_inject_parent_job_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:
+            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
+            )
         job_object = self.hook.submit_job(
             project_id=self.project_id,
             region=self.region,

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

@@ -343,6 +343,7 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
             LifecycleStateChangeDatasetFacet,
             PreviousIdentifier,
         )
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
         from airflow.providers.openlineage.extractors import OperatorLineage
 
         objects = []
@@ -350,12 +351,7 @@ class GCSDeleteObjectsOperator(GoogleCloudBaseOperator):
             objects = self.objects
         elif self.prefix is not None:
             prefixes = [self.prefix] if isinstance(self.prefix, str) else self.prefix
-            for pref in prefixes:
-                # Use parent if not a file (dot not in name) and not a dir (ends with slash)
-                if "." not in pref.split("/")[-1] and not pref.endswith("/"):
-                    pref = Path(pref).parent.as_posix()
-                pref = "/" if pref in (".", "", "/") else pref.rstrip("/")
-                objects.append(pref)
+            objects = [extract_ds_name_from_gcs_path(pref) for pref in prefixes]
 
         bucket_url = f"gs://{self.bucket_name}"
         input_datasets = [
@@ -921,20 +917,15 @@ class GCSTimeSpanFileTransformOperator(GoogleCloudBaseOperator):
     def get_openlineage_facets_on_complete(self, task_instance):
         """Implement on_complete as execute() resolves object prefixes."""
         from airflow.providers.common.compat.openlineage.facet import Dataset
+        from airflow.providers.google.cloud.openlineage.utils import extract_ds_name_from_gcs_path
         from airflow.providers.openlineage.extractors import OperatorLineage
 
-        def _parse_prefix(pref):
-            # Use parent if not a file (dot not in name) and not a dir (ends with slash)
-            if "." not in pref.split("/")[-1] and not pref.endswith("/"):
-                pref = Path(pref).parent.as_posix()
-            return "/" if pref in (".", "/", "") else pref.rstrip("/")
-
         input_prefix, output_prefix = "/", "/"
         if self._source_prefix_interp is not None:
-            input_prefix = _parse_prefix(self._source_prefix_interp)
+            input_prefix = extract_ds_name_from_gcs_path(self._source_prefix_interp)
 
         if self._destination_prefix_interp is not None:
-            output_prefix = _parse_prefix(self._destination_prefix_interp)
+            output_prefix = extract_ds_name_from_gcs_path(self._destination_prefix_interp)
 
         return OperatorLineage(
             inputs=[