936 lines
32 KiB
Python
936 lines
32 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import sqlite3
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, date, datetime, timedelta, timezone
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
import aiosqlite
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ACTIVE_STATUSES = ("queued", "running")
|
||
MSK = timezone(timedelta(hours=3))
|
||
|
||
|
||
def utc_now() -> str:
|
||
return datetime.now(UTC).isoformat()
|
||
|
||
|
||
def msk_today() -> date:
|
||
return datetime.now(MSK).date()
|
||
|
||
|
||
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
|
||
enhanced_prompt: str | None
|
||
input_file_path: str | None
|
||
output_file_path: str | None
|
||
status: str
|
||
comfy_prompt_id: str | None
|
||
payment_charge_id: str | None
|
||
payment_refunded: bool
|
||
is_premium: bool
|
||
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"],
|
||
enhanced_prompt=row["enhanced_prompt"],
|
||
input_file_path=row["input_file_path"],
|
||
output_file_path=row["output_file_path"],
|
||
status=row["status"],
|
||
comfy_prompt_id=row["comfy_prompt_id"],
|
||
payment_charge_id=row["payment_charge_id"],
|
||
payment_refunded=bool(row["payment_refunded"]),
|
||
is_premium=bool(row["is_premium"]),
|
||
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)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class BlockedRequest:
|
||
id: int
|
||
user_id: int
|
||
timestamp: str
|
||
prompt: str
|
||
reason: str
|
||
mode: str | None
|
||
has_person: bool
|
||
|
||
@classmethod
|
||
def from_row(cls, row: aiosqlite.Row) -> BlockedRequest:
|
||
return cls(
|
||
id=row["id"],
|
||
user_id=row["user_id"],
|
||
timestamp=row["timestamp"],
|
||
prompt=row["prompt"],
|
||
reason=row["reason"],
|
||
mode=row["mode"],
|
||
has_person=bool(row["has_person"]),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class User:
|
||
user_id: int
|
||
username: str | None
|
||
credits: int
|
||
last_daily_credit: str
|
||
referrer_id: int | None
|
||
first_paid_at: str | None
|
||
created_at: str
|
||
is_admin: bool
|
||
is_blocked: bool
|
||
|
||
@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) -> User:
|
||
return cls(
|
||
user_id=int(row["user_id"]),
|
||
username=row["username"],
|
||
credits=int(row["credits"]),
|
||
last_daily_credit=row["last_daily_credit"],
|
||
referrer_id=row["referrer_id"],
|
||
first_paid_at=row["first_paid_at"],
|
||
created_at=row["created_at"],
|
||
is_admin=bool(row["is_admin"]),
|
||
is_blocked=bool(row["is_blocked"]),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PendingPayment:
|
||
payload: str
|
||
user_id: int
|
||
timestamp: str
|
||
mode: str
|
||
prompt: str
|
||
username: str | None
|
||
input_file_path: str | None
|
||
amount: int
|
||
status: str
|
||
is_premium: bool
|
||
telegram_payment_charge_id: str | None
|
||
request_id: int | None
|
||
|
||
@classmethod
|
||
def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
|
||
return cls(
|
||
payload=row["payload"],
|
||
user_id=row["user_id"],
|
||
timestamp=row["timestamp"],
|
||
mode=row["mode"],
|
||
prompt=row["prompt"],
|
||
username=row["username"],
|
||
input_file_path=row["input_file_path"],
|
||
amount=int(row["amount"]),
|
||
status=row["status"],
|
||
is_premium=bool(row["is_premium"]),
|
||
telegram_payment_charge_id=row["telegram_payment_charge_id"],
|
||
request_id=row["request_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,
|
||
enhanced_prompt TEXT,
|
||
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,
|
||
payment_charge_id TEXT,
|
||
payment_refunded INTEGER NOT NULL DEFAULT 0,
|
||
is_premium INTEGER NOT NULL DEFAULT 0,
|
||
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');
|
||
|
||
CREATE TABLE IF NOT EXISTS blocked_requests (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
user_id INTEGER NOT NULL,
|
||
timestamp TEXT NOT NULL,
|
||
prompt TEXT NOT NULL,
|
||
reason TEXT NOT NULL,
|
||
mode TEXT,
|
||
has_person INTEGER NOT NULL DEFAULT 0
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_blocked_user_time
|
||
ON blocked_requests(user_id, id);
|
||
|
||
CREATE TABLE IF NOT EXISTS pending_payments (
|
||
payload TEXT PRIMARY KEY,
|
||
user_id INTEGER NOT NULL,
|
||
timestamp TEXT NOT NULL,
|
||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
||
prompt TEXT NOT NULL,
|
||
username TEXT,
|
||
input_file_path TEXT,
|
||
amount INTEGER NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'pending'
|
||
CHECK (status IN ('pending', 'paid', 'cancelled', 'refunded')),
|
||
is_premium INTEGER NOT NULL DEFAULT 0,
|
||
telegram_payment_charge_id TEXT,
|
||
request_id INTEGER
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_pending_user_time
|
||
ON pending_payments(user_id, timestamp);
|
||
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
user_id INTEGER PRIMARY KEY,
|
||
username TEXT,
|
||
credits INTEGER NOT NULL DEFAULT 3,
|
||
last_daily_credit TEXT NOT NULL,
|
||
referrer_id INTEGER,
|
||
first_paid_at TEXT,
|
||
created_at TEXT NOT NULL,
|
||
is_admin INTEGER NOT NULL DEFAULT 0,
|
||
is_blocked INTEGER NOT NULL DEFAULT 0
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_users_referrer
|
||
ON users(referrer_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS app_settings (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT
|
||
);
|
||
"""
|
||
|
||
|
||
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")
|
||
if "enhanced_prompt" not in columns:
|
||
await conn.execute("ALTER TABLE requests ADD COLUMN enhanced_prompt TEXT")
|
||
if "payment_charge_id" not in columns:
|
||
await conn.execute("ALTER TABLE requests ADD COLUMN payment_charge_id TEXT")
|
||
if "payment_refunded" not in columns:
|
||
await conn.execute(
|
||
"ALTER TABLE requests ADD COLUMN payment_refunded INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
if "is_premium" not in columns:
|
||
await conn.execute(
|
||
"ALTER TABLE requests ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
async with conn.execute("PRAGMA table_info(pending_payments)") as cursor:
|
||
pending_columns = {row["name"] async for row in cursor}
|
||
if "is_premium" not in pending_columns:
|
||
await conn.execute(
|
||
"ALTER TABLE pending_payments ADD COLUMN is_premium INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
async with conn.execute("PRAGMA table_info(users)") as cursor:
|
||
user_columns = {row["name"] async for row in cursor}
|
||
if "is_admin" not in user_columns:
|
||
await conn.execute(
|
||
"ALTER TABLE users ADD COLUMN is_admin INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
if "is_blocked" not in user_columns:
|
||
await conn.execute(
|
||
"ALTER TABLE users ADD COLUMN is_blocked INTEGER NOT NULL DEFAULT 0"
|
||
)
|
||
await conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO users
|
||
(user_id, username, credits, last_daily_credit,
|
||
referrer_id, first_paid_at, created_at)
|
||
SELECT user_id, NULL, 3, ?, NULL, NULL, ?
|
||
FROM (SELECT DISTINCT user_id FROM requests)
|
||
""",
|
||
(msk_today().isoformat(), utc_now()),
|
||
)
|
||
|
||
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,
|
||
payment_charge_id: str | None = None,
|
||
is_premium: bool = False,
|
||
) -> 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, payment_charge_id,
|
||
is_premium)
|
||
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?, ?)
|
||
""",
|
||
(
|
||
user_id,
|
||
username,
|
||
utc_now(),
|
||
prompt_text,
|
||
input_file_path,
|
||
mode,
|
||
seed,
|
||
payment_charge_id,
|
||
int(is_premium),
|
||
),
|
||
)
|
||
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 get_user(self, user_id: int) -> User | None:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"SELECT * FROM users WHERE user_id = ?", (user_id,)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return User.from_row(row) if row is not None else None
|
||
|
||
async def ensure_user(
|
||
self, user_id: int, username: str | None = None
|
||
) -> User:
|
||
conn = self._connection()
|
||
today = msk_today().isoformat()
|
||
async with self._write_lock:
|
||
await conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO users
|
||
(user_id, username, credits, last_daily_credit, created_at)
|
||
VALUES (?, ?, 3, ?, ?)
|
||
""",
|
||
(user_id, username, today, utc_now()),
|
||
)
|
||
if username:
|
||
await conn.execute(
|
||
"UPDATE users SET username = ? WHERE user_id = ?",
|
||
(username, user_id),
|
||
)
|
||
await conn.commit()
|
||
user = await self.get_user(user_id)
|
||
if user is None:
|
||
raise RuntimeError("Не удалось прочитать пользователя")
|
||
return user
|
||
|
||
async def seed_admins(self, user_ids: frozenset[int] | set[int]) -> None:
|
||
if not user_ids:
|
||
return
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
async with conn.execute(
|
||
"SELECT value FROM app_settings WHERE key = 'admins_seeded'"
|
||
) as cursor:
|
||
if await cursor.fetchone() is not None:
|
||
return
|
||
today = msk_today().isoformat()
|
||
for user_id in user_ids:
|
||
await conn.execute(
|
||
"""
|
||
INSERT OR IGNORE INTO users
|
||
(user_id, username, credits, last_daily_credit, created_at)
|
||
VALUES (?, NULL, 3, ?, ?)
|
||
""",
|
||
(user_id, today, utc_now()),
|
||
)
|
||
await conn.execute(
|
||
"UPDATE users SET is_admin = 1 WHERE user_id = ?",
|
||
(user_id,),
|
||
)
|
||
await conn.execute(
|
||
"INSERT OR REPLACE INTO app_settings (key, value)"
|
||
" VALUES ('admins_seeded', '1')"
|
||
)
|
||
await conn.commit()
|
||
|
||
async def list_users_admin(self, limit: int = 500) -> list[User]:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"SELECT * FROM users ORDER BY user_id LIMIT ?", (limit,)
|
||
) as cursor:
|
||
rows = await cursor.fetchall()
|
||
return [User.from_row(row) for row in rows]
|
||
|
||
async def set_admin(self, user_id: int, enabled: bool) -> bool:
|
||
return await self._update_user_flag(
|
||
"is_admin", user_id, enabled
|
||
)
|
||
|
||
async def set_blocked(self, user_id: int, enabled: bool) -> bool:
|
||
return await self._update_user_flag(
|
||
"is_blocked", user_id, enabled
|
||
)
|
||
|
||
async def set_credits(self, user_id: int, credits: int) -> bool:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"UPDATE users SET credits = ? WHERE user_id = ?",
|
||
(max(0, credits), user_id),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def _update_user_flag(
|
||
self, column: str, user_id: int, enabled: bool
|
||
) -> bool:
|
||
if column not in {"is_admin", "is_blocked"}:
|
||
raise ValueError("Недопустимое поле")
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
f"UPDATE users SET {column} = ? WHERE user_id = ?",
|
||
(int(enabled), user_id),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def add_credits(self, user_id: int, amount: int) -> int | None:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"UPDATE users SET credits = credits + ? WHERE user_id = ?",
|
||
(amount, user_id),
|
||
)
|
||
await conn.commit()
|
||
if cursor.rowcount == 0:
|
||
return None
|
||
user = await self.get_user(user_id)
|
||
return user.credits if user is not None else None
|
||
|
||
async def consume_credit(self, user_id: int) -> int | None:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"UPDATE users SET credits = credits - 1"
|
||
" WHERE user_id = ? AND credits > 0",
|
||
(user_id,),
|
||
)
|
||
await conn.commit()
|
||
if cursor.rowcount == 0:
|
||
return None
|
||
user = await self.get_user(user_id)
|
||
return user.credits if user is not None else None
|
||
|
||
async def set_referrer(self, user_id: int, referrer_id: int) -> bool:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET referrer_id = ?
|
||
WHERE user_id = ? AND referrer_id IS NULL AND user_id != ?
|
||
""",
|
||
(referrer_id, user_id, referrer_id),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def mark_first_payment(self, user_id: int) -> bool:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"UPDATE users SET first_paid_at = ?"
|
||
" WHERE user_id = ? AND first_paid_at IS NULL",
|
||
(utc_now(), user_id),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def accrue_daily_credits(self, today: date | None = None) -> int:
|
||
conn = self._connection()
|
||
current = today or msk_today()
|
||
current_text = current.isoformat()
|
||
async with self._write_lock:
|
||
async with conn.execute(
|
||
"SELECT user_id, last_daily_credit FROM users"
|
||
" WHERE last_daily_credit < ?",
|
||
(current_text,),
|
||
) as cursor:
|
||
rows = await cursor.fetchall()
|
||
for row in rows:
|
||
try:
|
||
last = date.fromisoformat(row["last_daily_credit"])
|
||
days = max(1, (current - last).days)
|
||
except ValueError:
|
||
days = 1
|
||
await conn.execute(
|
||
"UPDATE users SET credits = credits + ?,"
|
||
" last_daily_credit = ? WHERE user_id = ?",
|
||
(days, current_text, int(row["user_id"])),
|
||
)
|
||
await conn.commit()
|
||
return len(rows)
|
||
|
||
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()
|
||
request = await self.get_request(request_id)
|
||
if request is None:
|
||
return 0
|
||
if request.is_premium:
|
||
query = """
|
||
SELECT COUNT(*) AS total FROM requests
|
||
WHERE is_premium = 1 AND status IN ('queued', 'running')
|
||
AND id < ?
|
||
"""
|
||
params = (request_id,)
|
||
else:
|
||
query = """
|
||
SELECT COUNT(*) AS total FROM requests
|
||
WHERE status IN ('queued', 'running')
|
||
AND (is_premium = 1 OR (is_premium = 0 AND id < ?))
|
||
"""
|
||
params = (request_id,)
|
||
async with conn.execute(query, params) as cursor:
|
||
row = await cursor.fetchone()
|
||
return int(row["total"]) if row is not None else 0
|
||
|
||
async def queue_counts(self) -> dict[str, int]:
|
||
conn = self._connection()
|
||
counts = {"premium": 0, "normal": 0}
|
||
async with conn.execute(
|
||
"""
|
||
SELECT is_premium, COUNT(*) AS total FROM requests
|
||
WHERE status IN ('queued', 'running')
|
||
GROUP BY is_premium
|
||
"""
|
||
) as cursor:
|
||
async for row in cursor:
|
||
key = "premium" if row["is_premium"] else "normal"
|
||
counts[key] = int(row["total"])
|
||
return counts
|
||
|
||
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_enhanced_prompt(self, request_id: int, prompt: str) -> None:
|
||
await self._update(
|
||
"UPDATE requests SET enhanced_prompt = ? WHERE id = ?",
|
||
(prompt, request_id),
|
||
)
|
||
|
||
async def set_refunded(self, request_id: int) -> None:
|
||
await self._update(
|
||
"UPDATE requests SET payment_refunded = 1 WHERE id = ?",
|
||
(request_id,),
|
||
)
|
||
|
||
async def create_pending_payment(
|
||
self,
|
||
*,
|
||
payload: str,
|
||
user_id: int,
|
||
mode: str,
|
||
prompt: str,
|
||
username: str | None,
|
||
input_file_path: str | None,
|
||
amount: int,
|
||
is_premium: bool = False,
|
||
) -> PendingPayment:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO pending_payments
|
||
(payload, user_id, timestamp, mode, prompt, username,
|
||
input_file_path, amount, status, is_premium)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?)
|
||
""",
|
||
(
|
||
payload,
|
||
user_id,
|
||
utc_now(),
|
||
mode,
|
||
prompt,
|
||
username,
|
||
input_file_path,
|
||
amount,
|
||
int(is_premium),
|
||
),
|
||
)
|
||
await conn.commit()
|
||
pending = await self.get_pending_payment(payload)
|
||
if pending is None:
|
||
raise RuntimeError("Не удалось прочитать счёт")
|
||
return pending
|
||
|
||
async def get_pending_payment(self, payload: str) -> PendingPayment | None:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"SELECT * FROM pending_payments WHERE payload = ?", (payload,)
|
||
) as cursor:
|
||
row = await cursor.fetchone()
|
||
return PendingPayment.from_row(row) if row is not None else None
|
||
|
||
async def set_pending_choice(
|
||
self, payload: str, is_premium: bool, amount: int
|
||
) -> bool:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE pending_payments
|
||
SET is_premium = ?, amount = ?
|
||
WHERE payload = ? AND status = 'pending' AND amount = 0
|
||
""",
|
||
(int(is_premium), amount, payload),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE pending_payments
|
||
SET status = 'paid', telegram_payment_charge_id = ?
|
||
WHERE payload = ? AND status = 'pending'
|
||
""",
|
||
(charge_id, payload),
|
||
)
|
||
await conn.commit()
|
||
return cursor.rowcount > 0
|
||
|
||
async def mark_pending_refunded(self, payload: str) -> None:
|
||
await self._update(
|
||
"UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
|
||
(payload,),
|
||
)
|
||
|
||
async def mark_pending_cancelled(self, payload: str) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE pending_payments SET status = 'cancelled'
|
||
WHERE payload = ? AND status = 'pending'
|
||
""",
|
||
(payload,),
|
||
)
|
||
|
||
async def attach_pending_request(self, payload: str, request_id: int) -> None:
|
||
await self._update(
|
||
"UPDATE pending_payments SET request_id = ? WHERE payload = ?",
|
||
(request_id, payload),
|
||
)
|
||
|
||
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 log_blocked(
|
||
self,
|
||
*,
|
||
user_id: int,
|
||
prompt: str,
|
||
reason: str,
|
||
mode: str | None = None,
|
||
has_person: bool = False,
|
||
) -> int:
|
||
conn = self._connection()
|
||
async with self._write_lock:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO blocked_requests
|
||
(user_id, timestamp, prompt, reason, mode, has_person)
|
||
VALUES (?, ?, ?, ?, ?, ?)
|
||
""",
|
||
(user_id, utc_now(), prompt, reason, mode, int(has_person)),
|
||
)
|
||
await conn.commit()
|
||
return int(cursor.lastrowid)
|
||
|
||
async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]:
|
||
conn = self._connection()
|
||
async with conn.execute(
|
||
"SELECT * FROM blocked_requests ORDER BY id DESC LIMIT ?", (limit,)
|
||
) as cursor:
|
||
rows = await cursor.fetchall()
|
||
return [BlockedRequest.from_row(row) for row in rows]
|
||
|
||
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()
|