SupplyAgent/supply_infra/db/repositories/pipeline_outbox_repo.py

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any

from sqlalchemy import select

from supply_infra.db.models.pipeline_outbox import PipelineOutbox
from supply_infra.db.repositories.base import BaseRepository
from supply_infra.pipeline.enums import OutboxStatus


class PipelineOutboxRepository(BaseRepository[PipelineOutbox]):
    model = PipelineOutbox

    def get_by_idempotency_key(self, key: str) -> PipelineOutbox | None:
        stmt = select(PipelineOutbox).where(PipelineOutbox.idempotency_key == key)
        return self.session.scalar(stmt)

    def get(self, outbox_id: str, *, for_update: bool = False) -> PipelineOutbox | None:
        stmt = select(PipelineOutbox).where(PipelineOutbox.outbox_id == outbox_id)
        if for_update:
            stmt = stmt.with_for_update()
        return self.session.scalar(stmt)

    def enqueue(
        self,
        *,
        run_id: str,
        step_run_id: str,
        effect_type: str,
        idempotency_key: str,
        payload_hash: str,
        payload: dict[str, Any] | None,
        payload_uri: str | None = None,
        dry_run: bool = False,
    ) -> PipelineOutbox:
        existing = self.get_by_idempotency_key(idempotency_key)
        if existing is not None:
            if existing.payload_hash != payload_hash:
                raise ValueError(
                    "idempotency key reused with a different payload: "
                    f"{idempotency_key}"
                )
            return existing
        return self.add(
            PipelineOutbox(
                outbox_id=str(uuid.uuid4()),
                run_id=run_id,
                step_run_id=step_run_id,
                effect_type=effect_type,
                idempotency_key=idempotency_key,
                payload_hash=payload_hash,
                payload_json=payload,
                payload_uri=payload_uri,
                status=OutboxStatus.PENDING.value,
                dry_run=dry_run,
            )
        )

    def mark_sending(self, row: PipelineOutbox, *, now: datetime) -> None:
        row.status = OutboxStatus.SENDING.value
        row.attempt_count = int(row.attempt_count or 0) + 1
        row.last_attempt_at = now
        row.next_retry_at = None
        row.record_error = None
        self.session.flush()

    def mark_created(
        self,
        row: PipelineOutbox,
        *,
        external_id: str,
        external_name: str | None,
        response: dict[str, Any] | None,
    ) -> None:
        row.status = OutboxStatus.CREATED.value
        row.external_id = external_id
        row.external_name = external_name
        row.response_json = response
        row.record_error = None
        self.session.flush()

    def mark_retryable_failed(
        self,
        row: PipelineOutbox,
        *,
        error: str,
        response: dict[str, Any] | None = None,
    ) -> None:
        row.status = OutboxStatus.RETRYABLE_FAILED.value
        row.record_error = error
        if response is not None:
            row.response_json = response
        self.session.flush()

    def mark_ambiguous(self, row: PipelineOutbox, *, error: str) -> None:
        row.status = OutboxStatus.AMBIGUOUS.value
        row.record_error = error
        self.session.flush()

    def mark_succeeded(
        self,
        row: PipelineOutbox,
        *,
        response: dict[str, Any] | None = None,
    ) -> None:
        row.status = OutboxStatus.SUCCEEDED.value
        row.record_error = None
        if response is not None:
            row.response_json = response
        self.session.flush()

    def list_for_run(self, run_id: str) -> list[PipelineOutbox]:
        stmt = (
            select(PipelineOutbox)
            .where(PipelineOutbox.run_id == run_id)
            .order_by(PipelineOutbox.created_at)
        )
        return list(self.session.scalars(stmt).all())