fix(agents): use in-memory PK pools for DML (avoid full-table scans on 54M-row tables)

api/agent_ops.py

@@ -23,8 +23,8 @@ from __future__ import annotations
 import asyncio
 import os
 import random
-import time
 import uuid
+from collections import deque
 from datetime import datetime, timezone
 from typing import Any
 
@@ -118,14 +118,34 @@ def _tag(run_id: str) -> str:
     return f"{AGENT_TAG} {run_id}"
 
 
+# In-memory pools of agent-created primary keys per source. All UPDATE/DELETE
+# operate by PK on these (instant), so we never scan the large unindexed
+# `notes` column. Bounded by _POOL_CAP; orphaned tagged rows after a restart
+# are harmless (clearly marked) and simply not re-tracked.
+_pools: dict[str, deque] = {
+    "postgres": deque(maxlen=_POOL_CAP),
+    "mysql": deque(maxlen=_POOL_CAP),
+    "mongodb": deque(maxlen=_POOL_CAP),
+}
+
+
+def _resolve_op(op: str, pool: deque) -> str:
+    """UPDATE/DELETE require an existing agent row; otherwise fall back to INSERT."""
+    if op in ("update", "delete") and not pool:
+        return "insert"
+    if op == "insert" and len(pool) >= _POOL_CAP:
+        return "delete"  # keep the agent footprint bounded
+    return op
+
+
 def _pg_dml(op: str) -> str:
+    pool = _pools["postgres"]
+    op = _resolve_op(op, pool)
     conn = _pg_conn()
     try:
         conn.autocommit = True
         with conn.cursor() as cur:
-            cur.execute("SELECT count(*) FROM public.sales_orders WHERE notes LIKE %s", (AGENT_TAG + "%",))
-            pool = cur.fetchone()[0]
-            if op == "insert" or (op == "delete" and pool == 0) or (op == "update" and pool == 0):
+            if op == "insert":
                 rid = uuid.uuid4().hex[:8]
                 cur.execute(
                     """INSERT INTO public.sales_orders
@@ -139,38 +159,29 @@ def _pg_dml(op: str) -> str:
                     ),
                 )
                 oid = cur.fetchone()[0]
+                pool.append(oid)
                 return f"INSERT sales_orders order_id={oid} ({_tag(rid)})"
             if op == "update":
+                oid = random.choice(list(pool))
                 cur.execute(
-                    """UPDATE public.sales_orders SET order_status=%s, amount=round(amount*%s,2)
-                       WHERE order_id IN (
-                         SELECT order_id FROM public.sales_orders WHERE notes LIKE %s
-                         ORDER BY order_id DESC LIMIT 1)
-                       RETURNING order_id""",
-                    (random.choice(ORDER_STATUS), round(random.uniform(0.9, 1.2), 2), AGENT_TAG + "%"),
+                    "UPDATE public.sales_orders SET order_status=%s, amount=round(amount*%s,2) WHERE order_id=%s",
+                    (random.choice(ORDER_STATUS), round(random.uniform(0.9, 1.2), 2), oid),
                 )
-                row = cur.fetchone()
-                return f"UPDATE sales_orders order_id={row[0]}" if row else "UPDATE sales_orders (no agent rows)"
-            # delete
-            cur.execute(
-                """DELETE FROM public.sales_orders WHERE order_id IN (
-                     SELECT order_id FROM public.sales_orders WHERE notes LIKE %s
-                     ORDER BY order_id ASC LIMIT 1) RETURNING order_id""",
-                (AGENT_TAG + "%",),
-            )
-            row = cur.fetchone()
-            return f"DELETE sales_orders order_id={row[0]}" if row else "DELETE sales_orders (no agent rows)"
+                return f"UPDATE sales_orders order_id={oid}"
+            oid = pool.popleft()
+            cur.execute("DELETE FROM public.sales_orders WHERE order_id=%s", (oid,))
+            return f"DELETE sales_orders order_id={oid}"
     finally:
         conn.close()
 
 
 def _mysql_dml(op: str) -> str:
+    pool = _pools["mysql"]
+    op = _resolve_op(op, pool)
     conn = _mysql_conn()
     try:
         with conn.cursor() as cur:
-            cur.execute("SELECT count(*) FROM hr.employee_events WHERE notes LIKE %s", (AGENT_TAG + "%",))
-            pool = cur.fetchone()[0]
-            if op == "insert" or (op in ("update", "delete") and pool == 0):
+            if op == "insert":
                 rid = uuid.uuid4().hex[:8]
                 cur.execute(
                     """INSERT INTO hr.employee_events
@@ -182,28 +193,29 @@ def _mysql_dml(op: str) -> str:
                         round(random.uniform(-5000, 15000), 2), datetime.now(timezone.utc), _tag(rid),
                     ),
                 )
