310-a — Pod Safety + Atomic Claim¶

# Claiming
locked_by:  Optional[str]      = Field(None)  # set atomically on claim
locked_at:  Optional[datetime] = Field(None)  # most recent claim time (updates on retry)

# Retry scaffolding (used live by Plan A's poller; max_attempts honored by claim_job)
attempts:     int               = Field(0)
max_attempts: int               = Field(3)
run_after:    Optional[datetime] = Field(None)  # NULL = run immediately

# Grouping / dispatch
pipeline_id: Optional[uuid.UUID] = Field(None)  # groups jobs in a multi-step run
job_type:    Optional[str]        = Field(None)  # "csv_ingest", "factor_ingest",
                                                  # "emission_recalc", "unit_sync", etc.

ALTER TABLE data_ingestion_jobs
  ADD COLUMN locked_by    VARCHAR(255),
  ADD COLUMN locked_at    TIMESTAMPTZ,
  ADD COLUMN attempts     INTEGER NOT NULL DEFAULT 0,
  ADD COLUMN max_attempts INTEGER NOT NULL DEFAULT 3,
  ADD COLUMN run_after    TIMESTAMPTZ,
  ADD COLUMN pipeline_id  UUID,
  ADD COLUMN job_type     VARCHAR(100);

CREATE INDEX ix_data_ingestion_jobs_pending
    ON data_ingestion_jobs (run_after)
    WHERE state = 0      -- NOT_STARTED
      AND locked_by IS NULL;   -- supports the safety poller

# backend/app/tasks/_pod_id.py
import os, socket
POD_ID: str = os.environ.get("HOSTNAME", socket.gethostname())

# backend/app/repositories/data_ingestion.py

async def claim_job(self, job_id: int, pod_id: str) -> bool:
    """
    Atomically claim a job for execution.

    Returns True if this caller claimed it; False if another pod beat us, the job
    is no longer eligible, or attempts >= max_attempts.

    Two-step transaction:
      1. Unset is_current on previous current row (if any) for the same
         (module_type_id, data_entry_type_id, target_type, ingestion_method, year)
         combo, EXCEPT the row we're claiming.
      2. Atomic UPDATE on the target row. The partial unique index trips here
         if a concurrent claimer already set is_current=TRUE on a different row
         for the same combo.
    """
    job = await self.get_job_by_id(job_id)
    if job is None:
        return False

    # Step 1: clear previous is_current for this combo
    where = and_(
        col(DataIngestionJob.is_current),
        col(DataIngestionJob.id) != job_id,
        col(DataIngestionJob.target_type) == job.target_type,
        col(DataIngestionJob.ingestion_method) == job.ingestion_method,
        col(DataIngestionJob.year) == job.year,
    )
    if job.module_type_id is None:
        where = and_(where, col(DataIngestionJob.module_type_id).is_(None))
    else:
        where = and_(where, col(DataIngestionJob.module_type_id) == job.module_type_id)
    if job.data_entry_type_id is None:
        where = and_(where, col(DataIngestionJob.data_entry_type_id).is_(None))
    else:
        where = and_(
            where, col(DataIngestionJob.data_entry_type_id) == job.data_entry_type_id
        )
    await self.session.execute(
        update(DataIngestionJob).where(where).values(is_current=False)
    )

    # Step 2: atomic claim. Unique index trips on race.
    try:
        result = await self.session.execute(
            update(DataIngestionJob)
            .where(
                col(DataIngestionJob.id) == job_id,
                col(DataIngestionJob.state).in_(
                    [IngestionState.NOT_STARTED, IngestionState.QUEUED]
                ),
                col(DataIngestionJob.locked_by).is_(None),
                col(DataIngestionJob.attempts) < col(DataIngestionJob.max_attempts),
            )
            .values(
                locked_by=pod_id,
                locked_at=func.now(),
                state=IngestionState.RUNNING,
                is_current=True,
                attempts=col(DataIngestionJob.attempts) + 1,
            )
            .returning(DataIngestionJob.id)
        )
        await self.session.commit()
        return result.scalar_one_or_none() is not None
    except IntegrityError:
        # Unique-index trip: another pod claimed for this combo
        await self.session.rollback()
        return False

await job_repo.update_ingestion_job(state=IngestionState.RUNNING, ...)
await job_repo.mark_job_as_current(job)
await job_session.commit()

claimed = await job_repo.claim_job(job_id, POD_ID)
if not claimed:
    logger.info(f"Job {job_id} already claimed by another pod or not eligible")
    return

# Resets a job stuck in RUNNING after a pod crash.
# Allowed only if locked_at < now() - settings.STALE_JOB_TIMEOUT.
UPDATE data_ingestion_jobs
   SET state     = 'NOT_STARTED',
       locked_by = NULL,
       locked_at = NULL,
       attempts  = 0,         -- give it a fresh quota
       run_after = now()
 WHERE id = :job_id
   AND state = 'RUNNING'
   AND locked_at < now() - :stale_timeout
RETURNING id;

# backend/app/tasks/_poller.py
import asyncio
from app.core.config import settings
from app.db import SessionLocal
from app.models.data_ingestion import DataIngestionJob, IngestionState
from app.repositories.data_ingestion import DataIngestionRepository
from app.tasks._pod_id import POD_ID
# Plan C: from app.tasks.runner import run_job
# In Plan A, run_job is the legacy dispatch (run_ingestion / run_recalculation / etc.)
# selected by job_type. Plan C consolidates this.

POLL_INTERVAL_SECONDS = 10

async def poll_pending_jobs() -> None:
    """Pick up jobs that were created but never scheduled (e.g. crashed pod)."""
    while True:
        try:
            async with SessionLocal() as session:
                stmt = (
                    select(DataIngestionJob)
                    .where(
                        DataIngestionJob.state == IngestionState.NOT_STARTED,
                        or_(
                            DataIngestionJob.run_after.is_(None),
                            DataIngestionJob.run_after <= func.now(),
                        ),
                        DataIngestionJob.locked_by.is_(None),
                        DataIngestionJob.attempts < DataIngestionJob.max_attempts,
                    )
                    .with_for_update(skip_locked=True)
                    .limit(10)
                )
                jobs = (await session.execute(stmt)).scalars().all()
                for job in jobs:
                    asyncio.create_task(run_job(job.id))
        except Exception as exc:
            logger.warning(f"Poller iteration failed: {exc}", exc_info=True)
        await asyncio.sleep(POLL_INTERVAL_SECONDS)


def start_poller(app: FastAPI) -> None:
    @app.on_event("startup")
    async def _startup():
        if settings.RUN_BACKGROUND_POLLER:
            app.state.poller_task = asyncio.create_task(poll_pending_jobs())

    @app.on_event("shutdown")
    async def _shutdown():
        task = getattr(app.state, "poller_task", None)
        if task:
            task.cancel()

# backend/app/core/config.py
class Settings(BaseSettings):
    ...
    STALE_JOB_TIMEOUT_MINUTES: int = 30
    RUN_BACKGROUND_POLLER: bool = True