Reprocessing Dead-Letter Geotasks Safely

from __future__ import annotations

import json
from dataclasses import dataclass
from typing import Any, Optional, Sequence

import psycopg
import shapely
from prefect import flow, task, get_run_logger

ATTEMPT_CAP = 5

CLAIM_BATCH = """
UPDATE geometry_dlq
   SET attempts = attempts + 1, redriven_at = now()
 WHERE id IN (
       SELECT id FROM geometry_dlq
        WHERE resolved_at IS NULL
          AND attempts < %(cap)s
          AND (%(failure_class)s IS NULL OR failure_class = %(failure_class)s)
        ORDER BY failed_at
        LIMIT %(batch)s
        FOR UPDATE SKIP LOCKED)      -- two drains never take the same row
RETURNING id, work_key, source_uri, feature_id, geom_wkb,
          declared_crs, params, attempts
"""


@dataclass(frozen=True)
class RedriveItem:
    row_id: int
    work_key: str
    source_uri: str
    feature_id: Optional[str]
    geometry: Any
    declared_crs: str
    params: dict
    attempts: int


def claim_batch(
    conn: psycopg.Connection, batch: int = 200, failure_class: Optional[str] = None
) -> list[RedriveItem]:
    with conn.cursor() as cur:
        cur.execute(CLAIM_BATCH, {"cap": ATTEMPT_CAP, "batch": batch,
                                  "failure_class": failure_class})
        return [
            RedriveItem(
                row_id=r[0], work_key=r[1], source_uri=r[2], feature_id=r[3],
                # from_wkb, not from_wkt: the stored bytes are the exact geometry
                # that failed, down to the last coordinate digit.
                geometry=shapely.from_wkb(r[4]) if r[4] else None,
                declared_crs=r[5], params=r[6], attempts=r[7],
            )
            for r in cur.fetchall()
        ]


@task
def redrive_one(item: RedriveItem, conn: psycopg.Connection) -> bool:
    """Run the ORIGINAL task for one entry. Returns True when it finally succeeds."""
    logger = get_run_logger()
    try:
        # The same callable the pipeline uses — never a re-implementation. A second
        # code path is a second set of bugs, and it will drift.
        load_feature(
            geometry=item.geometry,
            declared_crs=item.declared_crs,
            work_key=item.work_key,          # same identity: the ledger de-duplicates
            **item.params,
        )
    except Exception as exc:
        logger.warning("redrive failed for %s (attempt %d): %s",
                       item.work_key, item.attempts, exc)
        with conn.cursor() as cur:
            if item.attempts >= ATTEMPT_CAP:
                cur.execute(
                    "UPDATE geometry_dlq SET resolved_at = now(), "
                    "error_text = %s WHERE id = %s",
                    (f"abandoned at attempt cap: {exc}"[:2000], item.row_id),
                )
        return False

    with conn.cursor() as cur:
        cur.execute("UPDATE geometry_dlq SET resolved_at = now() WHERE id = %s",
                    (item.row_id,))
    return True


@flow(name="drain-geometry-dlq")
def drain(batch: int = 200, failure_class: Optional[str] = None) -> dict[str, int]:
    with psycopg.connect(DSN) as conn:
        items = claim_batch(conn, batch=batch, failure_class=failure_class)
        results = [redrive_one(item, conn) for item in items]
        conn.commit()
    return {"claimed": len(items), "resolved": sum(results),
            "still_failing": len(items) - sum(results)}

def test_redrive_does_not_duplicate(conn, ledger, failing_feature) -> None:
    # First pass fails and dead-letters.
    with pytest.raises(GEOSException):
        load_feature(**failing_feature, work_key="wk-1")
    assert dlq_depth(conn) == 1

    # Repair the cause, then drain. The ledger's work_key stops a double write.
    install_geometry_repair_rule()
    outcome = drain(batch=10)
    assert outcome == {"claimed": 1, "resolved": 1, "still_failing": 0}
    assert feature_count(conn, "wk-1") == 1        # not 2

    # A second drain finds nothing: the row is resolved, not merely quiet.
    assert drain(batch=10)["claimed"] == 0


def test_cap_abandons_rather_than_looping(conn, poison_feature) -> None:
    for _ in range(ATTEMPT_CAP + 2):
        drain(batch=10)
    with conn.cursor() as cur:
        cur.execute("SELECT attempts, resolved_at, error_text FROM geometry_dlq")
        attempts, resolved_at, error_text = cur.fetchone()
    assert attempts == ATTEMPT_CAP
    assert resolved_at is not None
    assert "abandoned" in error_text

SELECT failure_class,
       count(*) FILTER (WHERE resolved_at IS NOT NULL) AS resolved,
       count(*)                                        AS total,
       round(100.0 * count(*) FILTER (WHERE resolved_at IS NOT NULL) / count(*), 1)
         AS pct_resolved
FROM   geometry_dlq
WHERE  redriven_at > now() - interval '1 day'
GROUP  BY failure_class
ORDER  BY 4;