apache-airflow-providers-amazon 9.5.0rc2__tar.gz → 9.5.0rc3__tar.gz

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: apache-airflow-providers-amazon
-Version: 9.5.0rc2
+Version: 9.5.0rc3
 Summary: Provider package apache-airflow-providers-amazon for Apache Airflow
 Keywords: airflow-provider,amazon,airflow,integration
 Author-email: Apache Software Foundation <dev@airflow.apache.org>

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/pyproject.toml

@@ -25,7 +25,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "apache-airflow-providers-amazon"
-version = "9.5.0.rc2"
+version = "9.5.0.rc3"
 description = "Provider package apache-airflow-providers-amazon for Apache Airflow"
 readme = "README.rst"
 authors = [
@@ -164,8 +164,9 @@ dev = [
     "apache-airflow-providers-ssh",
     # Additional devel dependencies (do not remove this line and add extra development dependencies)
     "aiobotocore>=2.13.0",
+    "apache-airflow-providers-postgres",
     "aws_xray_sdk>=2.12.0",
-    "moto[cloudformation,glue]>=5.0.0",
+    "moto[cloudformation,glue]>=5.1.2",
     "mypy-boto3-appflow>=1.37.0",
     "mypy-boto3-rds>=1.34.90",
     "mypy-boto3-redshift-data>=1.34.0",

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/auth_manager/aws_auth_manager.py

@@ -27,15 +27,6 @@ from fastapi import FastAPI
 
 from airflow.api_fastapi.app import AUTH_MANAGER_FASTAPI_APP_PREFIX
 from airflow.api_fastapi.auth.managers.base_auth_manager import BaseAuthManager
-from airflow.api_fastapi.auth.managers.models.resource_details import (
-    AccessView,
-    BackfillDetails,
-    ConnectionDetails,
-    DagAccessEntity,
-    DagDetails,
-    PoolDetails,
-    VariableDetails,
-)
 from airflow.cli.cli_config import CLICommand, DefaultHelpParser, GroupCommand
 from airflow.configuration import conf
 from airflow.exceptions import AirflowOptionalProviderFeatureException
@@ -59,9 +50,16 @@ if TYPE_CHECKING:
         IsAuthorizedVariableRequest,
     )
     from airflow.api_fastapi.auth.managers.models.resource_details import (
+        AccessView,
         AssetAliasDetails,
         AssetDetails,
+        BackfillDetails,
         ConfigurationDetails,
+        ConnectionDetails,
+        DagAccessEntity,
+        DagDetails,
+        PoolDetails,
+        VariableDetails,
     )
     from airflow.api_fastapi.common.types import MenuItem
 
@@ -256,7 +254,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.CONNECTION,
-                "entity_id": cast(ConnectionDetails, request["details"]).conn_id
+                "entity_id": cast("ConnectionDetails", request["details"]).conn_id
                 if request.get("details")
                 else None,
             }
@@ -274,10 +272,10 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.DAG,
-                "entity_id": cast(DagDetails, request["details"]).id if request.get("details") else None,
+                "entity_id": cast("DagDetails", request["details"]).id if request.get("details") else None,
                 "context": {
                     "dag_entity": {
-                        "string": cast(DagAccessEntity, request["access_entity"]).value,
+                        "string": cast("DagAccessEntity", request["access_entity"]).value,
                     },
                 }
                 if request.get("access_entity")
@@ -297,7 +295,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.POOL,
-                "entity_id": cast(PoolDetails, request["details"]).name if request.get("details") else None,
+                "entity_id": cast("PoolDetails", request["details"]).name if request.get("details") else None,
             }
             for request in requests
         ]
@@ -313,7 +311,7 @@ class AwsAuthManager(BaseAuthManager[AwsAuthManagerUser]):
             {
                 "method": request["method"],
                 "entity_type": AvpEntities.VARIABLE,
-                "entity_id": cast(VariableDetails, request["details"]).key
+                "entity_id": cast("VariableDetails", request["details"]).key
                 if request.get("details")
                 else None,
             }

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/auth_manager/router/login.py

@@ -83,7 +83,9 @@ def login_callback(request: Request):
     url = conf.get("api", "base_url")
     token = get_auth_manager().generate_jwt(user)
     response = RedirectResponse(url=url, status_code=303)
-    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=True)
+
+    secure = conf.has_option("api", "ssl_cert")
+    response.set_cookie(COOKIE_NAME_JWT_TOKEN, token, secure=secure)
     return response
 
 

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/executors/ecs/ecs_executor.py

@@ -278,7 +278,7 @@ class AwsEcsExecutor(BaseExecutor):
         if not has_exit_codes:
             return ""
         reasons = [
-            f'{container["container_arn"]} - {container["reason"]}'
+            f"{container['container_arn']} - {container['reason']}"
             for container in containers
             if "reason" in container
         ]

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/athena_sql.py

@@ -146,10 +146,10 @@ class AthenaSQLHook(AwsBaseHook, DbApiHook):
         creds = self.get_credentials(region_name=conn_params["region_name"])
 
         return URL.create(
-            f'awsathena+{conn_params["driver"]}',
+            f"awsathena+{conn_params['driver']}",
             username=creds.access_key,
             password=creds.secret_key,
-            host=f'athena.{conn_params["region_name"]}.{conn_params["aws_domain"]}',
+            host=f"athena.{conn_params['region_name']}.{conn_params['aws_domain']}",
             port=443,
             database=conn_params["schema_name"],
             query={"aws_session_token": creds.token, **self.conn.extra_dejson},

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/batch_client.py

@@ -416,8 +416,7 @@ class BatchClientHook(AwsBaseHook):
                 )
         else:
             raise AirflowException(
-                f"AWS Batch job ({job_id}) description error: exceeded status_retries "
-                f"({self.status_retries})"
+                f"AWS Batch job ({job_id}) description error: exceeded status_retries ({self.status_retries})"
             )
 
     @staticmethod

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/batch_waiters.py

@@ -30,7 +30,7 @@ import json
 import sys
 from copy import deepcopy
 from pathlib import Path
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 import botocore.client
 import botocore.exceptions
@@ -144,7 +144,12 @@ class BatchWaitersHook(BatchClientHook):
         return self._waiter_model
 
     def get_waiter(
-        self, waiter_name: str, _: dict[str, str] | None = None, deferrable: bool = False, client=None
+        self,
+        waiter_name: str,
+        parameters: dict[str, str] | None = None,
+        config_overrides: dict[str, Any] | None = None,
+        deferrable: bool = False,
+        client=None,
     ) -> botocore.waiter.Waiter:
         """
         Get an AWS Batch service waiter, using the configured ``.waiter_model``.
@@ -175,7 +180,10 @@ class BatchWaitersHook(BatchClientHook):
             the name (including the casing) of the key name in the waiter
             model file (typically this is CamelCasing); see ``.list_waiters``.
 
-        :param _: unused, just here to match the method signature in base_aws
+        :param parameters: unused, just here to match the method signature in base_aws
+        :param config_overrides: unused, just here to match the method signature in base_aws
+        :param deferrable: unused, just here to match the method signature in base_aws
+        :param client: unused, just here to match the method signature in base_aws
 
         :return: a waiter object for the named AWS Batch service
         """

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/dms.py

@@ -292,7 +292,9 @@ class DmsHook(AwsBaseHook):
             return arn
 
         except ClientError as err:
-            err_str = f"Error: {err.get('Error','').get('Code','')}: {err.get('Error','').get('Message','')}"
+            err_str = (
+                f"Error: {err.get('Error', '').get('Code', '')}: {err.get('Error', '').get('Message', '')}"
+            )
             self.log.error("Error while creating replication config: %s", err_str)
             raise err
 

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/redshift_cluster.py

