Initial commit - ACME demo version

backend/beyond_metrics/io/s3.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import io
+import json
+from typing import Any, Dict, Tuple
+
+import boto3
+import pandas as pd
+from matplotlib.figure import Figure
+
+from .base import DataSource, ResultsSink
+
+
+def _split_s3_path(path: str) -> Tuple[str, str]:
+    """
+    Convierte 's3://bucket/key' en (bucket, key).
+    """
+    if not path.startswith("s3://"):
+        raise ValueError(f"Ruta S3 inválida: {path}")
+
+    without_scheme = path[len("s3://") :]
+    parts = without_scheme.split("/", 1)
+    if len(parts) != 2:
+        raise ValueError(f"Ruta S3 inválida: {path}")
+    return parts[0], parts[1]
+
+
+class S3DataSource(DataSource):
+    """
+    DataSource que lee CSV desde S3 usando boto3.
+    """
+
+    def __init__(self, boto3_client: Any | None = None) -> None:
+        self.s3 = boto3_client or boto3.client("s3")
+
+    def read_csv(self, path: str) -> pd.DataFrame:
+        bucket, key = _split_s3_path(path)
+        obj = self.s3.get_object(Bucket=bucket, Key=key)
+        body = obj["Body"].read()
+        buffer = io.BytesIO(body)
+        return pd.read_csv(buffer)
+
+
+class S3ResultsSink(ResultsSink):
+    """
+    ResultsSink que escribe JSON e imágenes en S3.
+    """
+
+    def __init__(self, boto3_client: Any | None = None) -> None:
+        self.s3 = boto3_client or boto3.client("s3")
+
+    def write_json(self, path: str, data: Dict[str, Any]) -> None:
+        bucket, key = _split_s3_path(path)
+        body = json.dumps(data, ensure_ascii=False, indent=2).encode("utf-8")
+        self.s3.put_object(Bucket=bucket, Key=key, Body=body)
+
+    def write_figure(self, path: str, fig: Figure) -> None:
+        bucket, key = _split_s3_path(path)
+        buf = io.BytesIO()
+        fig.savefig(buf, format="png", bbox_inches="tight")
+        buf.seek(0)
+        self.s3.put_object(Bucket=bucket, Key=key, Body=buf.getvalue(), ContentType="image/png")