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

airflow/providers/google/cloud/log/gcs_task_handler.py

@@ -27,11 +27,11 @@ from typing import TYPE_CHECKING
 
 import attrs
 
-# not sure why but mypy complains on missing `storage` but it is clearly there and is importable
-from google.cloud import storage  # type: ignore[attr-defined]
+# Make mypy happy by importing as aliases
+import google.cloud.storage as storage
 
 from airflow.configuration import conf
-from airflow.exceptions import AirflowNotFoundException
+from airflow.providers.common.compat.sdk import AirflowNotFoundException
 from airflow.providers.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
 from airflow.providers.google.cloud.utils.credentials_provider import (
     get_credentials_and_project_id,
@@ -43,9 +43,11 @@ from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
 if TYPE_CHECKING:
+    from io import TextIOWrapper
+
     from airflow.models.taskinstance import TaskInstance
     from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
-    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
+    from airflow.utils.log.file_task_handler import LogResponse, RawLogStream, StreamingLogResponse
 
 _DEFAULT_SCOPESS = frozenset(
     [
@@ -61,13 +63,15 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
     remote_base: str
     base_log_folder: Path = attrs.field(converter=Path)
     delete_local_copy: bool
+    project_id: str | None = None
+
+    gcp_key_path: str | None = None
+    gcp_keyfile_dict: dict | None = None
+    scopes: Collection[str] | None = _DEFAULT_SCOPESS
 
-    gcp_key_path: str | None
-    gcp_keyfile_dict: dict | None
-    scopes: Collection[str] | None
-    project_id: str
+    processors = ()
 
-    def upload(self, path: os.PathLike, ti: RuntimeTI):
+    def upload(self, path: os.PathLike | str, ti: RuntimeTI):
         """Upload the given log path to the remote storage."""
         path = Path(path)
         if path.is_absolute():
@@ -147,11 +151,26 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
             exc, "resp", {}
         ).get("status") == "404"
 
-    def read(self, relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
-        messages = []
-        logs = []
+    def read(self, relative_path: str, ti: RuntimeTI) -> LogResponse:
+        messages, log_streams = self.stream(relative_path, ti)
+        if not log_streams:
+            return messages, None
+
+        logs: list[str] = []
+        try:
+            # for each log_stream, exhaust the generator into a string
+            logs = ["".join(line for line in log_stream) for log_stream in log_streams]
+        except Exception as e:
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"Unable to read remote log {e}")
+
+        return messages, logs
+
+    def stream(self, relative_path: str, ti: RuntimeTI) -> StreamingLogResponse:
+        messages: list[str] = []
+        log_streams: list[RawLogStream] = []
         remote_loc = os.path.join(self.remote_base, relative_path)
-        uris = []
+        uris: list[str] = []
         bucket, prefix = _parse_gcs_url(remote_loc)
         blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
 
@@ -162,18 +181,29 @@ class GCSRemoteLogIO(LoggingMixin):  # noqa: D101
             else:
                 messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
         else:
-            return messages, None
+            return messages, []
 
         try:
             for key in sorted(uris):
                 blob = storage.Blob.from_string(key, self.client)
-                remote_log = blob.download_as_bytes().decode()
-                if remote_log:
-                    logs.append(remote_log)
+                stream = blob.open("r")
+                log_streams.append(self._get_log_stream(stream))
         except Exception as e:
             if not AIRFLOW_V_3_0_PLUS:
                 messages.append(f"Unable to read remote log {e}")
-        return messages, logs
+        return messages, log_streams
+
+    def _get_log_stream(self, stream: TextIOWrapper) -> RawLogStream:
+        """
+        Yield lines from the given stream.
+
+        :param stream: The opened stream to read from.
+        :yield: Lines of the log file.
+        """
+        try:
+            yield from stream
+        finally:
+            stream.close()
 
 
 class GCSTaskHandler(FileTaskHandler, LoggingMixin):
@@ -211,9 +241,15 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         gcp_keyfile_dict: dict | None = None,
         gcp_scopes: Collection[str] | None = _DEFAULT_SCOPESS,
         project_id: str = PROVIDE_PROJECT_ID,
+        max_bytes: int = 0,
+        backup_count: int = 0,
+        delay: bool = False,
         **kwargs,
-    ):
-        super().__init__(base_log_folder)
+    ) -> None:
+        # support log file size handling of FileTaskHandler
+        super().__init__(
+            base_log_folder=base_log_folder, max_bytes=max_bytes, backup_count=backup_count, delay=delay
+        )
         self.handler: logging.FileHandler | None = None
         self.log_relative_path = ""
         self.closed = False
@@ -265,7 +301,7 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         # Mark closed so we don't double write if close is called twice
         self.closed = True
 
-    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[LogSourceInfo, LogMessages]:
+    def _read_remote_logs(self, ti, try_number, metadata=None) -> LogResponse:
         # Explicitly getting log relative path is necessary as the given
         # task instance might be different than task instance passed in
         # in set_context method.
@@ -275,7 +311,7 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
 
         if logs is None:
             logs = []
-            if not AIRFLOW_V_3_0_PLUS:
+            if not AIRFLOW_V_3_0_PLUS and not messages:
                 messages.append(f"No logs found in GCS; ti={ti}")
 
         return messages, logs

airflow/providers/google/cloud/log/stackdriver_task_handler.py

@@ -35,17 +35,20 @@ from airflow.exceptions import AirflowProviderDeprecationWarning
 from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
 from airflow.providers.google.common.consts import CLIENT_INFO
 from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
-from airflow.utils.types import NOTSET, ArgNotSet
+
+try:
+    from airflow.sdk.definitions._internal.types import NOTSET, ArgNotSet
+except ImportError:
+    from airflow.utils.types import NOTSET, ArgNotSet  # type: ignore[attr-defined,no-redef]
+
+if not AIRFLOW_V_3_0_PLUS:
+    from airflow.utils.log.trigger_handler import ctx_indiv_trigger
 
 if TYPE_CHECKING:
     from google.auth.credentials import Credentials
 
     from airflow.models import TaskInstance
 
-
-if not AIRFLOW_V_3_0_PLUS:
-    from airflow.utils.log.trigger_handler import ctx_indiv_trigger
-
 DEFAULT_LOGGER_NAME = "airflow"
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
@@ -159,7 +162,7 @@ class StackdriverTaskHandler(logging.Handler):
         """Object responsible for sending data to Stackdriver."""
         # The Transport object is badly defined (no init) but in the docs client/name as constructor
         # arguments are a requirement for any class that derives from Transport class, hence ignore:
-        return self.transport_type(self._client, self.gcp_log_name)  # type: ignore[call-arg]
+        return self.transport_type(self._client, self.gcp_log_name)
 
     def _get_labels(self, task_instance=None):
         if task_instance:

airflow/providers/google/cloud/openlineage/CloudStorageTransferJobFacet.json

@@ -0,0 +1,68 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "CloudStorageTransferJobFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/JobFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "jobName": {
+              "type": "string",
+              "description": "Transfer job name assigned by GCP Storage Transfer Service."
+            },
+            "projectId": {
+              "type": "string",
+              "description": "GCP project ID."
+            },
+            "description": {
+              "type": "string",
+              "description": "Optional description of the transfer job."
+            },
+            "status": {
+              "type": "string",
+              "description": "Status of the transfer job (ENABLED, DISABLED)."
+            },
+            "sourceBucket": {
+              "type": "string",
+              "description": "Source AWS S3 bucket."
+            },
+            "sourcePath": {
+              "type": "string",
+              "description": "Prefix path inside the source bucket."
+            },
+            "targetBucket": {
+              "type": "string",
+              "description": "Target GCS bucket."
+            },
+            "targetPath": {
+              "type": "string",
+              "description": "Prefix path inside the target bucket."
+            },
+            "objectConditions": {
+              "type": "object",
+              "description": "Filtering conditions for objects transferred."
+            },
+            "transferOptions": {
+              "type": "object",
+              "description": "Transfer options such as overwrite or delete."
+            },
+            "schedule": {
+              "type": "object",
+              "description": "Transfer schedule details."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "cloudStorageTransferJob": {
+      "$ref": "#/$defs/CloudStorageTransferJobFacet"
+    }
+  }
+}

airflow/providers/google/cloud/openlineage/CloudStorageTransferRunFacet.json

@@ -0,0 +1,60 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "CloudStorageTransferRunFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "jobName": {
+              "type": "string",
+              "description": "Transfer job name associated with this run."
+            },
+            "operationName": {
+              "type": "string",
+              "description": "Transfer operation name if available."
+            },
+            "status": {
+              "type": "string",
+              "description": "Run status if available."
+            },
+            "startTime": {
+              "type": "string",
+              "description": "Start time of the transfer operation."
+            },
+            "endTime": {
+              "type": "string",
+              "description": "End time of the transfer operation."
+            },
+            "wait": {
+              "type": "boolean",
+              "description": "Whether the operator waited for completion."
+            },
+            "timeout": {
+              "type": ["number", "null"],
+              "description": "Timeout in seconds."
+            },
+            "deferrable": {
+              "type": "boolean",
+              "description": "Whether the operator used deferrable mode."
+            },
+            "deleteJobAfterCompletion": {
+              "type": "boolean",
+              "description": "Whether the transfer job was deleted after completion."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "cloudStorageTransferRun": {
+      "$ref": "#/$defs/CloudStorageTransferRunFacet"
+    }
+  }
+}

airflow/providers/google/cloud/openlineage/DataFusionRunFacet.json

@@ -0,0 +1,32 @@
+{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "$defs": {
+    "DataFusionRunFacet": {
+      "allOf": [
+        {
+          "$ref": "https://openlineage.io/spec/2-0-2/OpenLineage.json#/$defs/RunFacet"
+        },
+        {
+          "type": "object",
+          "properties": {
+            "runId": {
+              "type": "string",
+              "description": "Pipeline run ID assigned by Cloud Data Fusion."
+            },
+            "runtimeArgs": {
+              "type": "object",
+              "description": "Runtime arguments provided when starting the pipeline."
+            }
+          }
+        }
+      ],
+      "type": "object"
+    }
+  },
+  "type": "object",
+  "properties": {
+    "dataFusionRun": {
+      "$ref": "#/$defs/DataFusionRunFacet"
+    }
+  }
+}

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: