mo
a724615a9a
- DAG hadoop_to_trino (deployed to Airflow) + worker hadoop_to_trino.py (on the
Hadoop master) move HDFS historical_sales -> iceberg.hadoop.historical_sales_hdfs.
- api/movements.py: movement registry, Airflow trigger+watch, run tracking
(state/duration/rows), endpoints /api/movements, /{id}/run, /runs.
- agent_ops.py: ETL-agent loop autonomously triggers movements on an interval
and logs each run; /api/agent-ops/etl/toggle + etl status.
374 lines
16 KiB
Python
374 lines
16 KiB
Python
"""Autonomous agent DML operations against the source databases.
|
|
|
|
Agents continuously perform small, guard-railed INSERT/UPDATE/DELETE operations
|
|
on the operational source databases (PostgreSQL, MySQL, MongoDB) so that the
|
|
Debezium CDC connectors capture a steady stream of changes that flow through
|
|
Kafka and downstream.
|
|
|
|
Safety guardrails:
|
|
- Only a fixed whitelist of tables/collections is touched (the exact ones the
|
|
Debezium connectors capture: public.sales_orders, hr.employee_events,
|
|
supplychain.events).
|
|
- Every agent-created row is tagged (notes/payload marker, or an `atc_agent`
|
|
field for Mongo). UPDATE and DELETE operate ONLY on rows the agents created
|
|
themselves, never on seed/real data.
|
|
- A hard per-tick row limit and an env kill-switch bound the activity.
|
|
|
|
The work runs in a background loop started from the app lifespan. DB drivers
|
|
are synchronous, so each operation is executed via ``asyncio.to_thread``.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import os
|
|
import random
|
|
import uuid
|
|
from collections import deque
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
|
|
from fastapi import APIRouter, Body
|
|
from fastapi.responses import JSONResponse
|
|
|
|
router = APIRouter(prefix="/api/agent-ops", tags=["agent-ops"])
|
|
|
|
# ── Config (env-overridable) ────────────────────────────────────────────────
|
|
DB_HOST = os.getenv("SRC_DB_HOST", "10.0.21.51")
|
|
DB_USER = os.getenv("SRC_DB_USER", "mo")
|
|
DB_PASSWORD = os.getenv("SRC_DB_PASSWORD", "Dell2026!")
|
|
PG_DB = os.getenv("SRC_PG_DB", "postgres")
|
|
PG_PORT = int(os.getenv("SRC_PG_PORT", "5432"))
|
|
MYSQL_DB = os.getenv("SRC_MYSQL_DB", "hr")
|
|
MYSQL_PORT = int(os.getenv("SRC_MYSQL_PORT", "3306"))
|
|
MONGO_URI = os.getenv("SRC_MONGO_URI", f"mongodb://{DB_HOST}:27017/?replicaSet=rs0")
|
|
MONGO_DB = os.getenv("SRC_MONGO_DB", "supplychain")
|
|
|
|
AGENT_TAG = "ATC-AGENT" # marker stored on agent-created rows
|
|
DML_AGENT = "data-custodian"
|
|
|
|
_INTERVAL = float(os.getenv("AGENT_DML_INTERVAL_SECONDS", "45"))
|
|
_MAX_ROWS = int(os.getenv("AGENT_DML_MAX_ROWS", "5"))
|
|
# ETL-agent: autonomously trigger data movements on an interval.
|
|
_ETL_INTERVAL = float(os.getenv("ETL_AGENT_INTERVAL_SECONDS", "300"))
|
|
# Movements the ETL-agent cycles through autonomously (must exist in movements.py).
|
|
_ETL_ROTATION = [m for m in os.getenv("ETL_AGENT_MOVEMENTS", "hadoop_to_trino,gen_postgres,gen_mysql,gen_mongodb").split(",") if m]
|
|
# How many agent-created rows we keep around per source before favouring DELETE.
|
|
_POOL_CAP = int(os.getenv("AGENT_DML_POOL_CAP", "400"))
|
|
|
|
_state: dict[str, Any] = {
|
|
"enabled": os.getenv("AGENT_DML_ENABLED", "1") not in ("0", "false", "False", ""),
|
|
"interval": _INTERVAL,
|
|
"max_rows": _MAX_ROWS,
|
|
"started": False,
|
|
"ops_total": 0,
|
|
"last_op": None,
|
|
"last_error": None,
|
|
"by_source": {"postgres": 0, "mysql": 0, "mongodb": 0},
|
|
"by_op": {"insert": 0, "update": 0, "delete": 0},
|
|
}
|
|
|
|
_etl_state: dict[str, Any] = {
|
|
"enabled": os.getenv("ETL_AGENT_ENABLED", "1") not in ("0", "false", "False", ""),
|
|
"interval": _ETL_INTERVAL,
|
|
"rotation": _ETL_ROTATION,
|
|
"idx": 0,
|
|
"runs_total": 0,
|
|
"last": None,
|
|
"started": False,
|
|
}
|
|
|
|
# ── Realistic value domains ─────────────────────────────────────────────────
|
|
REGIONS = ["NA", "EU", "EMEA", "APAC", "LATAM"]
|
|
CHANNELS = ["online", "retail", "partner", "wholesale", "direct"]
|
|
CURRENCIES = ["USD", "EUR", "GBP"]
|
|
ORDER_STATUS = ["NEW", "PROCESSING", "SHIPPED", "DELIVERED", "RETURNED", "CANCELLED"]
|
|
DEPARTMENTS = ["Engineering", "Sales", "HR", "Finance", "Operations", "Marketing"]
|
|
ROLES = ["Analyst", "Manager", "Engineer", "Director", "Specialist"]
|
|
EVENT_TYPES = ["HIRE", "PROMOTION", "SALARY_CHANGE", "TRANSFER", "TERMINATION"]
|
|
MONGO_TYPES = ["ORDER", "SHIPMENT", "RETURN", "RESTOCK", "UPDATE"]
|
|
MONGO_SOURCES = ["CRM", "ERP", "WMS", "POS"]
|
|
|
|
|
|
# ── Feed / terminal helpers (lazy import to avoid circular import) ───────────
|
|
async def _emit(message: str, level: str = "info") -> None:
|
|
try:
|
|
from main import add_feed, publish_event
|
|
from agent_terminal import terminal_log
|
|
|
|
entry = add_feed(DML_AGENT, message, level)
|
|
await publish_event({"type": "feed", "entry": entry})
|
|
await terminal_log(DML_AGENT, message, level=level, phase="dml")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
# ── Synchronous DB operations (run in a thread) ──────────────────────────────
|
|
def _pg_conn():
|
|
import psycopg2
|
|
|
|
return psycopg2.connect(
|
|
host=DB_HOST, dbname=PG_DB, user=DB_USER, password=DB_PASSWORD, port=PG_PORT, connect_timeout=8
|
|
)
|
|
|
|
|
|
def _mysql_conn():
|
|
import pymysql
|
|
|
|
return pymysql.connect(
|
|
host=DB_HOST, user=DB_USER, password=DB_PASSWORD, database=MYSQL_DB,
|
|
port=MYSQL_PORT, connect_timeout=8, autocommit=True,
|
|
)
|
|
|
|
|
|
def _mongo_coll():
|
|
import pymongo
|
|
|
|
client = pymongo.MongoClient(MONGO_URI, serverSelectionTimeoutMS=8000)
|
|
return client, client[MONGO_DB]["events"]
|
|
|
|
|
|
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:
|
|
if op == "insert":
|
|
rid = uuid.uuid4().hex[:8]
|
|
cur.execute(
|
|
"""INSERT INTO public.sales_orders
|
|
(customer_id, product_id, region, sales_channel, order_ts, amount, currency, order_status, notes)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s,%s,%s) RETURNING order_id""",
|
|
(
|
|
random.randint(1, 50000), random.randint(1, 2000),
|
|
random.choice(REGIONS), random.choice(CHANNELS),
|
|
datetime.now(timezone.utc), round(random.uniform(10, 9999), 2),
|
|
random.choice(CURRENCIES), random.choice(ORDER_STATUS), _tag(rid),
|
|
),
|
|
)
|
|
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=%s",
|
|
(random.choice(ORDER_STATUS), round(random.uniform(0.9, 1.2), 2), oid),
|
|
)
|
|
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:
|
|
if op == "insert":
|
|
rid = uuid.uuid4().hex[:8]
|
|
cur.execute(
|
|
"""INSERT INTO hr.employee_events
|
|
(employee_id, department, role_name, region, event_type, salary_change, event_ts, notes)
|
|
VALUES (%s,%s,%s,%s,%s,%s,%s,%s)""",
|
|
(
|
|
random.randint(1, 20000), random.choice(DEPARTMENTS), random.choice(ROLES),
|
|
random.choice(REGIONS), random.choice(EVENT_TYPES),
|
|
round(random.uniform(-5000, 15000), 2), datetime.now(timezone.utc), _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 event_id=%s",
|
|
(round(random.uniform(-5000, 15000), 2), random.choice(EVENT_TYPES), eid),
|
|
)
|
|
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:
|
|
if op == "insert":
|
|
rid = uuid.uuid4().hex[:8]
|
|
doc = {
|
|
"event_id": str(uuid.uuid4()),
|
|
"type": random.choice(MONGO_TYPES),
|
|
"region": random.choice(REGIONS),
|
|
"source": random.choice(MONGO_SOURCES),
|
|
"amount": round(random.uniform(10, 50000), 4),
|
|
"ts": datetime.now(timezone.utc),
|
|
"payload": "X" * 200,
|
|
"atc_agent": True,
|
|
"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":
|
|
oid = random.choice(list(pool))
|
|
coll.update_one(
|
|
{"_id": oid},
|
|
{"$set": {"type": random.choice(MONGO_TYPES), "amount": round(random.uniform(10, 50000), 4)}},
|
|
)
|
|
return f"UPDATE events _id={oid}"
|
|
oid = pool.popleft()
|
|
coll.delete_one({"_id": oid})
|
|
return f"DELETE events _id={oid}"
|
|
finally:
|
|
client.close()
|
|
|
|
|
|
_DISPATCH = {"postgres": _pg_dml, "mysql": _mysql_dml, "mongodb": _mongo_dml}
|
|
_SRC_LABEL = {"postgres": "PostgreSQL sales_orders", "mysql": "MySQL employee_events", "mongodb": "MongoDB events"}
|
|
|
|
|
|
def _pick_op() -> str:
|
|
# Insert-heavy so a pool of agent rows exists for safe UPDATE/DELETE.
|
|
return random.choices(["insert", "update", "delete"], weights=[0.5, 0.3, 0.2], k=1)[0]
|
|
|
|
|
|
async def _run_one(source: str | None = None, op: str | None = None) -> dict[str, Any]:
|
|
source = source or random.choice(list(_DISPATCH.keys()))
|
|
op = op or _pick_op()
|
|
fn = _DISPATCH.get(source)
|
|
if not fn:
|
|
return {"ok": False, "error": f"unknown source {source}"}
|
|
try:
|
|
detail = await asyncio.to_thread(fn, op)
|
|
_state["ops_total"] += 1
|
|
_state["by_source"][source] = _state["by_source"].get(source, 0) + 1
|
|
actual_op = detail.split(" ", 1)[0].lower()
|
|
if actual_op in _state["by_op"]:
|
|
_state["by_op"][actual_op] += 1
|
|
_state["last_op"] = {"source": source, "op": op, "detail": detail, "ts": datetime.now(timezone.utc).isoformat()}
|
|
_state["last_error"] = None
|
|
await _emit(f"[agent-dml] {_SRC_LABEL[source]}: {detail} — Debezium will capture this change", "info")
|
|
return {"ok": True, "source": source, "op": op, "detail": detail}
|
|
except Exception as exc:
|
|
_state["last_error"] = str(exc)
|
|
await _emit(f"[agent-dml] {_SRC_LABEL.get(source, source)}: operation failed: {str(exc)[:120]}", "err")
|
|
return {"ok": False, "source": source, "op": op, "error": str(exc)}
|
|
|
|
|
|
async def agent_dml_loop() -> None:
|
|
"""Background loop: continuously perform guard-railed DML on the sources."""
|
|
_state["started"] = True
|
|
await asyncio.sleep(8) # let the app settle / DB reachable
|
|
await _emit("[agent-dml] Autonomous DML agent online — generating live changes for Debezium", "info")
|
|
while True:
|
|
try:
|
|
if _state["enabled"]:
|
|
await _run_one()
|
|
except Exception as exc: # never let the loop die
|
|
_state["last_error"] = str(exc)
|
|
await asyncio.sleep(max(5.0, float(_state["interval"])))
|
|
|
|
|
|
async def etl_agent_loop() -> None:
|
|
"""Background loop: ETL-agents autonomously trigger data movements and log
|
|
each run (rows, duration, status) to the feed."""
|
|
_etl_state["started"] = True
|
|
await asyncio.sleep(30) # let the platform settle
|
|
while True:
|
|
try:
|
|
if _etl_state["enabled"] and _etl_state["rotation"]:
|
|
from movements import trigger_and_watch, MOVEMENT_BY_ID
|
|
|
|
mid = _etl_state["rotation"][_etl_state["idx"] % len(_etl_state["rotation"])]
|
|
_etl_state["idx"] += 1
|
|
if mid in MOVEMENT_BY_ID:
|
|
result = await trigger_and_watch(mid, autonomous=True)
|
|
_etl_state["runs_total"] += 1
|
|
_etl_state["last"] = {"movement_id": mid, "state": result.get("state"),
|
|
"rows": result.get("rows"), "duration_s": result.get("duration_s"),
|
|
"ts": datetime.now(timezone.utc).isoformat()}
|
|
except Exception as exc:
|
|
_etl_state["last"] = {"error": str(exc), "ts": datetime.now(timezone.utc).isoformat()}
|
|
await asyncio.sleep(max(30.0, float(_etl_state["interval"])))
|
|
|
|
|
|
# ── Endpoints ────────────────────────────────────────────────────────────────
|
|
@router.get("/status")
|
|
async def status() -> JSONResponse:
|
|
pools = {k: len(v) for k, v in _pools.items()}
|
|
return JSONResponse({"ok": True, "pools": pools, "etl": _etl_state, **_state})
|
|
|
|
|
|
@router.post("/etl/toggle")
|
|
async def etl_toggle(body: dict[str, Any] = Body(default={})) -> JSONResponse:
|
|
if "enabled" in body:
|
|
_etl_state["enabled"] = bool(body["enabled"])
|
|
else:
|
|
_etl_state["enabled"] = not _etl_state["enabled"]
|
|
if "interval" in body:
|
|
try:
|
|
_etl_state["interval"] = max(30.0, float(body["interval"]))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
await _emit(f"[etl-agent] Autonomous ETL movements {'ENABLED' if _etl_state['enabled'] else 'PAUSED'} by operator", "warn")
|
|
return JSONResponse({"ok": True, "enabled": _etl_state["enabled"], "interval": _etl_state["interval"]})
|
|
|
|
|
|
@router.post("/toggle")
|
|
async def toggle(body: dict[str, Any] = Body(default={})) -> JSONResponse:
|
|
if "enabled" in body:
|
|
_state["enabled"] = bool(body["enabled"])
|
|
else:
|
|
_state["enabled"] = not _state["enabled"]
|
|
if "interval" in body:
|
|
try:
|
|
_state["interval"] = max(5.0, float(body["interval"]))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
await _emit(f"[agent-dml] Autonomous DML {'ENABLED' if _state['enabled'] else 'PAUSED'} by operator", "warn")
|
|
return JSONResponse({"ok": True, "enabled": _state["enabled"], "interval": _state["interval"]})
|
|
|
|
|
|
@router.post("/run-once")
|
|
async def run_once(body: dict[str, Any] = Body(default={})) -> JSONResponse:
|
|
source = body.get("source")
|
|
op = body.get("op")
|
|
result = await _run_one(source=source, op=op)
|
|
return JSONResponse(result)
|