390 lines
12 KiB
Python
390 lines
12 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import sqlite3
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, datetime
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import aiosqlite
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ACTIVE_STATUSES = ("queued", "running")
|
||
|
||
|
||
def utc_now() -> str:
|
||
return datetime.now(UTC).isoformat()
|
||
|
||
|
||
class UserAlreadyQueued(RuntimeError):
|
||
def __init__(self, user_id: int) -> None:
|
||
super().__init__(f"У пользователя {user_id} уже есть активный запрос")
|
||
self.user_id = user_id
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Request:
|
||
id: int
|
||
user_id: int
|
||
username: str | None
|
||
timestamp: str
|
||
prompt_text: str
|
||
input_file_path: str | None
|
||
output_file_path: str | None
|
||
status: str
|
||
comfy_prompt_id: str | None
|
||
mode: str
|
||
seed: int | None
|
||
error_text: str | None
|
||
finished_at: str | None
|
||
|
||
@property
|
||
def display_user(self) -> str:
|
||
return f"@{self.username}" if self.username else str(self.user_id)
|
||
|
||
@classmethod
|
||
def from_row(cls, row: aiosqlite.Row) -> Request:
|
||
return cls(
|
||
id=row["id"],
|
||
user_id=row["user_id"],
|
||
username=row["username"],
|
||
timestamp=row["timestamp"],
|
||
prompt_text=row["prompt_text"],
|
||
input_file_path=row["input_file_path"],
|
||
output_file_path=row["output_file_path"],
|
||
status=row["status"],
|
||
comfy_prompt_id=row["comfy_prompt_id"],
|
||
mode=row["mode"],
|
||
seed=row["seed"],
|
||
error_text=row["error_text"],
|
||
finished_at=row["finished_at"],
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class UserStat:
|
||
user_id: int
|
||
username: str | None
|
||
total: int
|
||
last_seen: str
|
||
|
||
@property
|
||
def display_user(self) -> str:
|
||
return f"@{self.username}" if self.username else str(self.user_id)
|
||
|
||
|
||
SCHEMA = """
|
||
CREATE TABLE IF NOT EXISTS requests (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
user_id INTEGER NOT NULL,
|
||
username TEXT,
|
||
timestamp TEXT NOT NULL,
|
||
prompt_text TEXT NOT NULL,
|
||
input_file_path TEXT,
|
||
output_file_path TEXT,
|
||
status TEXT NOT NULL DEFAULT 'queued'
|
||
CHECK (status IN ('queued', 'running', 'done', 'failed')),
|
||
comfy_prompt_id TEXT,
|
||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
||
seed INTEGER,
|
||
error_text TEXT,
|
||
finished_at TEXT
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_requests_user_status
|
||
ON requests(user_id, status);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_requests_status_id
|
||
ON requests(status, id);
|
||
|
||
CREATE UNIQUE INDEX IF NOT EXISTS uq_requests_one_active
|
||
ON requests(user_id) WHERE status IN ('queued', 'running');
|
||
"""
|
||
|
||
|
||
class Database:
|
||
def __init__(self, path: Path) -> None:
|
||
self._path = path
|
||
self._conn: aiosqlite.Connection | None = None
|
||
self._write_lock = asyncio.Lock()
|
||
|
||
@property
|
||
def path(self) -> Path:
|
||
return self._path
|
||
|
||
async def connect(self) -> None:
|
||
self._path.parent.mkdir(parents=True, exist_ok=True)
|
||
self._conn = await aiosqlite.connect(self._path)
|
||
self._conn.row_factory = aiosqlite.Row
|
||
await self._conn.execute("PRAGMA journal_mode=WAL")
|
||
await self._conn.execute("PRAGMA busy_timeout=5000")
|
||
await self._conn.execute("PRAGMA foreign_keys=ON")
|
||
await self._conn.executescript(SCHEMA)
|
||
await self._migrate()
|
||
await self._conn.commit()
|
||
|
||
async def _migrate(self) -> None:
|
||
conn = self._connection()
|
||
async with conn.execute("PRAGMA table_info(requests)") as cursor:
|
||
columns = {row["name"] async for row in cursor}
|
||
if "username" not in columns:
|
||
await conn.execute("ALTER TABLE requests ADD COLUMN username TEXT")
|
||
|
||
async def close(self) -> None:
|
||
if self._conn is not None:
|
||
await self._conn.close()
|
||
self._conn = None
|
||
|
||
def _connection(self) -> aiosqlite.Connection:
|
||
if self._conn is None:
|
||
raise RuntimeError("База данных не подключена")
|
||
return self._conn
|
||
|
||
async def create_request(
|
||
self,
|
||
*,
|
||
user_id: int,
|
||
username: str | None = None,
|
||
mode: str,
|
||
prompt_text: str,
|
||
input_file_path: str | None,
|
||
seed: int,
|
||
) -> Request:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
try:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO requests
|
||
(user_id, username, timestamp, prompt_text,
|
||
input_file_path, status, mode, seed)
|
||
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?)
|
||
""",
|
||
(
|
||
user_id,
|
||
username,
|
||
utc_now(),
|
||
prompt_text,
|
||
input_file_path,
|
||
mode,
|
||
seed,
|
||
),
|
||
)
|
||
request_id = cursor.lastrowid
|
||
await conn.commit()
|
||
except sqlite3.IntegrityError as exc:
|
||
await conn.rollback()
|
||
if "UNIQUE" in str(exc).upper():
|
||
raise UserAlreadyQueued(user_id) from exc
|
||
raise
|
||
request = await self.get_request(int(request_id))
|
||
if request is None:
|
||
raise RuntimeError("Не удалось прочитать созданный запрос")
|
||
return request
|
||
|
||
async def get_request(self, request_id: int) -> Request | None:
|
||
conn = self._connection()
|
||
async with conn.execute("SELECT * FROM requests WHERE id = ?", (request_id,)) as cursor:
|
||
row = await cursor.fetchone()
|
||
return Request.from_row(row) if row is not None else None
|
||
|
||
async def get_active_by_user(self, user_id: int) -> Request | None:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"""
|
||
SELECT * FROM requests
|
||
WHERE user_id = ? AND status IN ('queued', 'running')
|
||
ORDER BY id DESC LIMIT 1
|
||
""",
|
||
(user_id,),
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return Request.from_row(row) if row is not None else None
|
||
|
||
async def list_queued(self) -> list[Request]:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"SELECT * FROM requests WHERE status = 'queued' ORDER BY id"
|
||
) as cursor:
|
||
rows = await cursor.fetchall()
|
||
return [Request.from_row(row) for row in rows]
|
||
|
||
@staticmethod
|
||
def _filters(
|
||
user_id: int | None,
|
||
status: str | None,
|
||
mode: str | None,
|
||
) -> tuple[str, list[Any]]:
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
if user_id is not None:
|
||
clauses.append("user_id = ?")
|
||
params.append(user_id)
|
||
if status:
|
||
clauses.append("status = ?")
|
||
params.append(status)
|
||
if mode:
|
||
clauses.append("mode = ?")
|
||
params.append(mode)
|
||
where = f" WHERE {' AND '.join(clauses)}" if clauses else ""
|
||
return where, params
|
||
|
||
async def list_requests(
|
||
self,
|
||
*,
|
||
user_id: int | None = None,
|
||
status: str | None = None,
|
||
mode: str | None = None,
|
||
limit: int = 50,
|
||
offset: int = 0,
|
||
) -> list[Request]:
|
||
conn = self._connection()
|
||
where, params = self._filters(user_id, status, mode)
|
||
query = (
|
||
f"SELECT * FROM requests{where} ORDER BY id DESC LIMIT ? OFFSET ?"
|
||
)
|
||
async with conn.execute(query, (*params, limit, offset)) as cursor:
|
||
rows = await cursor.fetchall()
|
||
return [Request.from_row(row) for row in rows]
|
||
|
||
async def count_requests(
|
||
self,
|
||
*,
|
||
user_id: int | None = None,
|
||
status: str | None = None,
|
||
mode: str | None = None,
|
||
) -> int:
|
||
conn = self._connection()
|
||
where, params = self._filters(user_id, status, mode)
|
||
async with conn.execute(
|
||
f"SELECT COUNT(*) AS total FROM requests{where}", params
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return int(row["total"]) if row is not None else 0
|
||
|
||
async def stats(self) -> dict[str, Any]:
|
||
conn = self._connection()
|
||
by_status: dict[str, int] = {}
|
||
async with conn.execute(
|
||
"SELECT status, COUNT(*) AS total FROM requests GROUP BY status"
|
||
) as cursor:
|
||
async for row in cursor:
|
||
by_status[str(row["status"])] = int(row["total"])
|
||
async with conn.execute("SELECT COUNT(DISTINCT user_id) AS total FROM requests") as cursor:
|
||
row = await cursor.fetchone()
|
||
return {
|
||
"total": sum(by_status.values()),
|
||
"users": int(row["total"]) if row is not None else 0,
|
||
"by_status": by_status,
|
||
}
|
||
|
||
async def list_users(self, limit: int = 100) -> list[UserStat]:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"""
|
||
SELECT
|
||
user_id,
|
||
COUNT(*) AS total,
|
||
MAX(timestamp) AS last_seen,
|
||
(
|
||
SELECT username FROM requests AS r2
|
||
WHERE r2.user_id = requests.user_id AND r2.username IS NOT NULL
|
||
ORDER BY r2.id DESC LIMIT 1
|
||
) AS username
|
||
FROM requests
|
||
GROUP BY user_id
|
||
ORDER BY last_seen DESC
|
||
LIMIT ?
|
||
""",
|
||
(limit,),
|
||
) as cursor:
|
||
rows = await cursor.fetchall()
|
||
return [
|
||
UserStat(
|
||
user_id=int(row["user_id"]),
|
||
username=row["username"],
|
||
total=int(row["total"]),
|
||
last_seen=str(row["last_seen"]),
|
||
)
|
||
for row in rows
|
||
]
|
||
|
||
async def count_ahead(self, request_id: int) -> int:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"""
|
||
SELECT COUNT(*) AS total FROM requests
|
||
WHERE id < ? AND status IN ('queued', 'running')
|
||
""",
|
||
(request_id,),
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return int(row["total"]) if row is not None else 0
|
||
|
||
async def set_running(self, request_id: int, comfy_prompt_id: str | None = None) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE requests
|
||
SET status = 'running',
|
||
comfy_prompt_id = COALESCE(?, comfy_prompt_id),
|
||
error_text = NULL
|
||
WHERE id = ?
|
||
""",
|
||
(comfy_prompt_id, request_id),
|
||
)
|
||
|
||
async def set_comfy_prompt_id(self, request_id: int, comfy_prompt_id: str) -> None:
|
||
await self._update(
|
||
"UPDATE requests SET comfy_prompt_id = ? WHERE id = ?",
|
||
(comfy_prompt_id, request_id),
|
||
)
|
||
|
||
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE requests
|
||
SET status = 'done',
|
||
output_file_path = ?,
|
||
error_text = NULL,
|
||
finished_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(output_file_path, utc_now(), request_id),
|
||
)
|
||
|
||
async def set_failed(self, request_id: int, error_text: str) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE requests
|
||
SET status = 'failed',
|
||
error_text = ?,
|
||
finished_at = ?
|
||
WHERE id = ?
|
||
""",
|
||
(error_text[:1000], utc_now(), request_id),
|
||
)
|
||
|
||
async def recover(self, reason: str = "прервано перезапуском") -> list[Request]:
|
||
"""Помечает running как failed и возвращает queued для повторной постановки."""
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
await conn.execute(
|
||
"""
|
||
UPDATE requests
|
||
SET status = 'failed', error_text = ?, finished_at = ?
|
||
WHERE status = 'running'
|
||
""",
|
||
(reason, utc_now()),
|
||
)
|
||
await conn.commit()
|
||
return await self.list_queued()
|
||
|
||
async def _update(self, sql: str, params: tuple[Any, ...]) -> None:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
await conn.execute(sql, params)
|
||
await conn.commit()
|