a2a-source-edit: write agent.py
This commit is contained in:
139
agent.py
139
agent.py
@@ -514,36 +514,38 @@ def _connect_db() -> Any:
|
||||
|
||||
|
||||
async def _save_audit_and_receipt(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
|
||||
def work() -> None:
|
||||
from psycopg.types.json import Jsonb
|
||||
await asyncio.to_thread(_save_audit_and_receipt_sync, tenant, audit_id, payload, receipt)
|
||||
|
||||
with _connect_db() as conn:
|
||||
with conn.transaction():
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
|
||||
target_url = EXCLUDED.target_url,
|
||||
ok = EXCLUDED.ok,
|
||||
report = EXCLUDED.report,
|
||||
payload = EXCLUDED.payload,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET
|
||||
payload = EXCLUDED.payload,
|
||||
created_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
|
||||
)
|
||||
await asyncio.to_thread(work)
|
||||
|
||||
def _save_audit_and_receipt_sync(tenant: str, audit_id: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
with _connect_db() as conn:
|
||||
with conn.transaction():
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
|
||||
target_url = EXCLUDED.target_url,
|
||||
ok = EXCLUDED.ok,
|
||||
report = EXCLUDED.report,
|
||||
payload = EXCLUDED.payload,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, payload["target_url"], True, Jsonb(payload["report"]), Jsonb(payload)),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET
|
||||
payload = EXCLUDED.payload,
|
||||
created_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
|
||||
)
|
||||
|
||||
|
||||
async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], started_at: str) -> None:
|
||||
@@ -555,54 +557,57 @@ async def _persist_failure_if_possible(tenant: str, audit_id: str, target_url: s
|
||||
"created_at": _now_iso(),
|
||||
"code": payload.get("code"),
|
||||
}
|
||||
payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()}
|
||||
|
||||
def work() -> None:
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
with _connect_db() as conn:
|
||||
with conn.transaction():
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
|
||||
target_url = EXCLUDED.target_url,
|
||||
ok = EXCLUDED.ok,
|
||||
report = EXCLUDED.report,
|
||||
payload = EXCLUDED.payload,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, _safe_url_echo(target_url), False, Jsonb({}), Jsonb(payload)),
|
||||
)
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
|
||||
""",
|
||||
(tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
|
||||
)
|
||||
saved_payload = {**payload, "receipt": receipt, "created_at": started_at, "updated_at": _now_iso()}
|
||||
try:
|
||||
await asyncio.to_thread(work)
|
||||
await asyncio.to_thread(_persist_failure_sync, tenant, audit_id, target_url, saved_payload, receipt)
|
||||
except Exception: # noqa: BLE001
|
||||
return
|
||||
|
||||
|
||||
async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None:
|
||||
def work() -> dict[str, Any] | None:
|
||||
with _connect_db() as conn:
|
||||
def _persist_failure_sync(tenant: str, audit_id: str, target_url: str, payload: dict[str, Any], receipt: dict[str, Any]) -> None:
|
||||
from psycopg.types.json import Jsonb
|
||||
|
||||
with _connect_db() as conn:
|
||||
with conn.transaction():
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT payload FROM launch_audits WHERE tenant_key = %s AND audit_id = %s",
|
||||
(tenant, audit_id),
|
||||
"""
|
||||
INSERT INTO launch_audits (tenant_key, audit_id, target_url, ok, report, payload)
|
||||
VALUES (%s, %s, %s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, audit_id) DO UPDATE SET
|
||||
target_url = EXCLUDED.target_url,
|
||||
ok = EXCLUDED.ok,
|
||||
report = EXCLUDED.report,
|
||||
payload = EXCLUDED.payload,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
(tenant, audit_id, _safe_url_echo(target_url), False, Jsonb({}), Jsonb(payload)),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
return row[0] if row else None
|
||||
cur.execute(
|
||||
"""
|
||||
INSERT INTO launch_receipts (tenant_key, audit_id, receipt_id, payload)
|
||||
VALUES (%s, %s, %s, %s)
|
||||
ON CONFLICT (tenant_key, receipt_id) DO UPDATE SET payload = EXCLUDED.payload
|
||||
""",
|
||||
(tenant, audit_id, receipt["receipt_id"], Jsonb(receipt)),
|
||||
)
|
||||
|
||||
|
||||
async def _load_audit(tenant: str, audit_id: str) -> dict[str, Any] | None:
|
||||
try:
|
||||
return await asyncio.to_thread(work)
|
||||
return await asyncio.to_thread(_load_audit_sync, tenant, audit_id)
|
||||
except RuntimeError as exc:
|
||||
if str(exc) == "database_unavailable":
|
||||
return {"ok": False, "code": "database_unavailable", "message": "Persistent storage is not configured for this deployment.", "audit_id": audit_id}
|
||||
raise
|
||||
|
||||
|
||||
def _load_audit_sync(tenant: str, audit_id: str) -> dict[str, Any] | None:
|
||||
with _connect_db() as conn:
|
||||
with conn.cursor() as cur:
|
||||
cur.execute(
|
||||
"SELECT payload FROM launch_audits WHERE tenant_key = %s AND audit_id = %s",
|
||||
(tenant, audit_id),
|
||||
)
|
||||
row = cur.fetchone()
|
||||
return row[0] if row else None
|
||||
|
||||
Reference in New Issue
Block a user