o2o-site-AEO/solution/backend/crud/job_crud.py

"""작업 큐 CRUD — PostgreSQL 을 '제대로' 큐로 쓴다. (LPS `crud/job_crud.py` 이식)

 - 할당은 **단일 문장 원자 claim**: FOR UPDATE SKIP LOCKED 서브쿼리 + 같은 UPDATE + RETURNING.
   → 워커 컨테이너가 몇 개든 같은 잡 이중 할당이 원천 불가. fetch 와 claim 을 분리하지 않는다.
 - 모든 전이는 **조건부 CAS**(WHERE 에 status/worker_id 가드) + RETURNING.
 - 복구는 timeout 추측이 아니라 **lease 만료 소유권**(reaper 가 회수).
 - 재시도/백오프/dead-letter 를 큐에 내장.

전이가 조회/변경으로 나뉘지 않으므로(RETURNING) execute_lambda_write 로 실행한다.
큐 전이만 raw SQL 이다 — 다른 crud 는 전부 SQLAlchemy 표현식을 쓴다.
"""

import json

from sqlalchemy import text

from common.database.db_session_manager import DB_SESSION_MNG
from common.enums import DBType, DBWRType, JobStatus


# 잡 적재 시 워커를 즉시 깨우는 LISTEN/NOTIFY 채널(폴링 제거).
JOB_NOTIFY_CHANNEL = "web4ai_job"


def compute_backoff(attempts: int, base: float = 5.0, cap: float = 600.0) -> float:
    """지수 백오프(초). attempts 회 시도 후 다음 재시도까지 대기 = base * 2^(attempts-1), cap 상한."""
    return min(cap, base * (2 ** max(0, attempts - 1)))