-                return f"INSERT employee_events event_id={cur.lastrowid} ({_tag(rid)})"
+                eid = cur.lastrowid
+                pool.append(eid)
+                return f"INSERT employee_events event_id={eid} ({_tag(rid)})"
             if op == "update":
+                eid = random.choice(list(pool))
                 cur.execute(
-                    """UPDATE hr.employee_events SET salary_change=%s, event_type=%s
-                       WHERE notes LIKE %s ORDER BY event_id DESC LIMIT 1""",
-                    (round(random.uniform(-5000, 15000), 2), random.choice(EVENT_TYPES), AGENT_TAG + "%"),
+                    "UPDATE hr.employee_events SET salary_change=%s, event_type=%s WHERE event_id=%s",
+                    (round(random.uniform(-5000, 15000), 2), random.choice(EVENT_TYPES), eid),
                 )
-                return f"UPDATE employee_events ({cur.rowcount} row)"
-            cur.execute(
-                "DELETE FROM hr.employee_events WHERE notes LIKE %s ORDER BY event_id ASC LIMIT 1",
-                (AGENT_TAG + "%",),
-            )
-            return f"DELETE employee_events ({cur.rowcount} row)"
+                return f"UPDATE employee_events event_id={eid}"
+            eid = pool.popleft()
+            cur.execute("DELETE FROM hr.employee_events WHERE event_id=%s", (eid,))
+            return f"DELETE employee_events event_id={eid}"
     finally:
         conn.close()
 
 
 def _mongo_dml(op: str) -> str:
+    pool = _pools["mongodb"]
+    op = _resolve_op(op, pool)
     client, coll = _mongo_coll()
     try:
-        pool = coll.count_documents({"atc_agent": True}, limit=1)
-        if op == "insert" or (op in ("update", "delete") and pool == 0):
+        if op == "insert":
             rid = uuid.uuid4().hex[:8]
             doc = {
                 "event_id": str(uuid.uuid4()),
@@ -217,21 +229,18 @@ def _mongo_dml(op: str) -> str:
                 "agent_run": rid,
             }
             res = coll.insert_one(doc)
+            pool.append(res.inserted_id)
             return f"INSERT events _id={res.inserted_id} (agent_run={rid})"
         if op == "update":
-            doc = coll.find_one({"atc_agent": True}, sort=[("_id", -1)])
-            if not doc:
-                return "UPDATE events (no agent docs)"
+            oid = random.choice(list(pool))
             coll.update_one(
-                {"_id": doc["_id"]},
+                {"_id": oid},
                 {"$set": {"type": random.choice(MONGO_TYPES), "amount": round(random.uniform(10, 50000), 4)}},
             )
-            return f"UPDATE events _id={doc['_id']}"
-        doc = coll.find_one({"atc_agent": True}, sort=[("_id", 1)])
-        if not doc:
-            return "DELETE events (no agent docs)"
-        coll.delete_one({"_id": doc["_id"]})
-        return f"DELETE events _id={doc['_id']}"
+            return f"UPDATE events _id={oid}"
+        oid = pool.popleft()
+        coll.delete_one({"_id": oid})
+        return f"DELETE events _id={oid}"
     finally:
         client.close()
 
@@ -285,7 +294,8 @@ async def agent_dml_loop() -> None:
 # ── Endpoints ────────────────────────────────────────────────────────────────
 @router.get("/status")
 async def status() -> JSONResponse:
-    return JSONResponse({"ok": True, **_state})
+    pools = {k: len(v) for k, v in _pools.items()}
+    return JSONResponse({"ok": True, "pools": pools, **_state})
 
 
 @router.post("/toggle")