host
8006a6fcba
Bypass closed: automations._audit_row (direct audit_log INSERT, best-effort swallow) is deleted. run_due_automation now writes via cloud.auth.audit_on_conn, the same audit_log + audit_ledger pair AuthStore.audit() writes, sharing created_at. The audit pair runs before record_automation_run's commit, so task row, run row and audit pair commit or roll back together; chain-head IntegrityError retries the whole block up to 5 times, other errors roll back and raise (fail-closed). The claim stays committed, so a failed run is retried next interval, never recorded as a failed run. AuthStore.audit() now delegates to audit_on_conn; its pair, retry count and fail-closed semantics are unchanged.
193 lines
7.1 KiB
Python
193 lines
7.1 KiB
Python
"""自动化执行器入口 run_due_automation：抢单 + 配额门单测。"""
|
||
import os
|
||
import sqlite3
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
|
||
from cloud.auth import audit_on_conn
|
||
from cloud.automations import (
|
||
create_automation,
|
||
ensure_automations_tables,
|
||
list_automation_runs,
|
||
run_due_automation,
|
||
)
|
||
from cloud.ledger import GENESIS_HASH, compute_entry_hash
|
||
from cloud.server import init_db
|
||
|
||
|
||
def _tasks(db_path):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return conn.execute("SELECT id, prompt, status FROM tasks").fetchall()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _audit_count(db_path, action):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
return conn.execute(
|
||
"SELECT COUNT(*) FROM audit_log WHERE action = ?",
|
||
(action,)).fetchone()[0]
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
def _set_quota(db_path, user_id, micros):
|
||
conn = sqlite3.connect(db_path)
|
||
try:
|
||
conn.execute(
|
||
"CREATE TABLE IF NOT EXISTS user_quotas("
|
||
"user_id TEXT PRIMARY KEY, quota_usd_micros BIGINT, "
|
||
" updated_at DOUBLE PRECISION)")
|
||
conn.execute(
|
||
"INSERT OR REPLACE INTO user_quotas(user_id,"
|
||
" quota_usd_micros, updated_at) VALUES(?, ?, ?)",
|
||
(user_id, micros, time.time()))
|
||
conn.commit()
|
||
finally:
|
||
conn.close()
|
||
|
||
|
||
class AuditPairPgIdentityTestCase(unittest.TestCase):
|
||
"""模拟 PG：id 是 BIGINT identity，显式插 TEXT uuid 会报错。
|
||
|
||
假连接只接受不带 id 的 INSERT；旧 _audit_row（带 TEXT id）
|
||
在此假连接上应失败。接替它的 audit_on_conn 须同时做到：
|
||
audit_log 行不带 id、链条目与该行同 created_at、空链时
|
||
prev 为 GENESIS、entry_hash 按规范载荷重算一致；且本函数
|
||
不 commit（事务属于调用方，见 run_due_automation 文档）。
|
||
|
||
假连接类名须为 PgConnection：ledger.ensure_schema 靠类名
|
||
判定 PG 连接以跳过 SQLite 触发器，与真实 PG 适配层同款。
|
||
"""
|
||
|
||
class _FakeCursor:
|
||
def __init__(self, row=None):
|
||
self._row = row
|
||
self.lastrowid = 1
|
||
|
||
def fetchone(self):
|
||
return self._row
|
||
|
||
class PgConnection:
|
||
def __init__(self):
|
||
self.audit_rows = []
|
||
self.ledger_rows = []
|
||
self.committed = False
|
||
|
||
def execute(self, sql, params=()):
|
||
upper = sql.strip().upper()
|
||
if upper.startswith("INSERT INTO AUDIT_LOG(ID"):
|
||
raise ValueError("column \"id\" is of type bigint")
|
||
if upper.startswith("INSERT INTO AUDIT_LOG"):
|
||
self.audit_rows.append(params)
|
||
return AuditPairPgIdentityTestCase._FakeCursor()
|
||
if upper.startswith("SELECT ENTRY_HASH FROM AUDIT_LEDGER"):
|
||
return AuditPairPgIdentityTestCase._FakeCursor(None)
|
||
if upper.startswith("INSERT INTO AUDIT_LEDGER"):
|
||
self.ledger_rows.append(params)
|
||
return AuditPairPgIdentityTestCase._FakeCursor()
|
||
if upper.startswith("CREATE TABLE"):
|
||
return AuditPairPgIdentityTestCase._FakeCursor()
|
||
raise ValueError("unexpected sql in fake PG conn: %r" % sql)
|
||
|
||
def commit(self):
|
||
self.committed = True
|
||
|
||
def rollback(self):
|
||
pass
|
||
|
||
def test_pair_without_id_succeeds_on_pg_like_conn(self):
|
||
conn = self.PgConnection()
|
||
audit_on_conn(conn, "system", "p1", "automation_run_created",
|
||
"auto1 -> task1", created_at=123.0)
|
||
# 调用方拥有事务：本函数自己不 commit。
|
||
self.assertFalse(conn.committed)
|
||
self.assertEqual(len(conn.audit_rows), 1)
|
||
# 参数里没有 uuid：第一位是 system，第二位是 project_id。
|
||
self.assertEqual(conn.audit_rows[0][0], "system")
|
||
self.assertEqual(conn.audit_rows[0][1], "p1")
|
||
self.assertEqual(conn.audit_rows[0][4], 123.0)
|
||
self.assertEqual(len(conn.ledger_rows), 1)
|
||
ledger = conn.ledger_rows[0]
|
||
# audit_ledger 列序：user_id, project_id, task_id, action,
|
||
# detail, created_at, prev_hash, entry_hash。
|
||
self.assertEqual(ledger[0], "system")
|
||
self.assertEqual(ledger[1], "p1")
|
||
self.assertIsNone(ledger[2])
|
||
self.assertEqual(ledger[3], "automation_run_created")
|
||
self.assertEqual(ledger[4], "auto1 -> task1")
|
||
self.assertEqual(ledger[5], 123.0)
|
||
self.assertEqual(ledger[6], GENESIS_HASH)
|
||
self.assertEqual(
|
||
ledger[7],
|
||
compute_entry_hash(GENESIS_HASH, "system", "p1", None,
|
||
"automation_run_created",
|
||
"auto1 -> task1", 123.0))
|
||
|
||
|
||
class RunDueAutomationTestCase(unittest.TestCase):
|
||
def setUp(self):
|
||
self.tmp = tempfile.TemporaryDirectory()
|
||
self.db_path = os.path.join(self.tmp.name, "control.db")
|
||
init_db(self.db_path)
|
||
self.conn = sqlite3.connect(self.db_path)
|
||
self.conn.row_factory = sqlite3.Row
|
||
ensure_automations_tables(self.conn)
|
||
|
||
def tearDown(self):
|
||
self.conn.close()
|
||
self.tmp.cleanup()
|
||
|
||
def _auto(self, owner="owner1"):
|
||
return create_automation(self.conn, "proj", "daily", "do thing",
|
||
owner, interval_seconds=60)
|
||
|
||
def test_run_creates_task_within_quota(self):
|
||
auto = self._auto()
|
||
run = run_due_automation(self.conn, auto["id"])
|
||
self.assertIsNotNone(run)
|
||
self.assertEqual(run["status"], "created")
|
||
self.assertIsNotNone(run["task_id"])
|
||
tasks = _tasks(self.db_path)
|
||
self.assertEqual(len(tasks), 1)
|
||
self.assertEqual(tasks[0][1], "do thing")
|
||
self.assertEqual(tasks[0][2], "queued")
|
||
self.assertEqual(
|
||
_audit_count(self.db_path, "automation_run_created"), 1)
|
||
# 同一间隔内第二次跑：已抢过，返回 None，不重复建任务。
|
||
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
|
||
self.assertEqual(len(_tasks(self.db_path)), 1)
|
||
|
||
def test_run_skips_when_over_quota(self):
|
||
auto = self._auto()
|
||
_set_quota(self.db_path, "owner1", 0)
|
||
run = run_due_automation(self.conn, auto["id"])
|
||
self.assertIsNotNone(run)
|
||
self.assertEqual(run["status"], "skipped_quota")
|
||
self.assertIsNone(run["task_id"])
|
||
self.assertEqual(_tasks(self.db_path), [])
|
||
self.assertEqual(
|
||
_audit_count(self.db_path, "automation_run_quota_skipped"), 1)
|
||
# 跳过仍记了产出（不是失败），同间隔内不再重跑。
|
||
runs = list_automation_runs(self.conn, auto["id"])
|
||
self.assertEqual(len(runs), 1)
|
||
self.assertEqual(runs[0]["status"], "skipped_quota")
|
||
|
||
def test_run_not_due_returns_none(self):
|
||
auto = self._auto()
|
||
self.conn.execute(
|
||
"UPDATE automations SET last_run_at = ? WHERE id = ?",
|
||
(time.time(), auto["id"]))
|
||
self.conn.commit()
|
||
self.assertIsNone(run_due_automation(self.conn, auto["id"]))
|
||
self.assertEqual(_tasks(self.db_path), [])
|
||
self.assertEqual(
|
||
list_automation_runs(self.conn, auto["id"]), [])
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|