src/customer_intelligence/etl/load.py

"""Load transformed DataFrames into staging tables and run warehouse transforms."""

from datetime import date, timedelta
from pathlib import Path

import pandas as pd
from sqlalchemy import Engine, text
from tqdm import tqdm

STAGING_TABLE_MAP = {
    "orders": "staging.stg_orders",
    "order_items": "staging.stg_order_items",
    "order_payments": "staging.stg_order_payments",
    "order_reviews": "staging.stg_order_reviews",
    "customers": "staging.stg_customers",
    "sellers": "staging.stg_sellers",
    "products": "staging.stg_products",
    "category_translations": "staging.stg_product_category_translations",
    "geolocation": "staging.stg_geolocation",
}

WAREHOUSE_SQL = Path(__file__).parent / "warehouse_transform.sql"


class OlistLoader:
    def __init__(self, engine: Engine) -> None:
        self.engine = engine

    def truncate_and_load_staging(
        self, df: pd.DataFrame, table: str, chunksize: int = 10_000
    ) -> int:
        """TRUNCATE the staging table then bulk-insert the DataFrame."""
        with self.engine.begin() as conn:
            conn.execute(text(f"TRUNCATE TABLE {table}"))
        # psycopg3 caps bind parameters at 65535; stay safely under
        ncols = len(df.columns)
        safe_chunk = min(chunksize, max(1, 65535 // ncols))
        df.to_sql(
            table.split(".")[-1],
            self.engine,
            schema=table.split(".")[0],
            if_exists="append",
            index=False,
            method="multi",
            chunksize=safe_chunk,
        )
        return len(df)

    def load_all_staging(self, data: dict[str, pd.DataFrame]) -> dict[str, int]:
        counts = {}
        for key, df in tqdm(data.items(), desc="Loading staging tables"):
            table = STAGING_TABLE_MAP[key]
            n = self.truncate_and_load_staging(df, table)
            counts[key] = n
            print(f"  {table}: {n:,} rows")
        return counts

    def populate_dim_dates(self, start: date, end: date) -> int:
        """Generate and insert all dates in [start, end] into warehouse.dim_dates."""
        rows = []
        current = start
        while current <= end:
            rows.append(
                {
                    "date_id": current,
                    "year": current.year,
                    "quarter": (current.month - 1) // 3 + 1,
                    "month": current.month,
                    "month_name": current.strftime("%B"),
                    "week": current.isocalendar()[1],
                    "day_of_week": current.weekday(),
                    "day_name": current.strftime("%A"),
                    "is_weekend": current.weekday() >= 5,
                    "is_month_end": (current + timedelta(days=1)).month != current.month,
                }
            )
            current += timedelta(days=1)

        df = pd.DataFrame(rows)
        df.to_sql(
            "dim_dates",
            self.engine,
            schema="warehouse",
            if_exists="append",
            index=False,
            method="multi",
            chunksize=1000,
        )
        # Remove any duplicates that sneak in on re-runs
        with self.engine.begin() as conn:
            conn.execute(
                text(
                    "DELETE FROM warehouse.dim_dates a USING warehouse.dim_dates b "
                    "WHERE a.ctid < b.ctid AND a.date_id = b.date_id"
                )
            )
        return len(rows)

    def run_warehouse_transform(self) -> None:
        """Execute staging → warehouse SQL transforms."""
        sql = WAREHOUSE_SQL.read_text()
        with self.engine.begin() as conn:
            for stmt in sql.split(";"):
                stmt = stmt.strip()
                if stmt:
                    conn.execute(text(stmt))

    def run_data_quality_checks(self) -> dict[str, bool]:
        checks = {}
        with self.engine.connect() as conn:
            # All staging tables non-empty
            for key, table in STAGING_TABLE_MAP.items():
                result = conn.execute(text(f"SELECT COUNT(*) FROM {table}"))
                count = result.scalar()
                checks[f"{key}_staging_non_empty"] = count > 0

            # Warehouse tables populated
            for tbl in ["dim_customers", "dim_products", "dim_sellers", "fact_orders"]:
                result = conn.execute(
                    text(f"SELECT COUNT(*) FROM warehouse.{tbl}")
                )
                count = result.scalar()
                checks[f"{tbl}_non_empty"] = count > 0

            # fact_orders has no orphan customer_key
            result = conn.execute(
                text(
                    "SELECT COUNT(*) FROM warehouse.fact_orders f "
                    "LEFT JOIN warehouse.dim_customers c ON f.customer_key = c.customer_key "
                    "WHERE c.customer_key IS NULL"
                )
            )
            checks["no_orphan_customer_keys"] = result.scalar() == 0

        return checks