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

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

@@ -24,13 +24,17 @@ from attr import define, field
 from airflow.providers.google import __version__ as provider_version
 
 if TYPE_CHECKING:
-    from openlineage.client.generated.base import RunFacet
+    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:
+        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):
@@ -53,6 +57,100 @@ try:
                 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:
@@ -65,3 +163,6 @@ except ImportError:  # OpenLineage is not available
         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(

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

@@ -49,7 +49,7 @@ if TYPE_CHECKING:
     from google.cloud.bigquery.table import Table
 
     from airflow.providers.common.compat.openlineage.facet import Dataset
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 log = logging.getLogger(__name__)
@@ -214,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=[
@@ -222,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,7 @@ 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
@@ -42,7 +42,7 @@ if TYPE_CHECKING:
     from google.api_core.retry import Retry
     from google.protobuf.field_mask_pb2 import FieldMask
 
-    from airflow.utils.context import Context
+    from airflow.providers.common.compat.sdk import Context
 
 
 class AlloyDBBaseOperator(GoogleCloudBaseOperator):
@@ -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: