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

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

@@ -25,20 +25,27 @@ from functools import cached_property
 from pathlib import Path
 from typing import TYPE_CHECKING
 
+import attrs
+
+# 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.google.cloud.hooks.gcs import GCSHook, _parse_gcs_url
-from airflow.providers.google.cloud.utils.credentials_provider import get_credentials_and_project_id
+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.common.hooks.base_google import PROVIDE_PROJECT_ID
+from airflow.providers.google.version_compat import AIRFLOW_V_3_0_PLUS
 from airflow.utils.log.file_task_handler import FileTaskHandler
 from airflow.utils.log.logging_mixin import LoggingMixin
 
-# 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]
-
 if TYPE_CHECKING:
     from airflow.models.taskinstance import TaskInstance
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
+    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
 
 _DEFAULT_SCOPESS = frozenset(
     [
@@ -49,6 +56,128 @@ _DEFAULT_SCOPESS = frozenset(
 logger = logging.getLogger(__name__)
 
 
+@attrs.define
+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
+
+    processors = ()
+
+    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():
+            local_loc = path
+            remote_loc = os.path.join(self.remote_base, path.relative_to(self.base_log_folder))
+        else:
+            local_loc = self.base_log_folder.joinpath(path)
+            remote_loc = os.path.join(self.remote_base, path)
+
+        if local_loc.is_file():
+            # read log and remove old logs to get just the latest additions
+            log = local_loc.read_text()
+            has_uploaded = self.write(log, remote_loc)
+            if has_uploaded and self.delete_local_copy:
+                shutil.rmtree(os.path.dirname(local_loc))
+
+    @cached_property
+    def hook(self) -> GCSHook | None:
+        """Returns GCSHook if remote_log_conn_id configured."""
+        conn_id = conf.get("logging", "remote_log_conn_id", fallback=None)
+        if conn_id:
+            try:
+                return GCSHook(gcp_conn_id=conn_id)
+            except AirflowNotFoundException:
+                pass
+        return None
+
+    @cached_property
+    def client(self) -> storage.Client:
+        """Returns GCS Client."""
+        if self.hook:
+            credentials, project_id = self.hook.get_credentials_and_project_id()
+        else:
+            credentials, project_id = get_credentials_and_project_id(
+                key_path=self.gcp_key_path,
+                keyfile_dict=self.gcp_keyfile_dict,
+                scopes=self.scopes,
+                disable_logging=True,
+            )
+        return storage.Client(
+            credentials=credentials,
+            client_info=CLIENT_INFO,
+            project=self.project_id if self.project_id else project_id,
+        )
+
+    def write(self, log: str, remote_log_location: str) -> bool:
+        """
+        Write the log to the remote location and return `True`; fail silently and return `False` on error.
+
+        :param log: the log to write to the remote_log_location
+        :param remote_log_location: the log's location in remote storage
+        :return: whether the log is successfully written to remote location or not.
+        """
+        try:
+            blob = storage.Blob.from_string(remote_log_location, self.client)
+            old_log = blob.download_as_bytes().decode()
+            log = f"{old_log}\n{log}" if old_log else log
+        except Exception as e:
+            if not self.no_log_found(e):
+                self.log.warning("Error checking for previous log: %s", e)
+        try:
+            blob = storage.Blob.from_string(remote_log_location, self.client)
+            blob.upload_from_string(log, content_type="text/plain")
+        except Exception as e:
+            self.log.error("Could not write logs to %s: %s", remote_log_location, e)
+            return False
+        return True
+
+    @staticmethod
+    def no_log_found(exc):
+        """
+        Given exception, determine whether it is result of log not found.
+
+        :meta private:
+        """
+        return (exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]) or getattr(
+            exc, "resp", {}
+        ).get("status") == "404"
+
+    def read(self, relative_path: str, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
+        messages = []
+        logs = []
+        remote_loc = os.path.join(self.remote_base, relative_path)
+        uris = []
+        bucket, prefix = _parse_gcs_url(remote_loc)
+        blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
+
+        if blobs:
+            uris = [f"gs://{bucket}/{b.name}" for b in blobs]
+            if AIRFLOW_V_3_0_PLUS:
+                messages = uris
+            else:
+                messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
+        else:
+            return messages, None
+
+        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)
+        except Exception as e:
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"Unable to read remote log {e}")
+        return messages, logs
+
+
 class GCSTaskHandler(FileTaskHandler, LoggingMixin):
     """
     GCSTaskHandler is a python log handler that handles and reads task instance logs.
@@ -84,49 +213,29 @@ 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.remote_base = gcs_log_folder
         self.log_relative_path = ""
         self.closed = False
         self.upload_on_close = True
-        self.gcp_key_path = gcp_key_path
-        self.gcp_keyfile_dict = gcp_keyfile_dict
-        self.scopes = gcp_scopes
-        self.project_id = project_id
-        self.delete_local_copy = kwargs.get(
-            "delete_local_copy", conf.getboolean("logging", "delete_local_logs")
-        )
-
-    @cached_property
-    def hook(self) -> GCSHook | None:
-        """Returns GCSHook if remote_log_conn_id configured."""
-        conn_id = conf.get("logging", "remote_log_conn_id", fallback=None)
-        if conn_id:
-            try:
-                return GCSHook(gcp_conn_id=conn_id)
-            except AirflowNotFoundException:
-                pass
-        return None
-
-    @cached_property
-    def client(self) -> storage.Client:
-        """Returns GCS Client."""
-        if self.hook:
-            credentials, project_id = self.hook.get_credentials_and_project_id()
-        else:
-            credentials, project_id = get_credentials_and_project_id(
-                key_path=self.gcp_key_path,
-                keyfile_dict=self.gcp_keyfile_dict,
-                scopes=self.scopes,
-                disable_logging=True,
-            )
-        return storage.Client(
-            credentials=credentials,
-            client_info=CLIENT_INFO,
-            project=self.project_id if self.project_id else project_id,
+        self.io = GCSRemoteLogIO(
+            base_log_folder=base_log_folder,
+            remote_base=gcs_log_folder,
+            delete_local_copy=kwargs.get(
+                "delete_local_copy", conf.getboolean("logging", "delete_local_logs")
+            ),
+            gcp_key_path=gcp_key_path,
+            gcp_keyfile_dict=gcp_keyfile_dict,
+            scopes=gcp_scopes,
+            project_id=project_id,
         )
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None) -> None:
@@ -137,6 +246,8 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         if TYPE_CHECKING:
             assert self.handler is not None
 
+        self.ti = ti
+
         full_path = self.handler.baseFilename
         self.log_relative_path = Path(full_path).relative_to(self.local_base).as_posix()
         is_trigger_log_context = getattr(ti, "is_trigger_log_context", False)
@@ -156,86 +267,23 @@ class GCSTaskHandler(FileTaskHandler, LoggingMixin):
         if not self.upload_on_close:
             return
 
-        local_loc = os.path.join(self.local_base, self.log_relative_path)
-        remote_loc = os.path.join(self.remote_base, self.log_relative_path)
-        if os.path.exists(local_loc):
-            # read log and remove old logs to get just the latest additions
-            with open(local_loc) as logfile:
-                log = logfile.read()
-            gcs_write = self.gcs_write(log, remote_loc)
-            if gcs_write and self.delete_local_copy:
-                shutil.rmtree(os.path.dirname(local_loc))
+        if hasattr(self, "ti"):
+            self.io.upload(self.log_relative_path, self.ti)
 
         # Mark closed so we don't double write if close is called twice
         self.closed = True
 
-    def _add_message(self, msg):
-        filename, lineno, func, stackinfo = logger.findCaller()
-        record = logging.LogRecord("", logging.INFO, filename, lineno, msg + "\n", None, None, func=func)
-        return self.format(record)
-
-    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[list[str], list[str]]:
-        # Explicitly getting log relative path is necessary because this method
-        # is called from webserver from TaskLogReader, where we don't call set_context
-        # and can read logs for different TIs in each request
-        messages = []
-        logs = []
-        worker_log_relative_path = self._render_filename(ti, try_number)
-        remote_loc = os.path.join(self.remote_base, worker_log_relative_path)
-        uris = []
-        bucket, prefix = _parse_gcs_url(remote_loc)
-        blobs = list(self.client.list_blobs(bucket_or_name=bucket, prefix=prefix))
-
-        if blobs:
-            uris = [f"gs://{bucket}/{b.name}" for b in blobs]
-            messages.extend(["Found remote logs:", *[f"  * {x}" for x in sorted(uris)]])
-        else:
-            messages.append(f"No logs found in GCS; ti=%s {ti}")
-        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)
-        except Exception as e:
-            messages.append(f"Unable to read remote log {e}")
-        return messages, logs
+    def _read_remote_logs(self, ti, try_number, metadata=None) -> tuple[LogSourceInfo, LogMessages]:
+        # Explicitly getting log relative path is necessary as the given
+        # task instance might be different than task instance passed in
+        # in set_context method.
+        worker_log_rel_path = self._render_filename(ti, try_number)
 
-    def gcs_write(self, log, remote_log_location) -> bool:
-        """
-        Write the log to the remote location and return `True`; fail silently and return `False` on error.
+        messages, logs = self.io.read(worker_log_rel_path, ti)
 
-        :param log: the log to write to the remote_log_location
-        :param remote_log_location: the log's location in remote storage
-        :return: whether the log is successfully written to remote location or not.
-        """
-        try:
-            blob = storage.Blob.from_string(remote_log_location, self.client)
-            old_log = blob.download_as_bytes().decode()
-            log = f"{old_log}\n{log}" if old_log else log
-        except Exception as e:
-            if not self.no_log_found(e):
-                log += self._add_message(
-                    f"Error checking for previous log; if exists, may be overwritten: {e}"
-                )
-                self.log.warning("Error checking for previous log: %s", e)
-        try:
-            blob = storage.Blob.from_string(remote_log_location, self.client)
-            blob.upload_from_string(log, content_type="text/plain")
-        except Exception as e:
-            self.log.error("Could not write logs to %s: %s", remote_log_location, e)
-            return False
-        return True
-
-    @staticmethod
-    def no_log_found(exc):
-        """
-        Given exception, determine whether it is result of log not found.
+        if logs is None:
+            logs = []
+            if not AIRFLOW_V_3_0_PLUS:
+                messages.append(f"No logs found in GCS; ti={ti}")
 
-        :meta private:
-        """
-        if (exc.args and isinstance(exc.args[0], str) and "No such object" in exc.args[0]) or getattr(
-            exc, "resp", {}
-        ).get("status") == "404":
-            return True
-        return False
+        return messages, logs

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

@@ -25,25 +25,30 @@ from functools import cached_property
 from typing import TYPE_CHECKING
 from urllib.parse import urlencode
 
-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
 from google.cloud import logging as gcp_logging
 from google.cloud.logging import Resource
 from google.cloud.logging.handlers.transports import BackgroundThreadTransport, Transport
 from google.cloud.logging_v2.services.logging_service_v2 import LoggingServiceV2Client
 from google.cloud.logging_v2.types import ListLogEntriesRequest, ListLogEntriesResponse
 
-if TYPE_CHECKING:
-    from airflow.models import TaskInstance
-    from google.auth.credentials import Credentials
+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
 
+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
+
 DEFAULT_LOGGER_NAME = "airflow"
 _GLOBAL_RESOURCE = Resource(type="global", labels={})
 
@@ -157,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"
+    }
+  }
+}