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

airflow/providers/google/cloud/openlineage/facets.py

@@ -24,44 +24,145 @@ from attr import define, field
 from airflow.providers.google import __version__ as provider_version
 
 if TYPE_CHECKING:
-    from openlineage.client.generated.base import RunFacet
-else:
+    from openlineage.client.generated.base import JobFacet, RunFacet
+
+try:
+    try:
+        from openlineage.client.generated.base import RunFacet
+    except ImportError:  # Old OpenLineage client is used
+        from openlineage.client.facet import BaseFacet as RunFacet  # type: ignore[assignment]
     try:
-        try:
-            from openlineage.client.generated.base import RunFacet
-        except ImportError:  # Old OpenLineage client is used
-            from openlineage.client.facet import BaseFacet as RunFacet
-
-        @define
-        class BigQueryJobRunFacet(RunFacet):
-            """
-            Facet that represents relevant statistics of bigquery run.
-
-            :param cached: BigQuery caches query results. Rest of the statistics will not be provided for cached queries.
-            :param billedBytes: How many bytes BigQuery bills for.
-            :param properties: Full property tree of BigQUery run.
-            """
-
-            cached: bool
-            billedBytes: int | None = field(default=None)
-            properties: str | None = field(default=None)
-
-            @staticmethod
-            def _get_schema() -> str:
-                return (
-                    "https://raw.githubusercontent.com/apache/airflow/"
-                    f"providers-google/{provider_version}/airflow/providers/google/"
-                    "openlineage/BigQueryJobRunFacet.json"
-                )
-    except ImportError:  # OpenLineage is not available
-
-        def create_no_op(*_, **__) -> None:
-            """
-            Create a no-op placeholder.
-
-            This function creates and returns a None value, used as a placeholder when the OpenLineage client
-            library is available. It represents an action that has no effect.
-            """
-            return None
-
-        BigQueryJobRunFacet = create_no_op
+        from openlineage.client.generated.base import JobFacet
+    except ImportError:  # Old OpenLineage client is used
+        from openlineage.client.facet import BaseFacet as JobFacet  # type: ignore[assignment]
+
+    @define
+    class BigQueryJobRunFacet(RunFacet):
+        """
+        Facet that represents relevant statistics of bigquery run.
+
+        :param cached: BigQuery caches query results. Rest of the statistics will not be provided for cached queries.
+        :param billedBytes: How many bytes BigQuery bills for.
+        :param properties: Full property tree of BigQUery run.
+        """
+
+        cached: bool
+        billedBytes: int | None = field(default=None)
+        properties: str | None = field(default=None)
+
+        @staticmethod
+        def _get_schema() -> str:
+            return (
+                "https://raw.githubusercontent.com/apache/airflow/"
+                f"providers-google/{provider_version}/airflow/providers/google/"
+                "openlineage/BigQueryJobRunFacet.json"
+            )
+
+    @define
+    class CloudStorageTransferJobFacet(JobFacet):
+        """
+        Facet representing a Cloud Storage Transfer Service job configuration.
+
+        :param jobName: Unique name of the transfer job.
+        :param projectId: GCP project where the transfer job is defined.
+        :param description: User-provided description of the transfer job.
+        :param status: Current status of the transfer job (e.g. "ENABLED", "DISABLED").
+        :param sourceBucket: Name of the source bucket (e.g. AWS S3).
+        :param sourcePath: Prefix/path inside the source bucket.
+        :param targetBucket: Name of the destination bucket (e.g. GCS).
+        :param targetPath: Prefix/path inside the destination bucket.
+        :param objectConditions: Object selection rules (e.g. include/exclude prefixes).
+        :param transferOptions: Transfer options, such as overwrite behavior or whether to delete objects
+        from the source after transfer.
+        :param schedule: Schedule for the transfer job (if recurring).
+        """
+
+        jobName: str | None = field(default=None)
+        projectId: str | None = field(default=None)
+        description: str | None = field(default=None)
+        status: str | None = field(default=None)
+        sourceBucket: str | None = field(default=None)
+        sourcePath: str | None = field(default=None)
+        targetBucket: str | None = field(default=None)
+        targetPath: str | None = field(default=None)
+        objectConditions: dict | None = field(default=None)
+        transferOptions: dict | None = field(default=None)
+        schedule: dict | None = field(default=None)
+
+        @staticmethod
+        def _get_schema() -> str:
+            return (
+                "https://raw.githubusercontent.com/apache/airflow/"
+                f"providers-google/{provider_version}/airflow/providers/google/"
+                "openlineage/CloudStorageTransferJobFacet.json"
+            )
+
+    @define
+    class CloudStorageTransferRunFacet(RunFacet):
+        """
+        Facet representing a Cloud Storage Transfer Service job execution run.
+
+        :param jobName: Name of the transfer job being executed.
+        :param operationName: Name of the specific transfer operation instance.
+        :param status: Current status of the operation (e.g. "IN_PROGRESS", "SUCCESS", "FAILED").
+        :param startTime: Time when the transfer job execution started (ISO 8601 format).
+        :param endTime: Time when the transfer job execution finished (ISO 8601 format).
+        :param wait: Whether the operator waits for the job to complete before finishing.
+        :param timeout: Timeout (in seconds) for the transfer run to complete.
+        :param deferrable: Whether the operator defers execution until job completion.
+        :param deleteJobAfterCompletion: Whether the operator deletes the transfer job after the run completes.
+        """
+
+        jobName: str | None = field(default=None)
+        operationName: str | None = field(default=None)
+        status: str | None = field(default=None)
+        startTime: str | None = field(default=None)
+        endTime: str | None = field(default=None)
+        wait: bool = field(default=True)
+        timeout: float | None = field(default=None)
+        deferrable: bool = field(default=False)
+        deleteJobAfterCompletion: bool = field(default=False)
+
+        @staticmethod
+        def _get_schema() -> str:
+            return (
+                "https://raw.githubusercontent.com/apache/airflow/"
+                f"providers-google/{provider_version}/airflow/providers/google/"
+                "openlineage/CloudStorageTransferRunFacet.json"
+            )
+
+    @define
+    class DataFusionRunFacet(RunFacet):
+        """
+        Facet that represents relevant details of a Cloud Data Fusion pipeline run.
+
+        :param runId: The pipeline execution id.
+        :param runtimeArgs: Runtime arguments passed to the pipeline.
+        """
+
+        runId: str | None = field(default=None)
+        runtimeArgs: dict[str, str] | None = field(default=None)
+
+        @staticmethod
+        def _get_schema() -> str:
+            return (
+                "https://raw.githubusercontent.com/apache/airflow/"
+                f"providers-google/{provider_version}/airflow/providers/google/"
+                "openlineage/DataFusionRunFacet.json"
+            )
+
+except ImportError:  # OpenLineage is not available
+
+    def create_no_op(*_, **__) -> None:
+        """
+        Create a no-op placeholder.
+
+        This function creates and returns a None value, used as a placeholder when the OpenLineage client
+        library is available. It represents an action that has no effect.
+        """
+        return None
+
+    BigQueryJobRunFacet = create_no_op  # type: ignore[misc, assignment]
+    CloudStorageTransferJobFacet = create_no_op  # type: ignore[misc, assignment]
+    CloudStorageTransferRunFacet = create_no_op  # type: ignore[misc, assignment]
+    DataFusionRunFacet = create_no_op  # type: ignore[misc, assignment]

airflow/providers/google/cloud/openlineage/mixins.py

@@ -80,7 +80,7 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
         from airflow.providers.openlineage.sqlparser import SQLParser
 
         if not self.job_id:
-            self.log.warning("No BigQuery job_id was found by OpenLineage.")  # type: ignore[attr-defined]
+            self.log.warning("No BigQuery job_id was found by OpenLineage.")
             return OperatorLineage()
 
         if not self.hook:
@@ -92,14 +92,16 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
                 impersonation_chain=self.impersonation_chain,
             )
 
-        self.log.debug("Extracting data from bigquery job: `%s`", self.job_id)  # type: ignore[attr-defined]
+        self.log.debug("Extracting data from bigquery job: `%s`", self.job_id)
         inputs, outputs = [], []
         run_facets: dict[str, RunFacet] = {
             "externalQuery": ExternalQueryRunFacet(externalQueryId=self.job_id, source="bigquery")
         }
-        self._client = self.hook.get_client(project_id=self.hook.project_id, location=self.location)
+        self._client = self.hook.get_client(
+            project_id=self.project_id or self.hook.project_id, location=self.location
+        )
         try:
-            job_properties = self._client.get_job(job_id=self.job_id)._properties  # type: ignore
+            job_properties = self._client.get_job(job_id=self.job_id)._properties
 
             if get_from_nullable_chain(job_properties, ["status", "state"]) != "DONE":
                 raise ValueError(f"Trying to extract data from running bigquery job: `{self.job_id}`")
@@ -107,11 +109,11 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
             run_facets["bigQueryJob"] = self._get_bigquery_job_run_facet(job_properties)
 
             if get_from_nullable_chain(job_properties, ["statistics", "numChildJobs"]):
-                self.log.debug("Found SCRIPT job. Extracting lineage from child jobs instead.")  # type: ignore[attr-defined]
+                self.log.debug("Found SCRIPT job. Extracting lineage from child jobs instead.")
                 # SCRIPT job type has no input / output information but spawns child jobs that have one
                 # https://cloud.google.com/bigquery/docs/information-schema-jobs#multi-statement_query_job
                 for child_job_id in self._client.list_jobs(parent_job=self.job_id):
-                    child_job_properties = self._client.get_job(job_id=child_job_id)._properties  # type: ignore
+                    child_job_properties = self._client.get_job(job_id=child_job_id)._properties
                     child_inputs, child_outputs = self._get_inputs_and_outputs(child_job_properties)
                     inputs.extend(child_inputs)
                     outputs.extend(child_outputs)
@@ -119,7 +121,7 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
                 inputs, outputs = self._get_inputs_and_outputs(job_properties)
 
         except Exception as e:
-            self.log.warning("Cannot retrieve job details from BigQuery.Client. %s", e, exc_info=True)  # type: ignore[attr-defined]
+            self.log.warning("Cannot retrieve job details from BigQuery.Client. %s", e, exc_info=True)
             exception_msg = traceback.format_exc()
             run_facets.update(
                 {
@@ -173,7 +175,7 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
             if (
                 single_output.facets
                 and final_outputs[key].facets
-                and "columnLineage" in single_output.facets  # type: ignore
+                and "columnLineage" in single_output.facets
                 and "columnLineage" in final_outputs[key].facets  # type: ignore
             ):
                 single_output.facets["columnLineage"] = merge_column_lineage_facets(
@@ -188,10 +190,10 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
         return list(final_outputs.values())
 
     def _get_input_dataset(self, table: dict) -> InputDataset:
-        return cast(InputDataset, self._get_dataset(table, "input"))
+        return cast("InputDataset", self._get_dataset(table, "input"))
 
     def _get_output_dataset(self, table: dict) -> OutputDataset:
-        return cast(OutputDataset, self._get_dataset(table, "output"))
+        return cast("OutputDataset", self._get_dataset(table, "output"))
 
     def _get_dataset(self, table: dict, dataset_type: str) -> Dataset:
         project = table.get("projectId")
@@ -207,15 +209,14 @@ class _BigQueryInsertJobOperatorOpenLineageMixin:
                 name=dataset_name,
                 facets=dataset_facets,
             )
-        elif dataset_type == "output":
+        if dataset_type == "output":
             # Logic specific to creating OutputDataset (if needed)
             return OutputDataset(
                 namespace=BIGQUERY_NAMESPACE,
                 name=dataset_name,
                 facets=dataset_facets,
             )
-        else:
-            raise ValueError("Invalid dataset_type. Must be 'input' or 'output'")
+        raise ValueError("Invalid dataset_type. Must be 'input' or 'output'")
 
     def _get_table_facets_safely(self, table_name: str) -> dict[str, DatasetFacet]:
         try:

airflow/providers/google/cloud/openlineage/utils.py

@@ -26,6 +26,8 @@ from collections import defaultdict
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any
 
+from google.cloud.dataproc_v1 import Batch, RuntimeConfig
+
 from airflow.providers.common.compat.openlineage.facet import (
     ColumnLineageDatasetFacet,
     DatasetFacet,
@@ -42,13 +44,13 @@ from airflow.providers.common.compat.openlineage.utils.spark import (
     inject_transport_information_into_spark_properties,
 )
 from airflow.providers.google.cloud.hooks.gcs import _parse_gcs_url
-from google.cloud.dataproc_v1 import Batch, RuntimeConfig
 
 if TYPE_CHECKING:
-    from airflow.providers.common.compat.openlineage.facet import Dataset
-    from airflow.utils.context import Context
     from google.cloud.bigquery.table import Table
 
+    from airflow.providers.common.compat.openlineage.facet import Dataset
+    from airflow.providers.common.compat.sdk import Context
+
 
 log = logging.getLogger(__name__)
 
@@ -212,7 +214,20 @@ def extract_ds_name_from_gcs_path(path: str) -> str:
 
 def get_facets_from_bq_table(table: Table) -> dict[str, DatasetFacet]:
     """Get facets from BigQuery table object."""
+    return get_facets_from_bq_table_for_given_fields(table, selected_fields=None)
+
+
+def get_facets_from_bq_table_for_given_fields(
+    table: Table, selected_fields: list[str] | None
+) -> dict[str, DatasetFacet]:
+    """
+    Get facets from BigQuery table object for selected fields only.
+
+    If selected_fields is None, include all fields.
+    """
     facets: dict[str, DatasetFacet] = {}
+    selected_fields_set = set(selected_fields) if selected_fields else None
+
     if table.schema:
         facets["schema"] = SchemaDatasetFacet(
             fields=[
@@ -220,6 +235,7 @@ def get_facets_from_bq_table(table: Table) -> dict[str, DatasetFacet]:
                     name=schema_field.name, type=schema_field.field_type, description=schema_field.description
                 )
                 for schema_field in table.schema
+                if selected_fields_set is None or schema_field.name in selected_fields_set
             ]
         )
     if table.description:

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

@@ -21,7 +21,11 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from functools import cached_property
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
+
+from google.api_core.exceptions import NotFound
+from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
+from google.cloud import alloydb_v1
 
 from airflow.exceptions import AirflowException
 from airflow.providers.google.cloud.hooks.alloy_db import AlloyDbHook
@@ -31,18 +35,15 @@ from airflow.providers.google.cloud.links.alloy_db import (
     AlloyDBUsersLink,
 )
 from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
-from google.api_core.exceptions import NotFound
-from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
-from google.cloud import alloydb_v1
 
 if TYPE_CHECKING:
     import proto
-
-    from airflow.utils.context import Context
     from google.api_core.operation import Operation
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
+    from airflow.providers.common.compat.sdk import Context
+
 
 class AlloyDBBaseOperator(GoogleCloudBaseOperator):
     """
@@ -144,8 +145,7 @@ class AlloyDBWriteBaseOperator(AlloyDBBaseOperator):
         if self.validate_request:
             # Validation requests are only validated and aren't executed, thus no operation result is expected
             return None
-        else:
-            return self.hook.wait_for_operation(timeout=self.timeout, operation=operation)
+        return self.hook.wait_for_operation(timeout=self.timeout, operation=operation)
 
 
 class AlloyDBCreateClusterOperator(AlloyDBWriteBaseOperator):
@@ -228,15 +228,16 @@ class AlloyDBCreateClusterOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
-    def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
 
+    def execute(self, context: Context) -> dict | None:
+        AlloyDBClusterLink.persist(context=context)
         if cluster := self._get_cluster():
             return cluster
 
@@ -334,14 +335,16 @@ class AlloyDBUpdateClusterOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB cluster request.")
         else:
@@ -545,14 +548,16 @@ class AlloyDBCreateInstanceOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if instance := self._get_instance():
             return instance
 
@@ -654,14 +659,16 @@ class AlloyDBUpdateInstanceOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBClusterLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBClusterLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB instance request.")
         else:
@@ -861,14 +868,16 @@ class AlloyDBCreateUserOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBUsersLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBUsersLink.persist(context=context)
         if (_user := self._get_user()) is not None:
             return _user
 
@@ -968,14 +977,16 @@ class AlloyDBUpdateUserOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "location_id": self.location,
+            "cluster_id": self.cluster_id,
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBUsersLink.persist(
-            context=context,
-            task_instance=self,
-            location_id=self.location,
-            cluster_id=self.cluster_id,
-            project_id=self.project_id,
-        )
+        AlloyDBUsersLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB user request.")
         else:
@@ -1159,12 +1170,14 @@ class AlloyDBCreateBackupOperator(AlloyDBWriteBaseOperator):
             return result
         return None
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBBackupsLink.persist(
-            context=context,
-            task_instance=self,
-            project_id=self.project_id,
-        )
+        AlloyDBBackupsLink.persist(context=context)
         if backup := self._get_backup():
             return backup
 
@@ -1259,12 +1272,14 @@ class AlloyDBUpdateBackupOperator(AlloyDBWriteBaseOperator):
         self.update_mask = update_mask
         self.allow_missing = allow_missing
 
+    @property
+    def extra_links_params(self) -> dict[str, Any]:
+        return {
+            "project_id": self.project_id,
+        }
+
     def execute(self, context: Context) -> dict | None:
-        AlloyDBBackupsLink.persist(
-            context=context,
-            task_instance=self,
-            project_id=self.project_id,
-        )
+        AlloyDBBackupsLink.persist(context=context)
         if self.validate_request:
             self.log.info("Validating an Update AlloyDB backup request.")
         else: