213 lines
8.3 KiB
Python
213 lines
8.3 KiB
Python
|
"""
|
||
|
|
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()
|