@@ -67,7 +67,7 @@ class RedshiftHook(AwsBaseHook):
             for the cluster that is being created.
         :param params: Remaining AWS Create cluster API params.
         """
-        response = self.get_conn().create_cluster(
+        response = self.conn.create_cluster(
             ClusterIdentifier=cluster_identifier,
             NodeType=node_type,
             MasterUsername=master_username,
@@ -87,9 +87,9 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         """
         try:
-            response = self.get_conn().describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
+            response = self.conn.describe_clusters(ClusterIdentifier=cluster_identifier)["Clusters"]
             return response[0]["ClusterStatus"] if response else None
-        except self.get_conn().exceptions.ClusterNotFoundFault:
+        except self.conn.exceptions.ClusterNotFoundFault:
             return "cluster_not_found"
 
     async def cluster_status_async(self, cluster_identifier: str) -> str:
@@ -115,7 +115,7 @@ class RedshiftHook(AwsBaseHook):
         """
         final_cluster_snapshot_identifier = final_cluster_snapshot_identifier or ""
 
-        response = self.get_conn().delete_cluster(
+        response = self.conn.delete_cluster(
             ClusterIdentifier=cluster_identifier,
             SkipFinalClusterSnapshot=skip_final_cluster_snapshot,
             FinalClusterSnapshotIdentifier=final_cluster_snapshot_identifier,
@@ -131,7 +131,7 @@ class RedshiftHook(AwsBaseHook):
 
         :param cluster_identifier: unique identifier of a cluster
         """
-        response = self.get_conn().describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
+        response = self.conn.describe_cluster_snapshots(ClusterIdentifier=cluster_identifier)
         if "Snapshots" not in response:
             return None
         snapshots = response["Snapshots"]
@@ -149,7 +149,7 @@ class RedshiftHook(AwsBaseHook):
         :param cluster_identifier: unique identifier of a cluster
         :param snapshot_identifier: unique identifier for a snapshot of a cluster
         """
-        response = self.get_conn().restore_from_cluster_snapshot(
+        response = self.conn.restore_from_cluster_snapshot(
             ClusterIdentifier=cluster_identifier, SnapshotIdentifier=snapshot_identifier
         )
         return response["Cluster"] if response["Cluster"] else None
@@ -175,7 +175,7 @@ class RedshiftHook(AwsBaseHook):
         """
         if tags is None:
             tags = []
