dbx-main/app/modules/dispatch/db.py

"""dispatch_db PostgreSQL 저장소.

- 드라이버: psycopg 3 (`psycopg[binary,pool]`) — 다른 모듈과 동일 패턴.
- 연결 정보: 환경변수 `DISPATCH_DB_URL`
  (예: postgresql://dispatch_app:<pwd>@postgres-db:5432/dispatch_db)
- 스키마는 앱이 만들지 않는다. `scripts/sql/dispatch_db_init.sql` 을 superuser 가
  사전 적용한다. 앱 계정(dispatch_app)은 CRUD 권한만 받는다.
- 연결 풀은 lazy open — 부팅 시 DB 가 잠시 끊겨도 컨테이너가 죽지 않게.

박스 묶기/SKU 합산 등 순수 로직은 `store.py` 에 있고, 여기서는 DB I/O 와
조립만 담당한다.
"""

from __future__ import annotations

import logging
from datetime import date, datetime
from typing import Any

from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from app.timezone import KST

from . import store

logger = logging.getLogger("dispatch.db")


class DispatchStore:
    def __init__(self, dsn: str, *, min_size: int = 1, max_size: int = 5):
        self._pool = ConnectionPool(
            conninfo=dsn,
            min_size=min_size,
            max_size=max_size,
            kwargs={"row_factory": dict_row, "autocommit": True},
            open=False,
        )
        self._pool.open(wait=False)

    def close(self) -> None:
        self._pool.close()

    # ════════════════════════════════════════════════════════════
    # 배치
    # ════════════════════════════════════════════════════════════
    def create_batch(
        self,
        *,
        platform: str,
        dispatch_date: str,
        batch_name: str,
        picking_pdf_path: str = "",
        label_pdf_path: str = "",
        order_export_path: str = "",
    ) -> dict[str, Any]:
        if not (dispatch_date or "").strip():
            raise ValueError("출고 날짜는 필수입니다.")
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                INSERT INTO dispatch_batches
                    (platform, dispatch_date, batch_name,
                     picking_pdf_path, label_pdf_path, order_export_path)
                VALUES (%s,%s,%s,%s,%s,%s)
                RETURNING *
                """,
                (
                    (platform or "TikTok").strip(),
                    dispatch_date.strip(),
                    (batch_name or "").strip(),
                    (picking_pdf_path or "").strip(),
                    (label_pdf_path or "").strip(),
                    (order_export_path or "").strip(),
                ),
            ).fetchone()
        return self._serialize(row)

    def add_parcels(self, *, batch_id: int, parcels: list[dict[str, Any]]) -> int:
        """파서 결과(parcels)를 박스+상품으로 일괄 저장. 반환: 저장한 박스 수.

        한 트랜잭션. 같은 박스 안 같은 SKU 는 파서가 이미 합산해서 들어온다.
        """
        saved = 0
        with self._pool.connection() as conn:
            with conn.transaction():
                for p in parcels:
                    parcel_row = conn.execute(
                        """
                        INSERT INTO dispatch_parcels
                            (batch_id, seq, order_id, package_id,
                             tracking_id, shipping_provider,
                             recipient_name, recipient_phone, recipient_address)
                        VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s)
                        RETURNING id
                        """,
                        (
                            batch_id,
                            int(p.get("seq") or (saved + 1)),
                            (p.get("order_id") or "").strip(),
                            (p.get("package_id") or "").strip(),
                            (p.get("tracking_id") or "").strip(),
                            (p.get("shipping_provider") or "").strip(),
                            (p.get("recipient_name") or "").strip(),
                            (p.get("recipient_phone") or "").strip(),
                            (p.get("recipient_address") or "").strip(),
                        ),
                    ).fetchone()
                    parcel_id = parcel_row["id"]
                    for it in p.get("items", []):
                        conn.execute(
                            """
                            INSERT INTO dispatch_items
                                (parcel_id, seller_sku, product_name, quantity)
                            VALUES (%s,%s,%s,%s)
                            """,
                            (
                                parcel_id,
                                (it.get("seller_sku") or "").strip(),
                                (it.get("product_name") or "").strip(),
                                int(it.get("quantity") or 1),
                            ),
                        )
                    saved += 1
        return saved

    def list_batches(self, *, limit: int = 200) -> list[dict[str, Any]]:
        """배치 목록 + 박스 수/완료(택배스캔) 수 요약."""
        with self._pool.connection() as conn:
            rows = conn.execute(
                """
                SELECT b.*,
                       COUNT(p.id) AS parcel_count,
                       COUNT(p.id) FILTER (
                           WHERE p.label_attached AND p.handed_to_kagayaku
                       ) AS done_count
                FROM dispatch_batches b
                LEFT JOIN dispatch_parcels p ON p.batch_id = b.id
                GROUP BY b.id
                ORDER BY b.dispatch_date DESC, b.id DESC
                LIMIT %s
                """,
                (int(limit),),
            ).fetchall()
        return [self._serialize(r) for r in rows]

    def get_batch(self, *, batch_id: int) -> dict[str, Any] | None:
        with self._pool.connection() as conn:
            row = conn.execute(
                "SELECT * FROM dispatch_batches WHERE id = %s", (batch_id,)
            ).fetchone()
        return self._serialize(row) if row else None

    def delete_batch(self, *, batch_id: int) -> None:
        """배치 + 하위 박스/상품/로그 전체 삭제(CASCADE)."""
        with self._pool.connection() as conn:
            cur = conn.execute("DELETE FROM dispatch_batches WHERE id = %s", (batch_id,))
            if cur.rowcount == 0:
                raise KeyError(batch_id)

    # ════════════════════════════════════════════════════════════
    # 박스(parcel)
    # ════════════════════════════════════════════════════════════
    def list_parcels(self, *, batch_id: int) -> list[dict[str, Any]]:
        """배치의 박스 전체 + 각 박스의 상품 목록(seq 순)."""
        with self._pool.connection() as conn:
            parcel_rows = conn.execute(
                "SELECT * FROM dispatch_parcels WHERE batch_id = %s "
                "ORDER BY seq ASC, id ASC",
                (batch_id,),
            ).fetchall()
            item_rows = conn.execute(
                """
                SELECT i.* FROM dispatch_items i
                JOIN dispatch_parcels p ON p.id = i.parcel_id
                WHERE p.batch_id = %s
                ORDER BY i.seller_sku ASC, i.id ASC
                """,
                (batch_id,),
            ).fetchall()
        items_by_parcel: dict[int, list[dict[str, Any]]] = {}
        for r in item_rows:
            items_by_parcel.setdefault(r["parcel_id"], []).append(self._serialize(r))
        parcels: list[dict[str, Any]] = []
        for r in parcel_rows:
            p = self._serialize(r)
            p["items"] = items_by_parcel.get(r["id"], [])
            parcels.append(p)
        return parcels

    def sku_summary(self, *, batch_id: int) -> list[dict[str, Any]]:
        """배치 전체 SKU별 총 수량(피킹 요약). seller_sku 오름차순."""
        with self._pool.connection() as conn:
            rows = conn.execute(
                """
                SELECT i.seller_sku,
                       MAX(i.product_name)  AS product_name,
                       SUM(i.quantity)::int AS total_qty
                FROM dispatch_items i
                JOIN dispatch_parcels p ON p.id = i.parcel_id
                WHERE p.batch_id = %s
                GROUP BY i.seller_sku
                ORDER BY i.seller_sku ASC
                """,
                (batch_id,),
            ).fetchall()
        return [
            {
                "seller_sku": r["seller_sku"],
                "product_name": r["product_name"] or "",
                "total_qty": int(r["total_qty"] or 0),
            }
            for r in rows
        ]

    def toggle_status(
        self, *, parcel_id: int, field: str, worker_name: str = ""
    ) -> dict[str, Any]:
        """박스의 작업 상태 boolean 1개를 토글하고 로그를 남긴다.

        반환: {"parcel_id", "field", "value"(토글 후 값)}.
        field 는 store.STATUS_FIELDS 화이트리스트 검증(SQL 식별자 안전).
        """
        if field not in store.STATUS_FIELDS:
            raise ValueError(f"허용되지 않는 작업 상태: {field}")
        with self._pool.connection() as conn:
            with conn.transaction():
                cur = conn.execute(
                    # field 는 화이트리스트라 인젝션 위험 없음. RETURNING 으로 새 값 확인.
                    f"UPDATE dispatch_parcels SET {field} = NOT {field} "
                    "WHERE id = %s RETURNING " + field,
                    (parcel_id,),
                ).fetchone()
                if cur is None:
                    raise KeyError(parcel_id)
                new_value = bool(cur[field])
                conn.execute(
                    """
                    INSERT INTO dispatch_logs
                        (parcel_id, action, old_value, new_value, worker_name)
                    VALUES (%s,%s,%s,%s,%s)
                    """,
                    (
                        parcel_id,
                        field,
                        str(not new_value).lower(),
                        str(new_value).lower(),
                        (worker_name or "").lower().strip(),
                    ),
                )
        return {"parcel_id": parcel_id, "field": field, "value": new_value}

    def parcel_batch_id(self, *, parcel_id: int) -> int | None:
        with self._pool.connection() as conn:
            row = conn.execute(
                "SELECT batch_id FROM dispatch_parcels WHERE id = %s", (parcel_id,)
            ).fetchone()
        return int(row["batch_id"]) if row else None

    # ════════════════════════════════════════════════════════════
    # Kagayaku 전달 요약
    # ════════════════════════════════════════════════════════════
    def handover_summary(self, *, batch_id: int) -> dict[str, Any]:
        """배송사별 박스 수 + Tracking ID 목록(전달 리스트/인쇄용)."""
        with self._pool.connection() as conn:
            provider_rows = conn.execute(
                """
                SELECT COALESCE(NULLIF(shipping_provider, ''), '(미지정)') AS provider,
                       COUNT(*)::int AS box_count
                FROM dispatch_parcels
                WHERE batch_id = %s
                GROUP BY provider
                ORDER BY box_count DESC, provider ASC
                """,
                (batch_id,),
            ).fetchall()
            tracking_rows = conn.execute(
                """
                SELECT seq, order_id, tracking_id, shipping_provider
                FROM dispatch_parcels
                WHERE batch_id = %s
                ORDER BY seq ASC, id ASC
                """,
                (batch_id,),
            ).fetchall()
        total = sum(int(r["box_count"]) for r in provider_rows)
        return {
            "total_boxes": total,
            "by_provider": [
                {"provider": r["provider"], "box_count": int(r["box_count"])}
                for r in provider_rows
            ],
            "tracking_list": [self._serialize(r) for r in tracking_rows],
        }

    # ════════════════════════════════════════════════════════════
    # 직렬화
    # ════════════════════════════════════════════════════════════
    @staticmethod
    def _serialize(row: dict[str, Any] | None) -> dict[str, Any] | None:
        if row is None:
            return None
        out = dict(row)
        for k, v in list(out.items()):
            if isinstance(v, datetime):
                out[k] = v.astimezone(KST).isoformat(timespec="seconds")
            elif isinstance(v, date):
                out[k] = v.isoformat()
        for k in ("id", "batch_id", "parcel_id", "seq", "quantity",
                  "parcel_count", "done_count", "total_qty", "box_count"):
            if k in out and out[k] is not None:
                try:
                    out[k] = int(out[k])
                except (TypeError, ValueError):
                    pass
        for k in store.STATUS_FIELDS:
            if k in out:
                out[k] = bool(out[k])
        return out