""" Shared Storage Helper for MinIO / S3 Landing & Bronze Staging. Supports S3 / MinIO uploads with graceful local disk staging fallback. """ import os import io import json import logging from datetime import datetime from typing import Any, Optional, Dict import pandas as pd logger = logging.getLogger(__name__) # Default base directory for local landing staging fallback DEFAULT_LOCAL_STAGING = os.getenv("DLH_LOCAL_STAGING_DIR", "/tmp/dlh_landing") def get_minio_s3_client(): """Attempt to create a boto3 S3 client for MinIO if boto3 is installed and configured.""" try: import boto3 endpoint = os.getenv("MINIO_ENDPOINT", "http://minio:9000") access_key = os.getenv("MINIO_ROOT_USER", os.getenv("AWS_ACCESS_KEY_ID", "minio_admin")) secret_key = os.getenv("MINIO_ROOT_PASSWORD", os.getenv("AWS_SECRET_ACCESS_KEY", "minio_password")) client = boto3.client( "s3", endpoint_url=endpoint, aws_access_key_id=access_key, aws_secret_access_key=secret_key, region_name="us-east-1" ) return client except Exception as e: logger.debug("boto3 client not available or initialization failed: %s", e) return None def build_partitioned_path(prefix: str, filename: str, dt: Optional[datetime] = None) -> str: """Generate Hive-style partitioned path: prefix/year=YYYY/month=MM/day=DD/filename""" target_dt = dt or datetime.utcnow() clean_prefix = prefix.strip("/") year_part = f"year={target_dt.strftime('%Y')}" month_part = f"month={target_dt.strftime('%m')}" day_part = f"day={target_dt.strftime('%d')}" if clean_prefix: return f"{clean_prefix}/{year_part}/{month_part}/{day_part}/{filename}" return f"{year_part}/{month_part}/{day_part}/{filename}" class StorageManager: """Manages saving JSON, CSV, and Parquet to MinIO S3 or local fallback staging.""" def __init__(self, local_staging_dir: str = DEFAULT_LOCAL_STAGING): self.local_staging_dir = local_staging_dir os.makedirs(self.local_staging_dir, exist_ok=True) self.s3_client = get_minio_s3_client() def save_json(self, data: Any, bucket: str, key: str) -> Dict[str, Any]: """Save JSON payload to S3 MinIO or local staging.""" json_bytes = json.dumps(data, indent=2, default=str).encode("utf-8") s3_uri = f"s3://{bucket}/{key}" # 1. Try S3 upload if client is available uploaded_to_s3 = False if self.s3_client: try: self.s3_client.put_object( Bucket=bucket, Key=key, Body=json_bytes, ContentType="application/json" ) uploaded_to_s3 = True logger.info("Successfully uploaded JSON to %s", s3_uri) # Maintain canonical unpartitioned root file for downstream Silver tasks base_name = os.path.basename(key).split("_202")[0] canonical_key = f"{base_name}.json" if canonical_key != key: self.s3_client.put_object( Bucket=bucket, Key=canonical_key, Body=json_bytes, ContentType="application/json" ) # Also write canonical parquet version for tabular consumers if isinstance(data, dict): readings = data.get("readings") if readings and isinstance(readings, list): df_json = pd.DataFrame(readings) if "station_id" in data: df_json["stasiun_id"] = data["station_id"] elif "station_id" not in df_json.columns: df_json["stasiun_id"] = "SPKU_DKI_01" if "timestamp" in df_json.columns and "waktu" not in df_json.columns: df_json["waktu"] = df_json["timestamp"] pq_buf = io.BytesIO() df_json.to_parquet(pq_buf, index=False) self.s3_client.put_object( Bucket=bucket, Key=f"{base_name}.parquet", Body=pq_buf.getvalue(), ContentType="application/vnd.apache.parquet" ) except Exception as e: logger.warning("MinIO S3 upload failed for %s, falling back to local: %s", s3_uri, e) # 2. Local disk fallback local_path = os.path.join(self.local_staging_dir, bucket, key) os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, "wb") as f: f.write(json_bytes) return { "s3_uri": s3_uri, "local_path": local_path, "bytes_written": len(json_bytes), "uploaded_to_s3": uploaded_to_s3, "format": "json" } def save_dataframe( self, df: pd.DataFrame, bucket: str, key: str, target_format: str = "parquet" ) -> Dict[str, Any]: """Save pandas DataFrame to Parquet or CSV on S3 MinIO or local staging.""" buffer = io.BytesIO() actual_format = target_format.lower() content_type = "application/octet-stream" if actual_format == "parquet": try: df.to_parquet(buffer, index=False) content_type = "application/vnd.apache.parquet" except Exception as e: logger.warning("Parquet engine unavailable (%s), falling back to CSV", e) buffer = io.BytesIO() df.to_csv(buffer, index=False) actual_format = "csv" content_type = "text/csv" if not key.endswith(".csv"): key = key.rsplit(".", 1)[0] + ".csv" else: df.to_csv(buffer, index=False) actual_format = "csv" content_type = "text/csv" payload_bytes = buffer.getvalue() s3_uri = f"s3://{bucket}/{key}" uploaded_to_s3 = False if self.s3_client: try: self.s3_client.put_object( Bucket=bucket, Key=key, Body=payload_bytes, ContentType=content_type ) uploaded_to_s3 = True logger.info("Successfully uploaded %s to %s", actual_format.upper(), s3_uri) # Maintain canonical unpartitioned root file for downstream Silver tasks base_name = os.path.basename(key).split("_202")[0].replace("public_", "") canonical_key = f"{base_name}.{actual_format}" if canonical_key != key: self.s3_client.put_object( Bucket=bucket, Key=canonical_key, Body=payload_bytes, ContentType=content_type ) # If format is CSV, also write a Parquet copy for Iceberg / DuckDB consumers if actual_format == "csv": try: pq_buf = io.BytesIO() df.to_parquet(pq_buf, index=False) self.s3_client.put_object( Bucket=bucket, Key=f"{base_name}.parquet", Body=pq_buf.getvalue(), ContentType="application/vnd.apache.parquet" ) except Exception: pass except Exception as e: logger.warning("MinIO S3 upload failed for %s, falling back to local: %s", s3_uri, e) local_path = os.path.join(self.local_staging_dir, bucket, key) os.makedirs(os.path.dirname(local_path), exist_ok=True) with open(local_path, "wb") as f: f.write(payload_bytes) return { "s3_uri": s3_uri, "local_path": local_path, "rows_written": len(df), "bytes_written": len(payload_bytes), "uploaded_to_s3": uploaded_to_s3, "format": actual_format } # Global singleton instance storage_manager = StorageManager()