class JobQueue:
    DB = DBType.MAIN.value

    async def _tx(self, fn):
        """쓰기 트랜잭션 — 값 반환이 필요한 큐 전이 전용 진입점."""
        return await DB_SESSION_MNG.execute_lambda_write(self.DB, fn)

    # ---- 적재 ----
    async def enqueue(
        self,
        job_type: int,
        payload: dict,
        priority: int = 100,
        dedupe_key: str | None = None,
        max_attempts: int = 3,
    ) -> str | None:
        """잡 적재. dedupe_key 가 활성(PENDING/RUNNING) 중복이면 삽입 없이 None 반환."""
        sql = text("""
            INSERT INTO job.jobs (job_type, priority, payload, dedupe_key, max_attempts)
            VALUES (:t, :p, CAST(:payload AS jsonb), :dk, :ma)
            ON CONFLICT (dedupe_key) WHERE status IN (1, 2) AND dedupe_key IS NOT NULL
            DO NOTHING
            RETURNING job_id
        """)

        async def run(s):
            row = (await s.execute(sql, {
                "t": job_type, "p": priority, "payload": json.dumps(payload, ensure_ascii=False),
                "dk": dedupe_key, "ma": max_attempts,
            })).first()
            if row:
                # 커밋 시 전달됨 → LISTEN 중인 유휴 워커를 즉시 깨운다(중복 스킵 시엔 알림 안 함).
                await s.execute(text("SELECT pg_notify(:ch, '')"), {"ch": JOB_NOTIFY_CHANNEL})
            return str(row[0]) if row else None

        return await self._tx(run)

    # ---- 원자적 claim ----
    async def claim(self, worker_id: str, lease_sec: int = 120) -> dict | None:
        """대기 잡 1건을 원자적으로 점유. 없으면 None.
        FOR UPDATE SKIP LOCKED 로 잠근 행을 같은 UPDATE 에서 RUNNING 으로 전이 → 이중 할당 불가."""
        sql = text("""
            UPDATE job.jobs SET
                status = 2,
                worker_id = :wid,
                lease_until = now() + make_interval(secs => :lease),
                run_started_at = now(),
                attempts = attempts + 1,
                updated_at = now()
            WHERE job_id = (
                SELECT job_id FROM job.jobs
                WHERE status = 1 AND run_after <= now()
                ORDER BY priority ASC, created_at ASC
                FOR UPDATE SKIP LOCKED
                LIMIT 1
            )
            RETURNING job_id, job_type, payload, attempts, max_attempts
        """)

        async def run(s):
            row = (await s.execute(sql, {"wid": worker_id, "lease": lease_sec})).mappings().first()
            if not row:
                return None
            d = dict(row)
            d["job_id"] = str(d["job_id"])
            if isinstance(d.get("payload"), str):
                d["payload"] = json.loads(d["payload"])
            return d

        return await self._tx(run)

    # ---- 완료/실패 (소유권 가드) ----
    async def complete(self, job_id: str, worker_id: str, result: dict | None = None) -> bool:
        sql = text("""
            UPDATE job.jobs SET status = 3, result = CAST(:result AS jsonb),
                                lease_until = NULL, worker_id = NULL, updated_at = now()
            WHERE job_id = CAST(:id AS uuid) AND status = 2 AND worker_id = :wid
            RETURNING job_id
        """)

        async def run(s):
            row = (await s.execute(sql, {
                "id": job_id, "wid": worker_id,
                "result": json.dumps(result, ensure_ascii=False) if result is not None else None,
            })).first()
            return row is not None

        return await self._tx(run)

    async def fail(self, job_id: str, worker_id: str, error: str, backoff_sec: float = 5.0) -> int | None:
        """실패 처리. 시도 남으면 PENDING(run_after=백오프)으로 재큐, 소진되면 DEAD(dead-letter).
        전이 후 status(JobStatus 값)를 반환. 소유 불일치면 None."""
        sql = text("""
            UPDATE job.jobs SET
                status = CASE WHEN attempts >= max_attempts THEN 4 ELSE 1 END,
                run_after = CASE WHEN attempts >= max_attempts THEN run_after
                                 ELSE now() + make_interval(secs => :backoff) END,
                last_error = :err,
                lease_until = NULL,
                worker_id = NULL,
                updated_at = now()
            WHERE job_id = CAST(:id AS uuid) AND status = 2 AND worker_id = :wid
            RETURNING status
        """)

        async def run(s):
            row = (await s.execute(sql, {
                "id": job_id, "wid": worker_id, "err": error[:2000], "backoff": backoff_sec,
            })).first()
            return int(row[0]) if row else None

        return await self._tx(run)

    # ---- lease 갱신(heartbeat) / 회수(reaper) ----
    async def renew_lease(self, job_id: str, worker_id: str, lease_sec: int = 120) -> bool:
        sql = text("""
            UPDATE job.jobs SET lease_until = now() + make_interval(secs => :lease), updated_at = now()
            WHERE job_id = CAST(:id AS uuid) AND worker_id = :wid AND status = 2
            RETURNING job_id
        """)

        async def run(s):
            row = (await s.execute(sql, {"id": job_id, "wid": worker_id, "lease": lease_sec})).first()
            return row is not None

        return await self._tx(run)

    async def reap(self) -> list[str]:
        """만료된 lease(워커 사망 등)의 RUNNING 잡을 회수. 시도 남으면 즉시 재큐, 소진되면 DEAD.
        회수된 job_id 목록 반환."""
        sql = text("""
            UPDATE job.jobs SET
                status = CASE WHEN attempts >= max_attempts THEN 4 ELSE 1 END,
                run_after = now(),
                last_error = COALESCE(last_error, '') || ' [lease-expired reclaim]',
                lease_until = NULL,
                worker_id = NULL,
                updated_at = now()
            WHERE status = 2 AND lease_until IS NOT NULL AND lease_until < now()
            RETURNING job_id
        """)

        async def run(s):
            rows = (await s.execute(sql)).all()
            return [str(r[0]) for r in rows]

        return await self._tx(run)

    # ---- 단건 조회 (상태 폴링) ----
    async def get(self, job_id: str) -> dict | None:
        """잡 단건 조회(읽기). 없으면 None. status 는 정수(JobStatus 값)."""
        sql = text("""
            SELECT job_id, job_type, status, priority, attempts, max_attempts,
                   payload, result, last_error, run_after, run_started_at, created_at, updated_at
            FROM job.jobs WHERE job_id = CAST(:id AS uuid)
        """)

        async def run(s):
            row = (await s.execute(sql, {"id": job_id})).mappings().first()
            if not row:
                return None
            d = dict(row)
            d["job_id"] = str(d["job_id"])
            for key in ("payload", "result"):
                if isinstance(d.get(key), str):
                    d[key] = json.loads(d[key])
            return d

        return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run)

    async def find_active(self, dedupe_key: str) -> dict | None:
        """dedupe_key 로 활성(PENDING/RUNNING) 잡을 찾는다.
        enqueue 가 중복으로 None 을 돌려줬을 때, 이미 돌고 있는 잡의 id 를 알려주기 위함."""
        sql = text("""
            SELECT job_id, job_type, status FROM job.jobs
            WHERE dedupe_key = :dk AND status IN (1, 2)
            LIMIT 1
        """)

        async def run(s):
            row = (await s.execute(sql, {"dk": dedupe_key})).mappings().first()
            if not row:
                return None
            d = dict(row)
            d["job_id"] = str(d["job_id"])
            return d

        return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run)

    # ---- 관측(관리 API/알림용) ----
    async def counts(self) -> dict[str, int]:
        """상태별 잡 개수."""
        async def run(s):
            rows = (await s.execute(text("SELECT status, count(*) FROM job.jobs GROUP BY status"))).all()
            by_val = {int(st): int(c) for st, c in rows}
            return {js.name: by_val.get(js.value, 0) for js in JobStatus}

        return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run)

    async def ops(self) -> dict:
        """운영 스냅샷(모니터링·알림용): 상태별 카운트 + 큐 지연(가장 오래된 PENDING 나이) +
        최근 1시간 DEAD + stuck(좀비 신호).

        stuck 은 두 축 — lease 만료(워커 사망인데 reaper 미회수) OR 실행 10분 초과(핸들러 행 —
        heartbeat 가 lease 를 계속 갱신해 lease 축엔 안 잡히므로 run_started_at 으로 따로 본다)."""
        sql = text("""
            SELECT
              count(*) FILTER (WHERE status = 1) AS pending,
              count(*) FILTER (WHERE status = 2) AS running,
              count(*) FILTER (WHERE status = 3) AS done,
              count(*) FILTER (WHERE status = 4) AS dead,
              count(*) FILTER (WHERE status = 4 AND updated_at > now() - interval '1 hour') AS dead_1h,
              count(*) FILTER (WHERE status = 2 AND (
                  (lease_until IS NOT NULL AND lease_until < now())
                  OR run_started_at < now() - interval '10 minutes'
              )) AS stuck_running,
              COALESCE(EXTRACT(EPOCH FROM (now() - min(created_at) FILTER (WHERE status = 1)))::int, 0) AS oldest_pending_sec,
              count(*) FILTER (WHERE last_error LIKE 'JobDeadlineExceeded%'
                  AND updated_at > now() - interval '1 hour') AS deadline_1h
            FROM job.jobs
        """)

        async def run(s):
            return {k: int(v) for k, v in dict((await s.execute(sql)).mappings().first()).items()}

        return await DB_SESSION_MNG.execute_lambda(self.DB, DBWRType.DB_READ.value, run)

    async def requeue(self, job_id: str) -> str | None:
        """DEAD 잡 재큐(관리자 액션): attempts 리셋 + PENDING 전이 + 워커 깨움.
        DEAD 가 아니거나 없으면 None. 같은 dedupe_key 의 활성 잡이 있으면 부분 유니크 위반."""
        sql = text("""
            UPDATE job.jobs SET status = 1, attempts = 0, run_after = now(),
                                lease_until = NULL, worker_id = NULL, run_started_at = NULL,
                                last_error = NULL, updated_at = now()
            WHERE job_id = CAST(:jid AS uuid) AND status = 4
            RETURNING job_id
        """)

        async def run(s):
            row = (await s.execute(sql, {"jid": job_id})).first()
            if row:
                await s.execute(text("SELECT pg_notify(:ch, '')"), {"ch": JOB_NOTIFY_CHANNEL})
            return str(row[0]) if row else None

        return await self._tx(run)