feat: initial commit for merica-pipeline delivery
This commit is contained in:
MERICA Admin
@@ -0,0 +1,212 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user