196 lines
8.6 KiB
Python
196 lines
8.6 KiB
Python
"""Бот-соавтор: правит редактируемые ячейки листа по WebSocket с заданной частотой.
|
||
|
||
CLI: python -m perfkit.coauthor --form-id N --sheet S [--direction D] --rate 1|5 --duration SEC
|
||
[--user rf_coauthor] [--scale 1] [--token T]
|
||
|
||
Протокол (api/src/api/v1/websocket.py, web/.../RealtimeContext.js):
|
||
ws://HOST/api/v1/ws/form/{id}/sheet/{sheet}?direction=D
|
||
-> {"event":"user_login","data":{"token":...}}
|
||
-> {"event":"cell_edit_start","data":{"line_id","column"}} (эхо отправителю, блокировка ячейки)
|
||
-> {"event":"cell_updated","data":{"line_id","line_id_code","column","value"}}
|
||
(эхо отправителю с "is_self": true; при отказе то же событие с полем "error")
|
||
-> {"event":"cell_edit_end","data":{"line_id","column"}}
|
||
Редактируемые колонки берутся из строки META листа (GET sheet: data.editable, роль EXECUTOR_RF),
|
||
правятся числовые колонки у строк типа INPUT.
|
||
|
||
Коды выхода: 0 - были подтверждённые эхом правки; 2 - нет эха первой правки за --ack-timeout (в т.ч. неверный токен
|
||
или закрытое соединение); 3 - ошибка подготовки (вход по HTTP, нет редактируемых ячеек).
|
||
Печатает "READY" после первого эха и в конце последней строкой JSON {"sent","echoed","errors"}.
|
||
SIGTERM/SIGINT завершают работу штатно (итог печатается).
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import random
|
||
import signal
|
||
import sys
|
||
import time
|
||
from urllib.parse import urlencode
|
||
|
||
import httpx
|
||
import websockets
|
||
|
||
from . import RESULTS_DIR
|
||
from .manifest import load as load_manifest
|
||
|
||
API = os.environ.get("PERF_API", "http://localhost:8000")
|
||
TEXT_HINTS = ("comment", "justification", "date", "note") # нечисловые колонки не трогаем
|
||
|
||
|
||
def log(msg: str) -> None:
|
||
print(msg, file=sys.stderr, flush=True)
|
||
|
||
|
||
def pick_cells(rows: list[dict]) -> list[tuple[int, str]]:
|
||
"""(line_id, column) для INPUT-строк в числовых редактируемых колонках META."""
|
||
meta = next((r for r in rows if r.get("row_type") == "META"), None)
|
||
if not meta:
|
||
return []
|
||
cols = [e["column"] for e in meta["data"].get("editable", []) if not any(h in e["column"] for h in TEXT_HINTS)]
|
||
cells = []
|
||
for r in rows:
|
||
if r.get("row_type") != "INPUT":
|
||
continue
|
||
d = r.get("data", {})
|
||
lid = d.get("line_id") or d.get("id")
|
||
if lid is None:
|
||
continue
|
||
cells.extend((lid, c) for c in cols)
|
||
return cells
|
||
|
||
|
||
async def fetch_cells(args) -> list[tuple[int, str]]:
|
||
m = load_manifest(RESULTS_DIR / f"manifest-s{args.scale}.json")
|
||
u = m.users[args.user]
|
||
async with httpx.AsyncClient(base_url=API + "/api/v1", timeout=120) as c:
|
||
r = await c.post("/auth/login", json={"username": u.username, "password": u.password})
|
||
r.raise_for_status()
|
||
tok = r.json()["access_token"]
|
||
params = {"direction": args.direction} if args.direction else None
|
||
r = await c.get(f"/form/{args.form_id}/sheet/{args.sheet}", params=params, headers={"Authorization": f"Bearer {tok}"})
|
||
r.raise_for_status()
|
||
args._http_token = tok
|
||
return pick_cells(r.json()["result"])
|
||
|
||
|
||
def ws_url(args) -> str:
|
||
base = API.replace("http://", "ws://").replace("https://", "wss://")
|
||
q = "?" + urlencode({"direction": args.direction}) if args.direction else ""
|
||
return f"{base}/api/v1/ws/form/{args.form_id}/sheet/{args.sheet}{q}"
|
||
|
||
|
||
async def run(args) -> int:
|
||
t0 = time.monotonic()
|
||
try:
|
||
cells = await fetch_cells(args)
|
||
except Exception as e: # noqa: BLE001
|
||
log(f"coauthor: подготовка не удалась: {e!r}")
|
||
return 3
|
||
if not cells:
|
||
log("coauthor: нет редактируемых ячеек для роли пользователя")
|
||
return 3
|
||
|
||
st = {"sent": 0, "echoed": 0, "errors": 0}
|
||
first_echo = asyncio.Event()
|
||
stop = asyncio.Event()
|
||
loop = asyncio.get_running_loop()
|
||
for sig in (signal.SIGTERM, signal.SIGINT):
|
||
loop.add_signal_handler(sig, stop.set)
|
||
|
||
token = args.token if args.token is not None else args._http_token
|
||
try:
|
||
ws = await asyncio.wait_for(websockets.connect(ws_url(args), max_size=None), 5)
|
||
except Exception as e: # noqa: BLE001
|
||
log(f"coauthor: WS недоступен: {e!r}")
|
||
return 2
|
||
|
||
async def reader():
|
||
try:
|
||
async for raw in ws:
|
||
msg = json.loads(raw)
|
||
if msg.get("event") != "cell_updated" or not msg.get("is_self"):
|
||
continue
|
||
if msg.get("error"):
|
||
st["errors"] += 1
|
||
if st["errors"] <= 3:
|
||
log(f"coauthor: отказ сервера: {msg['error']}")
|
||
continue
|
||
st["echoed"] += 1
|
||
if not first_echo.is_set():
|
||
first_echo.set()
|
||
print("READY", flush=True)
|
||
except websockets.ConnectionClosed as e:
|
||
log(f"coauthor: соединение закрыто: {e.rcvd}")
|
||
finally:
|
||
stop.set()
|
||
|
||
rd = asyncio.create_task(reader())
|
||
await ws.send(json.dumps({"event": "user_login", "data": {"token": token}}))
|
||
|
||
async def edit():
|
||
lid, col = random.choice(cells)
|
||
key = {"line_id": lid, "column": col}
|
||
await ws.send(json.dumps({"event": "cell_edit_start", "data": key}))
|
||
await ws.send(json.dumps({"event": "cell_updated", "data": {**key, "line_id_code": lid, "value": random.randint(1000, 999_999)}}))
|
||
st["sent"] += 1
|
||
await ws.send(json.dumps({"event": "cell_edit_end", "data": key}))
|
||
|
||
rc = 0
|
||
try:
|
||
# фаза 1: до первого эха, не дольше ack_timeout с момента старта процесса
|
||
while not first_echo.is_set():
|
||
if stop.is_set() or time.monotonic() - t0 > args.ack_timeout:
|
||
log("coauthor: нет эха первой правки (неверный токен, нет прав или окно закрыто)")
|
||
rc = 2
|
||
break
|
||
await edit()
|
||
try:
|
||
await asyncio.wait_for(first_echo.wait(), 0.7)
|
||
except asyncio.TimeoutError:
|
||
pass
|
||
if rc == 0:
|
||
# фаза 2: ровная частота rate правок/с до конца duration
|
||
period = 1.0 / args.rate
|
||
end = time.monotonic() + args.duration
|
||
nxt = time.monotonic() + period
|
||
while not stop.is_set() and time.monotonic() < end:
|
||
await asyncio.sleep(max(0.0, min(nxt, end) - time.monotonic()))
|
||
if time.monotonic() >= end or stop.is_set():
|
||
break
|
||
await edit()
|
||
nxt += period
|
||
await asyncio.sleep(0.5) # дождаться эха последних правок
|
||
if st["echoed"] == 0:
|
||
rc = 2
|
||
except websockets.ConnectionClosed:
|
||
rc = 2 if st["echoed"] == 0 else 0
|
||
finally:
|
||
try:
|
||
await ws.close()
|
||
except Exception: # noqa: BLE001
|
||
pass
|
||
rd.cancel()
|
||
print(json.dumps(st), flush=True)
|
||
return rc
|
||
|
||
|
||
def main(argv=None) -> int:
|
||
ap = argparse.ArgumentParser(description=__doc__.split("\n")[0])
|
||
ap.add_argument("--form-id", type=int, required=True)
|
||
ap.add_argument("--sheet", required=True)
|
||
ap.add_argument("--direction", default=None)
|
||
ap.add_argument("--rate", type=float, default=1.0, help="правок в секунду (1 или 5)")
|
||
ap.add_argument("--duration", type=float, default=10.0, help="секунд работы; SIGTERM завершает раньше")
|
||
ap.add_argument("--user", default="rf_coauthor", help="ключ пользователя в манифесте")
|
||
ap.add_argument("--scale", type=int, default=int(os.environ.get("PERF_SCALE", 1)))
|
||
ap.add_argument("--token", default=None, help="токен для WS вместо полученного входом (проверка неверного токена)")
|
||
ap.add_argument("--ack-timeout", type=float, default=5.0, help="сколько ждать эха первой правки")
|
||
return asyncio.run(run(ap.parse_args(argv)))
|
||
|
||
|
||
if __name__ == "__main__":
|
||
sys.exit(main())
|