fix(prune): add pg prune

Author: kshitijrajsharma

infra/run-artifact-refresh.sh

@@ -11,7 +11,8 @@ dsn="${OSMSG_PSQL_DSN:-postgresql://osmsg:osmsg@db:5432/osmsg}"
 repo="${OSMSG_HISTORY_REPO:-kshitijrajsharma/osmsg-history}"
 
 echo "[artifact-refresh] advancing ${artifact_host} from ${repo}"
-docker compose run --rm -e HF_XET_HIGH_PERFORMANCE=1 -v "${artifact_host}:/artifact" --entrypoint osmsg worker \
+# Cap Xet download concurrency so its reconstruction buffers stay well within the container memory limit.
+docker compose run --rm -e HF_XET_CLIENT_AC_MAX_DOWNLOAD_CONCURRENCY=4 -v "${artifact_host}:/artifact" --entrypoint osmsg worker \
   maintain refresh --artifact-dir /artifact --repo "${repo}"
 
 echo "[artifact-refresh] reloading the API onto the advanced frontier"

osmsg/prune.py

@@ -11,23 +11,45 @@
 DEFAULT_OVERLAP = dt.timedelta(days=2)
 
 
-def prune_pg(dsn: str, cutoff: dt.datetime) -> tuple[int, int]:
-    """Delete changesets and their changeset_stats older than cutoff; child rows first for the FK. The
-    DSN is interpolated into ATTACH, so it must be trusted."""
+def _attach(dsn: str) -> duckdb.DuckDBPyConnection:
     conn = duckdb.connect()
     conn.execute("INSTALL postgres")
     conn.execute("LOAD postgres")
     conn.execute(f"ATTACH '{dsn.replace(chr(39), chr(39) * 2)}' AS pg (TYPE postgres)")
+    return conn
+
+
+def _pg_execute(conn: duckdb.DuckDBPyConnection, sql: str) -> None:
+    """Run one statement natively on the attached Postgres, so a bulk DELETE is a single indexed
+    statement server-side instead of DuckDB's per-row ctid batches."""
+    conn.execute(f"CALL postgres_execute('pg', $osmsg_stmt${sql}$osmsg_stmt$)")
+
+
+def prune_pg(dsn: str, cutoff: dt.datetime) -> tuple[int, int]:
+    """Delete changesets and their changeset_stats older than cutoff; child rows first for the FK. The
+    DSN is interpolated into ATTACH, so it must be trusted. Counting and deleting use separate
+    connections because a read pins the connection read-only, which would block the native deletes."""
     iso = cutoff.astimezone(dt.UTC).isoformat()
-    old_cs = f"SELECT changeset_id FROM pg.changesets WHERE created_at < TIMESTAMPTZ '{iso}'"
-    stats_row = conn.execute(f"SELECT count(*) FROM pg.changeset_stats WHERE changeset_id IN ({old_cs})").fetchone()
-    cs_row = conn.execute(f"SELECT count(*) FROM pg.changesets WHERE created_at < TIMESTAMPTZ '{iso}'").fetchone()
+    older = f"created_at < TIMESTAMPTZ '{iso}'"
+
+    reader = _attach(dsn)
+    stats_row = reader.execute(
+        "SELECT count(*) FROM pg.changeset_stats s "
+        f"WHERE EXISTS (SELECT 1 FROM pg.changesets c WHERE c.changeset_id = s.changeset_id AND c.{older})"
+    ).fetchone()
+    cs_row = reader.execute(f"SELECT count(*) FROM pg.changesets WHERE {older}").fetchone()
+    reader.close()
     stats_n = stats_row[0] if stats_row else 0
     cs_n = cs_row[0] if cs_row else 0
+
     if cs_n:
-        conn.execute(f"DELETE FROM pg.changeset_stats WHERE changeset_id IN ({old_cs})")
-        conn.execute(f"DELETE FROM pg.changesets WHERE created_at < TIMESTAMPTZ '{iso}'")
-    conn.close()
+        writer = _attach(dsn)
+        _pg_execute(
+            writer,
+            f"DELETE FROM changeset_stats s USING changesets c WHERE s.changeset_id = c.changeset_id AND c.{older}",
+        )
+        _pg_execute(writer, f"DELETE FROM changesets WHERE {older}")
+        writer.close()
     return stats_n, cs_n
 
 

tests/test_prune.py

@@ -32,3 +32,31 @@ def test_prune_covered_raises_without_manifest(monkeypatch):
     monkeypatch.setattr(prune, "fetch_manifest", lambda url: None)
     with pytest.raises(OsmsgError):
         prune.prune_covered("dsn", "url")
+
+
+class _FakeConn:
+    def __init__(self, calls):
+        self._calls = calls
+
+    def execute(self, sql):
+        self._calls.append(sql)
+        return self
+
+    def fetchone(self):
+        return (5,)
+
+    def close(self):
+        pass
+
+
+def test_prune_pg_deletes_natively_via_postgres_execute(monkeypatch):
+    calls = []
+    monkeypatch.setattr(prune.duckdb, "connect", lambda *a, **k: _FakeConn(calls))
+    stats_n, cs_n = prune.prune_pg("dsn", dt.datetime(2026, 7, 30, tzinfo=UTC))
+
+    deletes = [c for c in calls if "postgres_execute" in c and "DELETE" in c]
+    assert len(deletes) == 2  # child (changeset_stats) then parent (changesets)
+    assert "changeset_stats" in deletes[0]
+    assert "DELETE FROM changesets WHERE" in deletes[1] and "changeset_stats" not in deletes[1]
+    assert all("2026-07-30" in d for d in deletes)
+    assert (stats_n, cs_n) == (5, 5)