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

"""expense_db PostgreSQL 저장소.

JSON 저장소(`ExpenseStore`)와 같은 인터페이스를 제공하여, 라우터 코드를
바꾸지 않고도 교체할 수 있다.

- 드라이버: psycopg 3 (`psycopg[binary,pool]`)
- 연결 정보: 환경변수 `EXPENSE_DB_URL` (예: postgresql://user:pwd@host:5432/expense_db)
- 스키마: `scripts/sql/expense_db_init.sql` 로 사전 초기화한다. 본 클래스는
  앱 부팅 시 `CREATE TABLE IF NOT EXISTS`로 보강만 한다.
"""

from __future__ import annotations

import uuid
from datetime import date, datetime, timezone
from typing import Any

from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool

from .store import CATEGORIES, METHODS, STATUSES

_DDL = """
CREATE TABLE IF NOT EXISTS expense_items (
    id           TEXT        PRIMARY KEY,
    owner        TEXT        NOT NULL,
    spent_at     DATE        NOT NULL,
    category     TEXT        NOT NULL,
    method       TEXT        NOT NULL,
    merchant     TEXT        NOT NULL DEFAULT '',
    amount       BIGINT      NOT NULL DEFAULT 0 CHECK (amount >= 0),
    memo         TEXT        NOT NULL DEFAULT '',
    status       TEXT        NOT NULL DEFAULT '작성중',
    created_at   TIMESTAMPTZ NOT NULL DEFAULT now(),
    updated_at   TIMESTAMPTZ NOT NULL DEFAULT now()
);
CREATE INDEX IF NOT EXISTS idx_expense_owner_spent ON expense_items (owner, spent_at DESC);
CREATE INDEX IF NOT EXISTS idx_expense_status      ON expense_items (status);
"""


