217 lines
10 KiB
Python
217 lines
10 KiB
Python
|
"""Small SQLite repository. Queue transitions are serialized transactions."""
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from datetime import datetime, timezone
|
||
|
|
import json
|
||
|
|
import sqlite3
|
||
|
|
import uuid
|
||
|
|
|
||
|
|
TERMINAL = {"succeeded", "failed", "cancelled", "timed_out", "interrupted"}
|
||
|
|
|
||
|
|
|
||
|
|
def now():
|
||
|
|
return datetime.now(timezone.utc).isoformat()
|
||
|
|
|
||
|
|
|
||
|
|
class Store:
|
||
|
|
def __init__(self, settings):
|
||
|
|
self.settings = settings
|
||
|
|
settings.prepare()
|
||
|
|
with self.connect() as db:
|
||
|
|
version = db.execute("PRAGMA user_version").fetchone()[0]
|
||
|
|
if version not in (0, 1):
|
||
|
|
raise RuntimeError("Unsupported database version")
|
||
|
|
db.executescript("""
|
||
|
|
PRAGMA journal_mode=WAL;
|
||
|
|
CREATE TABLE IF NOT EXISTS scripts(id TEXT PRIMARY KEY, data TEXT NOT NULL);
|
||
|
|
CREATE TABLE IF NOT EXISTS versions(
|
||
|
|
script_id TEXT NOT NULL, version INTEGER NOT NULL, data TEXT NOT NULL,
|
||
|
|
PRIMARY KEY(script_id,version));
|
||
|
|
CREATE TABLE IF NOT EXISTS runs(
|
||
|
|
id TEXT PRIMARY KEY, request_id TEXT NOT NULL UNIQUE,
|
||
|
|
script_id TEXT NOT NULL, state TEXT NOT NULL,
|
||
|
|
cancel_requested INTEGER NOT NULL DEFAULT 0, data TEXT NOT NULL);
|
||
|
|
CREATE TABLE IF NOT EXISTS health(key TEXT PRIMARY KEY, value TEXT NOT NULL);
|
||
|
|
PRAGMA user_version=1;
|
||
|
|
""")
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def connect(self, write=False):
|
||
|
|
db = sqlite3.connect(self.settings.db, timeout=10)
|
||
|
|
db.row_factory = sqlite3.Row
|
||
|
|
try:
|
||
|
|
if write:
|
||
|
|
db.execute("BEGIN IMMEDIATE")
|
||
|
|
yield db
|
||
|
|
db.commit()
|
||
|
|
except Exception:
|
||
|
|
db.rollback()
|
||
|
|
raise
|
||
|
|
finally:
|
||
|
|
db.close()
|
||
|
|
|
||
|
|
def scripts(self):
|
||
|
|
with self.connect() as db:
|
||
|
|
return [json.loads(r[0]) for r in db.execute("SELECT data FROM scripts ORDER BY id")]
|
||
|
|
|
||
|
|
def script(self, identity, versions=False):
|
||
|
|
with self.connect() as db:
|
||
|
|
row = db.execute("SELECT data FROM scripts WHERE id=?", (identity,)).fetchone()
|
||
|
|
if not row:
|
||
|
|
raise KeyError("脚本不存在")
|
||
|
|
item = json.loads(row[0])
|
||
|
|
if versions:
|
||
|
|
item["versions"] = [json.loads(r[0]) for r in db.execute(
|
||
|
|
"SELECT data FROM versions WHERE script_id=? ORDER BY version DESC", (identity,))]
|
||
|
|
return item
|
||
|
|
|
||
|
|
def save_script(self, item):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
db.execute("INSERT OR REPLACE INTO scripts VALUES(?,?)", (item["id"], json.dumps(item)))
|
||
|
|
|
||
|
|
def version(self, identity, version):
|
||
|
|
with self.connect() as db:
|
||
|
|
row = db.execute("SELECT data FROM versions WHERE script_id=? AND version=?", (identity, version)).fetchone()
|
||
|
|
if not row:
|
||
|
|
raise KeyError("版本不存在")
|
||
|
|
return json.loads(row[0])
|
||
|
|
|
||
|
|
def upload_version(self, item, content, filename):
|
||
|
|
# A transaction prevents two simultaneous uploads from assigning the same version.
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
current = db.execute("SELECT data FROM scripts WHERE id=?", (item["id"],)).fetchone()
|
||
|
|
if current:
|
||
|
|
item = json.loads(current[0])
|
||
|
|
version = item["version"] + 1
|
||
|
|
target = self.settings.library / item["id"] / "versions" / str(version) / filename
|
||
|
|
target.parent.mkdir(parents=True, exist_ok=False)
|
||
|
|
try:
|
||
|
|
with target.open("xb") as handle:
|
||
|
|
handle.write(content)
|
||
|
|
import hashlib
|
||
|
|
value = {"version": version, "sha256": hashlib.sha256(content).hexdigest(),
|
||
|
|
"created_at": now(), "filename": filename, "path": str(target)}
|
||
|
|
item.update(version=version, sha256=value["sha256"], source_path=str(target), updated_at=now())
|
||
|
|
db.execute("INSERT INTO versions VALUES(?,?,?)", (item["id"], version, json.dumps(value)))
|
||
|
|
db.execute("INSERT OR REPLACE INTO scripts VALUES(?,?)", (item["id"], json.dumps(item)))
|
||
|
|
except Exception:
|
||
|
|
target.unlink(missing_ok=True)
|
||
|
|
target.parent.rmdir()
|
||
|
|
raise
|
||
|
|
return item
|
||
|
|
|
||
|
|
@staticmethod
|
||
|
|
def unpack_run(row):
|
||
|
|
data = json.loads(row["data"])
|
||
|
|
data.update(state=row["state"], cancel_requested=bool(row["cancel_requested"]))
|
||
|
|
return data
|
||
|
|
|
||
|
|
def runs(self, script_id=None, limit=100):
|
||
|
|
with self.connect() as db:
|
||
|
|
query = "SELECT * FROM runs"
|
||
|
|
values = []
|
||
|
|
if script_id:
|
||
|
|
query += " WHERE script_id=?"
|
||
|
|
values.append(script_id)
|
||
|
|
query += " ORDER BY rowid DESC LIMIT ?"
|
||
|
|
values.append(limit)
|
||
|
|
return [self.unpack_run(r) for r in db.execute(query, values)]
|
||
|
|
|
||
|
|
def run(self, identity):
|
||
|
|
with self.connect() as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
|
||
|
|
if not row:
|
||
|
|
raise KeyError("运行记录不存在")
|
||
|
|
return self.unpack_run(row)
|
||
|
|
|
||
|
|
def run_for_request(self, request_id):
|
||
|
|
with self.connect() as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE request_id=?", (request_id,)).fetchone()
|
||
|
|
return self.unpack_run(row) if row else None
|
||
|
|
|
||
|
|
def queue(self, script, request_id, ip, execution=None):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
old = db.execute("SELECT * FROM runs WHERE request_id=?", (request_id,)).fetchone()
|
||
|
|
if old:
|
||
|
|
if old["script_id"] != script["id"]:
|
||
|
|
raise ValueError("请求编号已用于其他脚本")
|
||
|
|
return self.unpack_run(old)
|
||
|
|
pending = db.execute("SELECT count(*) FROM runs WHERE state IN ('queued','running')").fetchone()[0]
|
||
|
|
if pending >= 20:
|
||
|
|
raise ValueError("队列已满，请等待已有测试结束")
|
||
|
|
item = {"id": uuid.uuid4().hex, "script_id": script["id"], "script_name": script["name"],
|
||
|
|
"version": script["version"], "sha256": script["sha256"], "state": "queued",
|
||
|
|
"created_at": now(), "started_at": None, "finished_at": None, "exit_code": None,
|
||
|
|
"message": "等待执行", "config": script["config"], "artifacts": [],
|
||
|
|
"synthetic": script.get("synthetic", False), "source_ip": ip,
|
||
|
|
"source_path": script["source_path"], "language": script["language"],
|
||
|
|
"kind": script["kind"], "source_group": script["source_group"]}
|
||
|
|
if execution:
|
||
|
|
item["execution"] = execution
|
||
|
|
db.execute("INSERT INTO runs(id,request_id,script_id,state,data) VALUES(?,?,?,?,?)",
|
||
|
|
(item["id"], request_id, script["id"], "queued", json.dumps(item)))
|
||
|
|
return item
|
||
|
|
|
||
|
|
def claim(self):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
if db.execute("SELECT 1 FROM runs WHERE state='running'").fetchone():
|
||
|
|
return None
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE state='queued' ORDER BY rowid LIMIT 1").fetchone()
|
||
|
|
if not row:
|
||
|
|
return None
|
||
|
|
item = self.unpack_run(row)
|
||
|
|
item.update(started_at=now(), state="running", message="正在启动测试")
|
||
|
|
db.execute("UPDATE runs SET state='running',data=? WHERE id=?", (json.dumps(item), item["id"]))
|
||
|
|
return item
|
||
|
|
|
||
|
|
def finish(self, identity, state, **values):
|
||
|
|
assert state in TERMINAL
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
|
||
|
|
item = self.unpack_run(row)
|
||
|
|
item.update(values, state=state, finished_at=now())
|
||
|
|
db.execute("UPDATE runs SET state=?,data=? WHERE id=?", (state, json.dumps(item), identity))
|
||
|
|
return item
|
||
|
|
|
||
|
|
def running_message(self, identity, message):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE id=? AND state='running'", (identity,)).fetchone()
|
||
|
|
if row:
|
||
|
|
item = self.unpack_run(row)
|
||
|
|
item["message"] = message
|
||
|
|
db.execute("UPDATE runs SET data=? WHERE id=?", (json.dumps(item), identity))
|
||
|
|
|
||
|
|
def update_execution(self, identity, values):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE id=? AND state='running'", (identity,)).fetchone()
|
||
|
|
if not row:
|
||
|
|
raise ValueError("运行状态已变化")
|
||
|
|
item = self.unpack_run(row)
|
||
|
|
item["execution"] = {**item.get("execution", {}), **values}
|
||
|
|
db.execute("UPDATE runs SET data=? WHERE id=?", (json.dumps(item), identity))
|
||
|
|
return item
|
||
|
|
|
||
|
|
def cancel(self, identity):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
row = db.execute("SELECT * FROM runs WHERE id=?", (identity,)).fetchone()
|
||
|
|
if not row:
|
||
|
|
raise KeyError("运行记录不存在")
|
||
|
|
item = self.unpack_run(row)
|
||
|
|
if row["state"] == "queued":
|
||
|
|
item.update(state="cancelled", finished_at=now(), message="已取消排队")
|
||
|
|
db.execute("UPDATE runs SET state='cancelled',cancel_requested=1,data=? WHERE id=?",
|
||
|
|
(json.dumps(item), identity))
|
||
|
|
elif row["state"] == "running":
|
||
|
|
db.execute("UPDATE runs SET cancel_requested=1 WHERE id=?", (identity,))
|
||
|
|
return self.run(identity)
|
||
|
|
|
||
|
|
def heartbeat(self):
|
||
|
|
with self.connect(write=True) as db:
|
||
|
|
db.execute("INSERT OR REPLACE INTO health VALUES('worker',?)", (now(),))
|
||
|
|
|
||
|
|
def health(self):
|
||
|
|
with self.connect() as db:
|
||
|
|
row = db.execute("SELECT value FROM health WHERE key='worker'").fetchone()
|
||
|
|
last = row[0] if row else None
|
||
|
|
online = bool(last and (datetime.now(timezone.utc) - datetime.fromisoformat(last)).total_seconds() < 15)
|
||
|
|
return {"online": online, "last_seen": last}
|