-        response = self.get_conn().create_cluster_snapshot(
+        response = self.conn.create_cluster_snapshot(
             SnapshotIdentifier=snapshot_identifier,
             ClusterIdentifier=cluster_identifier,
             ManualSnapshotRetentionPeriod=retention_period,
@@ -192,11 +192,11 @@ class RedshiftHook(AwsBaseHook):
         :param snapshot_identifier: A unique identifier for the snapshot that you are requesting
         """
         try:
-            response = self.get_conn().describe_cluster_snapshots(
+            response = self.conn.describe_cluster_snapshots(
                 SnapshotIdentifier=snapshot_identifier,
             )
             snapshot = response.get("Snapshots")[0]
             snapshot_status: str = snapshot.get("Status")
             return snapshot_status
-        except self.get_conn().exceptions.ClusterSnapshotNotFoundFault:
+        except self.conn.exceptions.ClusterSnapshotNotFoundFault:
             return None

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/redshift_data.py

@@ -186,8 +186,7 @@ class RedshiftDataHook(AwsGenericHook["RedshiftDataAPIServiceClient"]):
                 RedshiftDataQueryFailedError if status == FAILED_STATE else RedshiftDataQueryAbortedError
             )
             raise exception_cls(
-                f"Statement {resp['Id']} terminated with status {status}. "
-                f"Response details: {pformat(resp)}"
+                f"Statement {resp['Id']} terminated with status {status}. Response details: {pformat(resp)}"
             )
 
         self.log.info("Query status: %s", status)

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/hooks/sagemaker.py

@@ -131,7 +131,7 @@ def secondary_training_status_message(
     status_strs = []
     for transition in transitions_to_print:
         message = transition["StatusMessage"]
-        time_utc = timezone.convert_to_utc(cast(datetime, job_description["LastModifiedTime"]))
+        time_utc = timezone.convert_to_utc(cast("datetime", job_description["LastModifiedTime"]))
         status_strs.append(f"{time_utc:%Y-%m-%d %H:%M:%S} {transition['Status']} - {message}")
 
     return "\n".join(status_strs)

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/links/athena.py

@@ -25,6 +25,5 @@ class AthenaQueryResultsLink(BaseAwsLink):
     name = "Query Results"
     key = "_athena_query_results"
     format_str = (
-        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#"
-        "/query-editor/history/{query_execution_id}"
+        BASE_AWS_CONSOLE_LINK + "/athena/home?region={region_name}#/query-editor/history/{query_execution_id}"
     )

{apache_airflow_providers_amazon-9.5.0rc2 → apache_airflow_providers_amazon-9.5.0rc3}/src/airflow/providers/amazon/aws/log/cloudwatch_task_handler.py

@@ -17,20 +17,31 @@
 # under the License.
 from __future__ import annotations
 
+import copy
+import json
+import logging
+import os
 from datetime import date, datetime, timedelta, timezone
 from functools import cached_property
+from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import attrs
 import watchtower
 
 from airflow.configuration import conf
 from airflow.providers.amazon.aws.hooks.logs import AwsLogsHook
 from airflow.providers.amazon.aws.utils import datetime_to_epoch_utc_ms
+from airflow.providers.amazon.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
 
 if TYPE_CHECKING:
-    from airflow.models import TaskInstance
+    import structlog.typing
+
+    from airflow.models.taskinstance import TaskInstance
+    from airflow.sdk.types import RuntimeTaskInstanceProtocol as RuntimeTI
+    from airflow.utils.log.file_task_handler import LogMessages, LogSourceInfo
 
 
 def json_serialize_legacy(value: Any) -> str | None:
@@ -62,6 +73,155 @@ def json_serialize(value: Any) -> str | None:
     return watchtower._json_serialize_default(value)
 
 
+@attrs.define(kw_only=True)
+class CloudWatchRemoteLogIO(LoggingMixin):  # noqa: D101
+    base_log_folder: Path = attrs.field(converter=Path)
+    remote_base: str = ""
+    delete_local_copy: bool = True
+
+    log_group_arn: str
+    log_stream_name: str = ""
+    log_group: str = attrs.field(init=False, repr=False)
+    region_name: str = attrs.field(init=False, repr=False)
+
+    @log_group.default
+    def _(self):
+        return self.log_group_arn.split(":")[6]
+
+    @region_name.default
+    def _(self):
+        return self.log_group_arn.split(":")[3]
+
+    @cached_property
+    def hook(self):
+        """Returns AwsLogsHook."""
+        return AwsLogsHook(
+            aws_conn_id=conf.get("logging", "remote_log_conn_id"), region_name=self.region_name
+        )
+
+    @cached_property
+    def handler(self) -> watchtower.CloudWatchLogHandler:
+        _json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer", fallback=None)
+        return watchtower.CloudWatchLogHandler(
+            log_group_name=self.log_group,
+            log_stream_name=self.log_stream_name,
+            use_queues=True,
+            boto3_client=self.hook.get_conn(),
+            json_serialize_default=_json_serialize or json_serialize_legacy,
+        )
+
+    @cached_property
+    def processors(self) -> tuple[structlog.typing.Processor, ...]:
+        from logging import getLogRecordFactory
+
+        import structlog.stdlib
+
+        logRecordFactory = getLogRecordFactory()
+        # The handler MUST be initted here, before the processor is actually used to log anything.
+        # Otherwise, logging that occurs during the creation of the handler can create infinite loops.
+        _handler = self.handler
+        from airflow.sdk.log import relative_path_from_logger
+
+        def proc(logger: structlog.typing.WrappedLogger, method_name: str, event: structlog.typing.EventDict):
+            if not logger or not (stream_name := relative_path_from_logger(logger)):
+                return event
+            # Only init the handler stream_name once. We cannot do it above when we init the handler because
+            # we don't yet know the log path at that point.
+            if not _handler.log_stream_name:
+                _handler.log_stream_name = stream_name.as_posix().replace(":", "_")
+            name = event.get("logger_name") or event.get("logger", "")
+            level = structlog.stdlib.NAME_TO_LEVEL.get(method_name.lower(), logging.INFO)
+            msg = copy.copy(event)
+            created = None
+            if ts := msg.pop("timestamp", None):
+                try:
+                    created = datetime.fromisoformat(ts)
+                except Exception:
+                    pass
+            record = logRecordFactory(
+                name, level, pathname="", lineno=0, msg=msg, args=(), exc_info=None, func=None, sinfo=None
+            )
+            if created is not None:
+                ct = created.timestamp()
+                record.created = ct
+                record.msecs = int((ct - int(ct)) * 1000) + 0.0  # Copied from stdlib logging
+            _handler.handle(record)
+            return event
+
+        return (proc,)
+
+    def close(self):
+        self.handler.close()
+
+    def upload(self, path: os.PathLike | str, ti: RuntimeTI):
+        # No-op, as we upload via the processor as we go
+        # But we need to give the handler time to finish off its business
+        self.close()
+        return
+
+    def read(self, relative_path, ti: RuntimeTI) -> tuple[LogSourceInfo, LogMessages | None]:
+        logs: LogMessages | None = []
+        messages = [
+            f"Reading remote log from Cloudwatch log_group: {self.log_group} log_stream: {relative_path}"
+        ]
+        try:
+            if AIRFLOW_V_3_0_PLUS:
+                from airflow.utils.log.file_task_handler import StructuredLogMessage
+
+                logs = [
+                    StructuredLogMessage.model_validate(log)
+                    for log in self.get_cloudwatch_logs(relative_path, ti)
+                ]
+            else:
+                logs = [self.get_cloudwatch_logs(relative_path, ti)]  # type: ignore[arg-value]
+        except Exception as e:
+            logs = None
+            messages.append(str(e))
+
+        return messages, logs
+
+    def get_cloudwatch_logs(self, stream_name: str, task_instance: RuntimeTI):
+        """
+        Return all logs from the given log stream.
+
+        :param stream_name: name of the Cloudwatch log stream to get all logs from
+        :param task_instance: the task instance to get logs about
+        :return: string of all logs from the given log stream
+        """
+        stream_name = stream_name.replace(":", "_")
+        # If there is an end_date to the task instance, fetch logs until that date + 30 seconds
+        # 30 seconds is an arbitrary buffer so that we don't miss any logs that were emitted
+        end_time = (
+            None
+            if (end_date := getattr(task_instance, "end_date", None)) is None
+            else datetime_to_epoch_utc_ms(end_date + timedelta(seconds=30))
+        )
+        events = self.hook.get_log_events(
+            log_group=self.log_group,
+            log_stream_name=stream_name,
+            end_time=end_time,
+        )
+        if AIRFLOW_V_3_0_PLUS:
+            return list(self._event_to_dict(e) for e in events)
+        return "\n".join(self._event_to_str(event) for event in events)
+
+    def _event_to_dict(self, event: dict) -> dict:
+        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc).isoformat()
+        message = event["message"]
+        try:
+            message = json.loads(message)
+            message["timestamp"] = event_dt
+            return message
+        except Exception:
+            return {"timestamp": event_dt, "event": message}
+
+    def _event_to_str(self, event: dict) -> str:
+        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
+        formatted_event_dt = event_dt.strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
+        message = event["message"]
+        return f"[{formatted_event_dt}] {message}"
+
+
 class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
     """
     CloudwatchTaskHandler is a python log handler that handles and reads task instance logs.
@@ -84,6 +244,11 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
         self.region_name = split_arn[3]
         self.closed = False
 
+        self.io = CloudWatchRemoteLogIO(
+            base_log_folder=base_log_folder,
+            log_group_arn=log_group_arn,
+        )
+
     @cached_property
     def hook(self):
         """Returns AwsLogsHook."""
@@ -97,14 +262,9 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
 
     def set_context(self, ti: TaskInstance, *, identifier: str | None = None):
         super().set_context(ti)
-        _json_serialize = conf.getimport("aws", "cloudwatch_task_handler_json_serializer", fallback=None)
-        self.handler = watchtower.CloudWatchLogHandler(
-            log_group_name=self.log_group,
-            log_stream_name=self._render_filename(ti, ti.try_number),
-            use_queues=not getattr(ti, "is_trigger_log_context", False),
-            boto3_client=self.hook.get_conn(),
-            json_serialize_default=_json_serialize or json_serialize_legacy,
-        )
+        self.io.log_stream_name = self._render_filename(ti, ti.try_number)
+
+        self.handler = self.io.handler
 
     def close(self):
         """Close the handler responsible for the upload of the local log file to Cloudwatch."""
@@ -120,49 +280,9 @@ class CloudwatchTaskHandler(FileTaskHandler, LoggingMixin):
         # Mark closed so we don't double write if close is called twice
         self.closed = True
 
-    def _read(self, task_instance, try_number, metadata=None):
+    def _read_remote_logs(
+        self, task_instance, try_number, metadata=None
+    ) -> tuple[LogSourceInfo, LogMessages]:
         stream_name = self._render_filename(task_instance, try_number)
-        try:
-            return (
-                f"*** Reading remote log from Cloudwatch log_group: {self.log_group} "
-                f"log_stream: {stream_name}.\n"
-                f"{self.get_cloudwatch_logs(stream_name=stream_name, task_instance=task_instance)}\n",
-                {"end_of_log": True},
-            )
-        except Exception as e:
-            log = (
-                f"*** Unable to read remote logs from Cloudwatch (log_group: {self.log_group}, log_stream: "
-                f"{stream_name})\n*** {e}\n\n"
-            )
-            self.log.error(log)
-            local_log, metadata = super()._read(task_instance, try_number, metadata)
-            log += local_log
-            return log, metadata
-
-    def get_cloudwatch_logs(self, stream_name: str, task_instance: TaskInstance) -> str:
-        """
-        Return all logs from the given log stream.
-
-        :param stream_name: name of the Cloudwatch log stream to get all logs from
-        :param task_instance: the task instance to get logs about
-        :return: string of all logs from the given log stream
-        """
-        # If there is an end_date to the task instance, fetch logs until that date + 30 seconds
-        # 30 seconds is an arbitrary buffer so that we don't miss any logs that were emitted
-        end_time = (
-            None
-            if task_instance.end_date is None
-            else datetime_to_epoch_utc_ms(task_instance.end_date + timedelta(seconds=30))
-        )
-        events = self.hook.get_log_events(
-            log_group=self.log_group,
-            log_stream_name=stream_name,
-            end_time=end_time,
-        )
-        return "\n".join(self._event_to_str(event) for event in events)
-
-    def _event_to_str(self, event: dict) -> str:
-        event_dt = datetime.fromtimestamp(event["timestamp"] / 1000.0, tz=timezone.utc)
-        formatted_event_dt = event_dt.strftime("%Y-%m-%d %H:%M:%S,%f")[:-3]
-        message = event["message"]
-        return f"[{formatted_event_dt}] {message}"
+        messages, logs = self.io.read(stream_name, task_instance)
+        return messages, logs or []