class ExpenseDBStore:
    """`ExpenseStore` 와 동일한 메서드 시그니처."""

    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=True,
        )
        self._ensure_schema()

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

    def _ensure_schema(self) -> None:
        with self._pool.connection() as conn:
            conn.execute(_DDL)

    # ── 조회 ──
    def list_for(self, email: str) -> list[dict[str, Any]]:
        email = email.lower().strip()
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM expense_items WHERE owner = %s "
                "ORDER BY spent_at DESC, created_at DESC",
                (email,),
            ).fetchall()
        return [self._serialize(r) for r in rows]

    def list_all(self) -> list[dict[str, Any]]:
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM expense_items ORDER BY created_at DESC"
            ).fetchall()
        return [self._serialize(r) for r in rows]

    def get(self, *, item_id: str, owner: str) -> dict[str, Any] | None:
        owner = owner.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                "SELECT * FROM expense_items WHERE id = %s AND owner = %s",
                (item_id, owner),
            ).fetchone()
        return self._serialize(row) if row else None

    # ── 변경 ──
    def create(self, *, owner: str, payload: dict[str, Any]) -> dict[str, Any]:
        owner = owner.lower().strip()
        norm = self._normalize(payload)
        if not norm["spent_at"]:
            raise ValueError("spent_at 필수")
        status = (payload.get("status") or "작성중").strip()
        item_id = uuid.uuid4().hex[:12]
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                INSERT INTO expense_items
                    (id, owner, spent_at, category, method, merchant, amount, memo, status)
                VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
                RETURNING *
                """,
                (
                    item_id,
                    owner,
                    norm["spent_at"],
                    norm["category"],
                    norm["method"],
                    norm["merchant"],
                    norm["amount"],
                    norm["memo"],
                    status,
                ),
            ).fetchone()
        return self._serialize(row)

    def update(
        self, *, item_id: str, owner: str, payload: dict[str, Any]
    ) -> dict[str, Any]:
        owner = owner.lower().strip()
        norm = self._normalize(payload)
        status = payload.get("status")
        with self._pool.connection() as conn:
            if status:
                row = conn.execute(
                    """
                    UPDATE expense_items
                       SET spent_at = %s, category = %s, method = %s,
                           merchant = %s, amount = %s, memo = %s, status = %s
                     WHERE id = %s AND owner = %s
                     RETURNING *
                    """,
                    (
                        norm["spent_at"] or None,
                        norm["category"],
                        norm["method"],
                        norm["merchant"],
                        norm["amount"],
                        norm["memo"],
                        status,
                        item_id,
                        owner,
                    ),
                ).fetchone()
            else:
                row = conn.execute(
                    """
                    UPDATE expense_items
                       SET spent_at = %s, category = %s, method = %s,
                           merchant = %s, amount = %s, memo = %s
                     WHERE id = %s AND owner = %s
                     RETURNING *
                    """,
                    (
                        norm["spent_at"] or None,
                        norm["category"],
                        norm["method"],
                        norm["merchant"],
                        norm["amount"],
                        norm["memo"],
                        item_id,
                        owner,
                    ),
                ).fetchone()
        if not row:
            raise KeyError(item_id)
        return self._serialize(row)

    def delete(self, *, item_id: str, owner: str) -> None:
        owner = owner.lower().strip()
        with self._pool.connection() as conn:
            cur = conn.execute(
                "DELETE FROM expense_items WHERE id = %s AND owner = %s",
                (item_id, owner),
            )
            if cur.rowcount == 0:
                raise KeyError(item_id)

    # ── 워크플로 ──
    def submit(self, *, item_id: str, owner: str) -> dict[str, Any]:
        return self._transition_owner(
            item_id=item_id, owner=owner, from_status="작성중", to_status="제출"
        )

    def revert_to_draft(self, *, item_id: str, owner: str) -> dict[str, Any]:
        """반려 또는 제출 상태에서 본인이 작성중으로 되돌림."""
        owner = owner.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                UPDATE expense_items
                   SET status = '작성중',
                       reject_reason = NULL,
                       decided_at = NULL,
                       approver_email = NULL
                 WHERE id = %s AND owner = %s
                   AND status IN ('제출', '반려')
                 RETURNING *
                """,
                (item_id, owner),
            ).fetchone()
        if not row:
            raise ValueError("작성중 으로 되돌릴 수 없는 상태입니다.")
        return self._serialize(row)

    def approve(self, *, item_id: str, approver_email: str) -> dict[str, Any]:
        return self._transition_approver(
            item_id=item_id,
            approver_email=approver_email,
            from_statuses=("제출",),
            to_status="승인",
        )

    def reject(
        self, *, item_id: str, approver_email: str, reason: str
    ) -> dict[str, Any]:
        if not reason.strip():
            raise ValueError("반려 사유 필수")
        approver_email = approver_email.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                UPDATE expense_items
                   SET status = '반려',
                       approver_email = %s,
                       decided_at = now(),
                       reject_reason = %s
                 WHERE id = %s AND status = '제출'
                 RETURNING *
                """,
                (approver_email, reason.strip(), item_id),
            ).fetchone()
        if not row:
            raise ValueError("제출 상태가 아니거나 항목 없음")
        return self._serialize(row)

    def settle(self, *, item_id: str, approver_email: str) -> dict[str, Any]:
        return self._transition_approver(
            item_id=item_id,
            approver_email=approver_email,
            from_statuses=("승인",),
            to_status="정산완료",
        )

    def _transition_owner(
        self, *, item_id: str, owner: str, from_status: str, to_status: str
    ) -> dict[str, Any]:
        owner = owner.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                UPDATE expense_items
                   SET status = %s
                 WHERE id = %s AND owner = %s AND status = %s
                 RETURNING *
                """,
                (to_status, item_id, owner, from_status),
            ).fetchone()
        if not row:
            raise ValueError(f"전이 불가: {from_status} → {to_status}")
        return self._serialize(row)

    def _transition_approver(
        self,
        *,
        item_id: str,
        approver_email: str,
        from_statuses: tuple[str, ...],
        to_status: str,
    ) -> dict[str, Any]:
        approver_email = approver_email.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                UPDATE expense_items
                   SET status = %s,
                       approver_email = %s,
                       decided_at = now()
                 WHERE id = %s AND status = ANY(%s)
                 RETURNING *
                """,
                (to_status, approver_email, item_id, list(from_statuses)),
            ).fetchone()
        if not row:
            raise ValueError(f"전이 불가 → {to_status}")
        return self._serialize(row)

    # ── 승인자 대기열 ──
    def list_pending_approval(self) -> list[dict[str, Any]]:
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM expense_items WHERE status = '제출' "
                "ORDER BY created_at ASC"
            ).fetchall()
        return [self._serialize(r) for r in rows]

    def get_any(self, *, item_id: str) -> dict[str, Any] | None:
        """승인자/관리자용 — owner 무시하고 단일 조회."""
        with self._pool.connection() as conn:
            row = conn.execute(
                "SELECT * FROM expense_items WHERE id = %s", (item_id,)
            ).fetchone()
        return self._serialize(row) if row else None

    # ── 첨부 ──
    def add_attachment(
        self,
        *,
        item_id: str,
        owner: str,
        kind: str,
        filename: str,
        stored_path: str,
        content_type: str,
        size_bytes: int,
    ) -> dict[str, Any]:
        if kind not in ("receipt", "other"):
            raise ValueError("kind 는 receipt|other")
        owner = owner.lower().strip()
        att_id = uuid.uuid4().hex[:12]
        with self._pool.connection() as conn:
            row = conn.execute(
                """
                INSERT INTO expense_attachments
                    (id, item_id, owner, kind, filename, stored_path,
                     content_type, size_bytes)
                VALUES (%s,%s,%s,%s,%s,%s,%s,%s)
                RETURNING *
                """,
                (
                    att_id,
                    item_id,
                    owner,
                    kind,
                    filename,
                    stored_path,
                    content_type,
                    size_bytes,
                ),
            ).fetchone()
        return self._att_serialize(row)

    def list_attachments(self, *, item_id: str) -> list[dict[str, Any]]:
        with self._pool.connection() as conn:
            rows = conn.execute(
                "SELECT * FROM expense_attachments WHERE item_id = %s "
                "ORDER BY uploaded_at ASC",
                (item_id,),
            ).fetchall()
        return [self._att_serialize(r) for r in rows]

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

    def delete_attachment(self, *, att_id: str, owner: str) -> dict[str, Any]:
        """삭제된 행 반환 (파일 정리용 stored_path 포함)."""
        owner = owner.lower().strip()
        with self._pool.connection() as conn:
            row = conn.execute(
                "DELETE FROM expense_attachments "
                "WHERE id = %s AND owner = %s RETURNING *",
                (att_id, owner),
            ).fetchone()
        if not row:
            raise KeyError(att_id)
        return self._att_serialize(row)

    @staticmethod
    def _att_serialize(row: dict[str, Any] | None) -> dict[str, Any] | None:
        if row is None:
            return None
        out = dict(row)
        v = out.get("uploaded_at")
        if isinstance(v, datetime):
            out["uploaded_at"] = v.astimezone(timezone.utc).isoformat(timespec="seconds")
        out["size_bytes"] = int(out.get("size_bytes", 0))
        return out

    # ── 집계 (월별) ──
    def monthly_summary(
        self, *, email: str, year: int
    ) -> list[dict[str, Any]]:
        email = email.lower().strip()
        with self._pool.connection() as conn:
            rows = conn.execute(
                """
                SELECT to_char(spent_at, 'YYYY-MM') AS month,
                       category,
                       COUNT(*) AS cnt,
                       COALESCE(SUM(amount), 0) AS total
                  FROM expense_items
                 WHERE owner = %s AND EXTRACT(YEAR FROM spent_at) = %s
                 GROUP BY 1, 2
                 ORDER BY 1, 2
                """,
                (email, year),
            ).fetchall()
        return [
            {
                "month": r["month"],
                "category": r["category"],
                "count": int(r["cnt"]),
                "total": int(r["total"]),
            }
            for r in rows
        ]

    def list_for_export(
        self,
        *,
        email: str | None,
        date_from: str | None = None,
        date_to: str | None = None,
    ) -> list[dict[str, Any]]:
        """엑셀 내보내기용. email=None 이면 전체 (승인자/관리자용)."""
        clauses = []
        params: list[Any] = []
        if email:
            clauses.append("owner = %s")
            params.append(email.lower().strip())
        if date_from:
            clauses.append("spent_at >= %s")
            params.append(date_from)
        if date_to:
            clauses.append("spent_at <= %s")
            params.append(date_to)
        where = ("WHERE " + " AND ".join(clauses)) if clauses else ""
        with self._pool.connection() as conn:
            rows = conn.execute(
                f"SELECT * FROM expense_items {where} "
                f"ORDER BY spent_at ASC, created_at ASC",
                params,
            ).fetchall()
        return [self._serialize(r) for r in rows]

    # ── 요약 ──
    def summary_for(self, email: str) -> dict[str, Any]:
        email = email.lower().strip()
        with self._pool.connection() as conn:
            head = conn.execute(
                "SELECT COUNT(*) AS count, COALESCE(SUM(amount), 0) AS total "
                "FROM expense_items WHERE owner = %s",
                (email,),
            ).fetchone()
            status_rows = conn.execute(
                "SELECT status, COUNT(*) AS c FROM expense_items "
                "WHERE owner = %s GROUP BY status",
                (email,),
            ).fetchall()
            cat_rows = conn.execute(
                "SELECT category, COALESCE(SUM(amount), 0) AS s "
                "FROM expense_items WHERE owner = %s GROUP BY category",
                (email,),
            ).fetchall()
        by_status = {s: 0 for s in STATUSES}
        for r in status_rows:
            by_status[r["status"]] = int(r["c"])
        by_category = {c: 0 for c in CATEGORIES}
        for r in cat_rows:
            by_category[r["category"]] = int(r["s"])
        return {
            "count": int(head["count"]) if head else 0,
            "total": int(head["total"]) if head else 0,
            "by_status": by_status,
            "by_category": by_category,
        }

    # ── helpers ──
    @staticmethod
    def _serialize(row: dict[str, Any] | None) -> dict[str, Any] | None:
        if row is None:
            return None
        out = dict(row)
        if isinstance(out.get("spent_at"), date):
            out["spent_at"] = out["spent_at"].isoformat()
        for k in ("created_at", "updated_at"):
            v = out.get(k)
            if isinstance(v, datetime):
                out[k] = v.astimezone(timezone.utc).isoformat(timespec="seconds")
        out["amount"] = int(out.get("amount", 0))
        return out

    @staticmethod
    def _normalize(
        payload: dict[str, Any], base: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        b = dict(base or {})
        b["spent_at"] = str(payload.get("spent_at") or b.get("spent_at") or "").strip()
        category = str(payload.get("category") or b.get("category") or "기타").strip()
        method = str(payload.get("method") or b.get("method") or "법인카드").strip()
        b["category"] = category if category in CATEGORIES else "기타"
        b["method"] = method if method in METHODS else "법인카드"
        b["merchant"] = str(payload.get("merchant") or b.get("merchant") or "").strip()
        try:
            b["amount"] = max(0, int(payload.get("amount") or b.get("amount") or 0))
        except (TypeError, ValueError):
            b["amount"] = 0
        b["memo"] = str(payload.get("memo") or b.get("memo") or "").strip()
        return b