1170 lines
40 KiB
Python
1170 lines
40 KiB
Python
from __future__ import annotations
|
||
|
||
import logging
|
||
from collections.abc import Mapping
|
||
from dataclasses import dataclass
|
||
from datetime import UTC, date, datetime, timedelta, timezone
|
||
from typing import Any
|
||
|
||
from psycopg import AsyncConnection, errors
|
||
from psycopg.rows import dict_row
|
||
from psycopg_pool import AsyncConnectionPool
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
ACTIVE_STATUSES = ("queued", "running")
|
||
MSK = timezone(timedelta(hours=3))
|
||
SCHEMA_LOCK_KEY = 0x4D50414B4F
|
||
|
||
|
||
def utc_now() -> datetime:
|
||
return datetime.now(UTC)
|
||
|
||
|
||
def msk_today() -> date:
|
||
return datetime.now(MSK).date()
|
||
|
||
|
||
def _as_iso(value: Any) -> str | None:
|
||
if value is None:
|
||
return None
|
||
if isinstance(value, datetime):
|
||
return value.isoformat()
|
||
if isinstance(value, date):
|
||
return value.isoformat()
|
||
return str(value)
|
||
|
||
|
||
def _normalize_dsn(dsn: str) -> str:
|
||
value = dsn.strip()
|
||
for prefix in (
|
||
"postgresql+psycopg://",
|
||
"postgresql+asyncpg://",
|
||
"postgres+psycopg://",
|
||
):
|
||
if value.startswith(prefix):
|
||
return "postgresql://" + value[len(prefix) :]
|
||
return value
|
||
|
||
|
||
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
|
||
context_base_prompt: str | None
|
||
context_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
|
||
hidden: 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: Mapping[str, Any]) -> Request:
|
||
return cls(
|
||
id=int(row["id"]),
|
||
user_id=int(row["user_id"]),
|
||
username=row["username"],
|
||
timestamp=_as_iso(row["timestamp"]) or "",
|
||
prompt_text=row["prompt_text"],
|
||
enhanced_prompt=row["enhanced_prompt"],
|
||
context_base_prompt=row["context_base_prompt"],
|
||
context_prompt=row["context_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"]),
|
||
hidden=bool(row["hidden"]),
|
||
mode=row["mode"],
|
||
seed=row["seed"],
|
||
error_text=row["error_text"],
|
||
finished_at=_as_iso(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: Mapping[str, Any]) -> BlockedRequest:
|
||
return cls(
|
||
id=int(row["id"]),
|
||
user_id=int(row["user_id"]),
|
||
timestamp=_as_iso(row["timestamp"]) or "",
|
||
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
|
||
age_confirmed: bool
|
||
age_declined: 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: Mapping[str, Any]) -> User:
|
||
return cls(
|
||
user_id=int(row["user_id"]),
|
||
username=row["username"],
|
||
credits=int(row["credits"]),
|
||
last_daily_credit=_as_iso(row["last_daily_credit"]) or "",
|
||
referrer_id=row["referrer_id"],
|
||
first_paid_at=_as_iso(row["first_paid_at"]),
|
||
created_at=_as_iso(row["created_at"]) or "",
|
||
is_admin=bool(row["is_admin"]),
|
||
is_blocked=bool(row["is_blocked"]),
|
||
age_confirmed=bool(row["age_confirmed"]),
|
||
age_declined=bool(row["age_declined"]),
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class PendingPayment:
|
||
payload: str
|
||
user_id: int
|
||
timestamp: str
|
||
mode: str
|
||
prompt: str
|
||
context_base_prompt: str | None
|
||
context_prompt: str | None
|
||
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: Mapping[str, Any]) -> PendingPayment:
|
||
return cls(
|
||
payload=row["payload"],
|
||
user_id=int(row["user_id"]),
|
||
timestamp=_as_iso(row["timestamp"]) or "",
|
||
mode=row["mode"],
|
||
prompt=row["prompt"],
|
||
context_base_prompt=row["context_base_prompt"],
|
||
context_prompt=row["context_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"],
|
||
)
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class UserContext:
|
||
user_id: int
|
||
base_prompt: str
|
||
current_prompt: str
|
||
last_request_id: int | None
|
||
last_mode: str | None
|
||
updated_at: str
|
||
|
||
@classmethod
|
||
def from_row(cls, row: Mapping[str, Any]) -> UserContext:
|
||
return cls(
|
||
user_id=int(row["user_id"]),
|
||
base_prompt=row["base_prompt"],
|
||
current_prompt=row["current_prompt"],
|
||
last_request_id=row["last_request_id"],
|
||
last_mode=row["last_mode"],
|
||
updated_at=_as_iso(row["updated_at"]) or "",
|
||
)
|
||
|
||
|
||
SCHEMA = """
|
||
CREATE TABLE IF NOT EXISTS requests (
|
||
id BIGSERIAL PRIMARY KEY,
|
||
user_id BIGINT NOT NULL,
|
||
username TEXT,
|
||
timestamp TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||
prompt_text TEXT NOT NULL,
|
||
enhanced_prompt TEXT,
|
||
context_base_prompt TEXT,
|
||
context_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 BOOLEAN NOT NULL DEFAULT FALSE,
|
||
is_premium BOOLEAN NOT NULL DEFAULT FALSE,
|
||
hidden BOOLEAN NOT NULL DEFAULT FALSE,
|
||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
||
seed BIGINT,
|
||
error_text TEXT,
|
||
finished_at TIMESTAMPTZ
|
||
);
|
||
|
||
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 BIGSERIAL PRIMARY KEY,
|
||
user_id BIGINT NOT NULL,
|
||
timestamp TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||
prompt TEXT NOT NULL,
|
||
reason TEXT NOT NULL,
|
||
mode TEXT,
|
||
has_person BOOLEAN NOT NULL DEFAULT FALSE
|
||
);
|
||
|
||
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 BIGINT NOT NULL,
|
||
timestamp TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||
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 BOOLEAN NOT NULL DEFAULT FALSE,
|
||
context_base_prompt TEXT,
|
||
context_prompt TEXT,
|
||
telegram_payment_charge_id TEXT,
|
||
request_id BIGINT
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_pending_user_time
|
||
ON pending_payments(user_id, timestamp);
|
||
|
||
CREATE TABLE IF NOT EXISTS users (
|
||
user_id BIGINT PRIMARY KEY,
|
||
username TEXT,
|
||
credits INTEGER NOT NULL DEFAULT 3,
|
||
last_daily_credit DATE NOT NULL,
|
||
referrer_id BIGINT,
|
||
first_paid_at TIMESTAMPTZ,
|
||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||
is_admin BOOLEAN NOT NULL DEFAULT FALSE,
|
||
is_blocked BOOLEAN NOT NULL DEFAULT FALSE,
|
||
age_confirmed BOOLEAN NOT NULL DEFAULT FALSE,
|
||
age_declined BOOLEAN NOT NULL DEFAULT FALSE
|
||
);
|
||
|
||
CREATE INDEX IF NOT EXISTS idx_users_referrer
|
||
ON users(referrer_id);
|
||
|
||
CREATE TABLE IF NOT EXISTS user_contexts (
|
||
user_id BIGINT PRIMARY KEY,
|
||
base_prompt TEXT NOT NULL,
|
||
current_prompt TEXT NOT NULL,
|
||
last_request_id BIGINT,
|
||
last_mode TEXT,
|
||
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
|
||
);
|
||
|
||
CREATE TABLE IF NOT EXISTS app_settings (
|
||
key TEXT PRIMARY KEY,
|
||
value TEXT
|
||
);
|
||
"""
|
||
|
||
SCHEMA_STATEMENTS = [part.strip() for part in SCHEMA.split(";") if part.strip()]
|
||
|
||
MIGRATIONS = (
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS username TEXT",
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS enhanced_prompt TEXT",
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_base_prompt TEXT",
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS context_prompt TEXT",
|
||
(
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS hidden BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_charge_id TEXT",
|
||
(
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS payment_refunded BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
(
|
||
"ALTER TABLE requests ADD COLUMN IF NOT EXISTS is_premium BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
(
|
||
"ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS is_premium BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
"ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_base_prompt TEXT",
|
||
"ALTER TABLE pending_payments ADD COLUMN IF NOT EXISTS context_prompt TEXT",
|
||
(
|
||
"ALTER TABLE users ADD COLUMN IF NOT EXISTS is_admin BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
(
|
||
"ALTER TABLE users ADD COLUMN IF NOT EXISTS is_blocked BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
(
|
||
"ALTER TABLE users ADD COLUMN IF NOT EXISTS age_confirmed BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
(
|
||
"ALTER TABLE users ADD COLUMN IF NOT EXISTS age_declined BOOLEAN"
|
||
" NOT NULL DEFAULT FALSE"
|
||
),
|
||
)
|
||
|
||
|
||
class Database:
|
||
def __init__(
|
||
self,
|
||
dsn: str,
|
||
*,
|
||
min_size: int = 1,
|
||
max_size: int = 10,
|
||
) -> None:
|
||
self._dsn = _normalize_dsn(dsn)
|
||
self._min_size = max(1, min_size)
|
||
self._max_size = max(self._min_size, max_size)
|
||
self._pool: AsyncConnectionPool | None = None
|
||
|
||
async def connect(self) -> None:
|
||
if not self._dsn:
|
||
raise RuntimeError(
|
||
"DATABASE_URL не задан. Добавь строку подключения PostgreSQL в .env"
|
||
)
|
||
if self._pool is not None:
|
||
await self.close()
|
||
self._pool = AsyncConnectionPool(
|
||
self._dsn,
|
||
min_size=self._min_size,
|
||
max_size=self._max_size,
|
||
open=False,
|
||
kwargs={"row_factory": dict_row},
|
||
timeout=10.0,
|
||
)
|
||
await self._pool.open(wait=True, timeout=15.0)
|
||
async with self._pool.connection() as conn:
|
||
await conn.execute("SELECT pg_advisory_lock(%s)", (SCHEMA_LOCK_KEY,))
|
||
try:
|
||
for statement in SCHEMA_STATEMENTS:
|
||
await conn.execute(statement)
|
||
await self._migrate(conn)
|
||
finally:
|
||
await conn.execute(
|
||
"SELECT pg_advisory_unlock(%s)", (SCHEMA_LOCK_KEY,)
|
||
)
|
||
|
||
async def _migrate(self, conn: AsyncConnection) -> None:
|
||
for statement in MIGRATIONS:
|
||
await conn.execute(statement)
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO users
|
||
(user_id, username, credits, last_daily_credit, created_at)
|
||
SELECT user_id, NULL, 3, %s, %s
|
||
FROM requests
|
||
GROUP BY user_id
|
||
ON CONFLICT (user_id) DO NOTHING
|
||
""",
|
||
(msk_today(), utc_now()),
|
||
)
|
||
|
||
async def close(self) -> None:
|
||
if self._pool is not None:
|
||
await self._pool.close()
|
||
self._pool = None
|
||
|
||
def _connection_pool(self) -> AsyncConnectionPool:
|
||
if self._pool is None:
|
||
raise RuntimeError("База данных не подключена")
|
||
return self._pool
|
||
|
||
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,
|
||
context_base_prompt: str | None = None,
|
||
context_prompt: str | None = None,
|
||
) -> Request:
|
||
try:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO requests
|
||
(user_id, username, timestamp, prompt_text,
|
||
context_base_prompt, context_prompt, input_file_path,
|
||
status, mode, seed, payment_charge_id, is_premium)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
user_id,
|
||
username,
|
||
utc_now(),
|
||
prompt_text,
|
||
context_base_prompt,
|
||
context_prompt,
|
||
input_file_path,
|
||
mode,
|
||
seed,
|
||
payment_charge_id,
|
||
is_premium,
|
||
),
|
||
)
|
||
row = await cursor.fetchone()
|
||
except errors.UniqueViolation as exc:
|
||
raise UserAlreadyQueued(user_id) from exc
|
||
if row is None:
|
||
raise RuntimeError("Не удалось создать запрос")
|
||
return Request.from_row(row)
|
||
|
||
async def get_request(self, request_id: int) -> Request | None:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM requests WHERE id = %s", (request_id,)
|
||
)
|
||
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:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
SELECT * FROM requests
|
||
WHERE user_id = %s AND status IN ('queued', 'running')
|
||
ORDER BY id DESC LIMIT 1
|
||
""",
|
||
(user_id,),
|
||
)
|
||
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:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM users WHERE user_id = %s", (user_id,)
|
||
)
|
||
row = await cursor.fetchone()
|
||
return User.from_row(row) if row is not None else None
|
||
|
||
async def get_user_context(self, user_id: int) -> UserContext | None:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM user_contexts WHERE user_id = %s", (user_id,)
|
||
)
|
||
row = await cursor.fetchone()
|
||
return UserContext.from_row(row) if row is not None else None
|
||
|
||
async def save_user_context(
|
||
self,
|
||
*,
|
||
user_id: int,
|
||
base_prompt: str,
|
||
current_prompt: str,
|
||
last_request_id: int | None,
|
||
last_mode: str | None,
|
||
) -> None:
|
||
async with self._connection_pool().connection() as conn:
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO user_contexts
|
||
(user_id, base_prompt, current_prompt, last_request_id,
|
||
last_mode, updated_at)
|
||
VALUES (%s, %s, %s, %s, %s, now())
|
||
ON CONFLICT (user_id) DO UPDATE
|
||
SET base_prompt = EXCLUDED.base_prompt,
|
||
current_prompt = EXCLUDED.current_prompt,
|
||
last_request_id = EXCLUDED.last_request_id,
|
||
last_mode = EXCLUDED.last_mode,
|
||
updated_at = now()
|
||
""",
|
||
(
|
||
user_id,
|
||
base_prompt,
|
||
current_prompt,
|
||
last_request_id,
|
||
last_mode,
|
||
),
|
||
)
|
||
|
||
async def clear_user_context(self, user_id: int) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"DELETE FROM user_contexts WHERE user_id = %s RETURNING user_id",
|
||
(user_id,),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def ensure_user(
|
||
self, user_id: int, username: str | None = None
|
||
) -> User:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO users
|
||
(user_id, username, credits, last_daily_credit, created_at)
|
||
VALUES (%s, %s, 3, %s, %s)
|
||
ON CONFLICT (user_id) DO UPDATE
|
||
SET username = COALESCE(EXCLUDED.username, users.username)
|
||
RETURNING *
|
||
""",
|
||
(user_id, username, msk_today(), utc_now()),
|
||
)
|
||
row = await cursor.fetchone()
|
||
if row is None:
|
||
raise RuntimeError("Не удалось прочитать пользователя")
|
||
return User.from_row(row)
|
||
|
||
async def seed_admins(self, user_ids: frozenset[int] | set[int]) -> None:
|
||
if not user_ids:
|
||
return
|
||
async with (
|
||
self._connection_pool().connection() as conn,
|
||
conn.transaction(),
|
||
):
|
||
cursor = await conn.execute(
|
||
"SELECT value FROM app_settings WHERE key = 'admins_seeded'"
|
||
)
|
||
if await cursor.fetchone() is not None:
|
||
return
|
||
today = msk_today()
|
||
now = utc_now()
|
||
for user_id in user_ids:
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO users
|
||
(user_id, username, credits, last_daily_credit,
|
||
created_at)
|
||
VALUES (%s, NULL, 3, %s, %s)
|
||
ON CONFLICT (user_id) DO NOTHING
|
||
""",
|
||
(user_id, today, now),
|
||
)
|
||
await conn.execute(
|
||
"UPDATE users SET is_admin = TRUE WHERE user_id = %s",
|
||
(user_id,),
|
||
)
|
||
await conn.execute(
|
||
"""
|
||
INSERT INTO app_settings (key, value)
|
||
VALUES ('admins_seeded', '1')
|
||
ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value
|
||
"""
|
||
)
|
||
|
||
async def list_users_admin(self, limit: int = 500) -> list[User]:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM users ORDER BY user_id LIMIT %s", (limit,)
|
||
)
|
||
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_age_confirmed(self, user_id: int, confirmed: bool) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users
|
||
SET age_confirmed = %s, age_declined = FALSE
|
||
WHERE user_id = %s
|
||
RETURNING user_id
|
||
""",
|
||
(confirmed, user_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def set_age_declined(self, user_id: int) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users
|
||
SET age_declined = TRUE, age_confirmed = FALSE
|
||
WHERE user_id = %s
|
||
RETURNING user_id
|
||
""",
|
||
(user_id,),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def set_credits(self, user_id: int, credits: int) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET credits = %s
|
||
WHERE user_id = %s
|
||
RETURNING user_id
|
||
""",
|
||
(max(0, credits), user_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def _update_user_flag(
|
||
self, column: str, user_id: int, enabled: bool
|
||
) -> bool:
|
||
if column not in {"is_admin", "is_blocked"}:
|
||
raise ValueError("Недопустимое поле")
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
f"""
|
||
UPDATE users SET {column} = %s
|
||
WHERE user_id = %s
|
||
RETURNING user_id
|
||
""",
|
||
(enabled, user_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def add_credits(self, user_id: int, amount: int) -> int | None:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET credits = credits + %s
|
||
WHERE user_id = %s
|
||
RETURNING credits
|
||
""",
|
||
(amount, user_id),
|
||
)
|
||
row = await cursor.fetchone()
|
||
return int(row["credits"]) if row is not None else None
|
||
|
||
async def consume_credit(self, user_id: int) -> int | None:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET credits = credits - 1
|
||
WHERE user_id = %s AND credits > 0
|
||
RETURNING credits
|
||
""",
|
||
(user_id,),
|
||
)
|
||
row = await cursor.fetchone()
|
||
return int(row["credits"]) if row is not None else None
|
||
|
||
async def set_referrer(self, user_id: int, referrer_id: int) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET referrer_id = %s
|
||
WHERE user_id = %s AND referrer_id IS NULL AND user_id != %s
|
||
RETURNING user_id
|
||
""",
|
||
(referrer_id, user_id, referrer_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def mark_first_payment(self, user_id: int) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE users SET first_paid_at = %s
|
||
WHERE user_id = %s AND first_paid_at IS NULL
|
||
RETURNING user_id
|
||
""",
|
||
(utc_now(), user_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def accrue_daily_credits(self, today: date | None = None) -> int:
|
||
current = today or msk_today()
|
||
async with (
|
||
self._connection_pool().connection() as conn,
|
||
conn.transaction(),
|
||
):
|
||
cursor = await conn.execute(
|
||
"""
|
||
SELECT user_id, last_daily_credit FROM users
|
||
WHERE last_daily_credit < %s
|
||
""",
|
||
(current,),
|
||
)
|
||
rows = await cursor.fetchall()
|
||
for row in rows:
|
||
last = row["last_daily_credit"]
|
||
if isinstance(last, datetime):
|
||
last = last.date()
|
||
if isinstance(last, str):
|
||
last = date.fromisoformat(last)
|
||
days = max(1, (current - last).days)
|
||
await conn.execute(
|
||
"""
|
||
UPDATE users
|
||
SET credits = credits + %s, last_daily_credit = %s
|
||
WHERE user_id = %s
|
||
""",
|
||
(days, current, row["user_id"]),
|
||
)
|
||
return len(rows)
|
||
|
||
async def list_queued(self) -> list[Request]:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM requests WHERE status = 'queued' ORDER BY id"
|
||
)
|
||
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,
|
||
include_hidden: bool = True,
|
||
) -> tuple[str, list[Any]]:
|
||
clauses: list[str] = []
|
||
params: list[Any] = []
|
||
if user_id is not None:
|
||
clauses.append("user_id = %s")
|
||
params.append(user_id)
|
||
if status:
|
||
clauses.append("status = %s")
|
||
params.append(status)
|
||
if mode:
|
||
clauses.append("mode = %s")
|
||
params.append(mode)
|
||
if not include_hidden:
|
||
clauses.append("hidden = FALSE")
|
||
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,
|
||
include_hidden: bool = True,
|
||
limit: int = 50,
|
||
offset: int = 0,
|
||
) -> list[Request]:
|
||
where, params = self._filters(
|
||
user_id, status, mode, include_hidden=include_hidden
|
||
)
|
||
query = (
|
||
f"SELECT * FROM requests{where} ORDER BY id DESC LIMIT %s OFFSET %s"
|
||
)
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(query, (*params, limit, offset))
|
||
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,
|
||
include_hidden: bool = True,
|
||
) -> int:
|
||
where, params = self._filters(
|
||
user_id, status, mode, include_hidden=include_hidden
|
||
)
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
f"SELECT COUNT(*) AS total FROM requests{where}", params
|
||
)
|
||
row = await cursor.fetchone()
|
||
return int(row["total"]) if row is not None else 0
|
||
|
||
async def stats(self, *, include_hidden: bool = True) -> dict[str, Any]:
|
||
where = "" if include_hidden else " WHERE hidden = FALSE"
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT status, COUNT(*) AS total FROM requests"
|
||
f"{where} GROUP BY status"
|
||
)
|
||
rows = await cursor.fetchall()
|
||
by_status = {
|
||
str(row["status"]): int(row["total"]) for row in rows
|
||
}
|
||
cursor = await conn.execute(
|
||
"SELECT COUNT(DISTINCT user_id) AS total FROM requests"
|
||
f"{where}"
|
||
)
|
||
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, *, include_hidden: bool = True
|
||
) -> list[UserStat]:
|
||
where = "" if include_hidden else " WHERE hidden = FALSE"
|
||
username_filter = (
|
||
"" if include_hidden else " AND r2.hidden = FALSE"
|
||
)
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
f"""
|
||
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
|
||
{username_filter}
|
||
ORDER BY r2.id DESC LIMIT 1
|
||
) AS username
|
||
FROM requests
|
||
{where}
|
||
GROUP BY user_id
|
||
ORDER BY last_seen DESC
|
||
LIMIT %s
|
||
""",
|
||
(limit,),
|
||
)
|
||
rows = await cursor.fetchall()
|
||
return [
|
||
UserStat(
|
||
user_id=int(row["user_id"]),
|
||
username=row["username"],
|
||
total=int(row["total"]),
|
||
last_seen=_as_iso(row["last_seen"]) or "",
|
||
)
|
||
for row in rows
|
||
]
|
||
|
||
async def count_ahead(self, request_id: int) -> int:
|
||
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 = TRUE AND status IN ('queued', 'running')
|
||
AND id < %s
|
||
"""
|
||
else:
|
||
query = """
|
||
SELECT COUNT(*) AS total FROM requests
|
||
WHERE status IN ('queued', 'running')
|
||
AND (is_premium = TRUE OR (is_premium = FALSE AND id < %s))
|
||
"""
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(query, (request_id,))
|
||
row = await cursor.fetchone()
|
||
return int(row["total"]) if row is not None else 0
|
||
|
||
async def queue_counts(self) -> dict[str, int]:
|
||
counts = {"premium": 0, "normal": 0}
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
SELECT is_premium, COUNT(*) AS total FROM requests
|
||
WHERE status IN ('queued', 'running')
|
||
GROUP BY is_premium
|
||
"""
|
||
)
|
||
rows = await cursor.fetchall()
|
||
for row in rows:
|
||
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(%s, comfy_prompt_id),
|
||
error_text = NULL
|
||
WHERE id = %s
|
||
""",
|
||
(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 = %s WHERE id = %s",
|
||
(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 = %s WHERE id = %s",
|
||
(prompt, request_id),
|
||
)
|
||
|
||
async def set_refunded(self, request_id: int) -> None:
|
||
await self._update(
|
||
"UPDATE requests SET payment_refunded = TRUE WHERE id = %s",
|
||
(request_id,),
|
||
)
|
||
|
||
async def set_hidden(self, request_id: int, hidden: bool) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"UPDATE requests SET hidden = %s WHERE id = %s RETURNING id",
|
||
(hidden, request_id),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
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,
|
||
context_base_prompt: str | None = None,
|
||
context_prompt: str | None = None,
|
||
) -> PendingPayment:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO pending_payments
|
||
(payload, user_id, timestamp, mode, prompt, username,
|
||
input_file_path, amount, status, is_premium,
|
||
context_base_prompt, context_prompt)
|
||
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s, %s, %s)
|
||
RETURNING *
|
||
""",
|
||
(
|
||
payload,
|
||
user_id,
|
||
utc_now(),
|
||
mode,
|
||
prompt,
|
||
username,
|
||
input_file_path,
|
||
amount,
|
||
is_premium,
|
||
context_base_prompt,
|
||
context_prompt,
|
||
),
|
||
)
|
||
row = await cursor.fetchone()
|
||
if row is None:
|
||
raise RuntimeError("Не удалось создать счёт")
|
||
return PendingPayment.from_row(row)
|
||
|
||
async def get_pending_payment(
|
||
self, payload: str
|
||
) -> PendingPayment | None:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM pending_payments WHERE payload = %s",
|
||
(payload,),
|
||
)
|
||
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:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE pending_payments
|
||
SET is_premium = %s, amount = %s
|
||
WHERE payload = %s AND status = 'pending' AND amount = 0
|
||
RETURNING payload
|
||
""",
|
||
(is_premium, amount, payload),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
UPDATE pending_payments
|
||
SET status = 'paid', telegram_payment_charge_id = %s
|
||
WHERE payload = %s AND status = 'pending'
|
||
RETURNING payload
|
||
""",
|
||
(charge_id, payload),
|
||
)
|
||
return await cursor.fetchone() is not None
|
||
|
||
async def mark_pending_refunded(self, payload: str) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE pending_payments SET status = 'refunded'
|
||
WHERE payload = %s
|
||
""",
|
||
(payload,),
|
||
)
|
||
|
||
async def mark_pending_cancelled(self, payload: str) -> None:
|
||
await self._update(
|
||
"""
|
||
UPDATE pending_payments SET status = 'cancelled'
|
||
WHERE payload = %s 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 = %s WHERE payload = %s",
|
||
(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 = %s,
|
||
error_text = NULL,
|
||
finished_at = %s
|
||
WHERE id = %s
|
||
""",
|
||
(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 = %s,
|
||
finished_at = %s
|
||
WHERE id = %s
|
||
""",
|
||
(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:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"""
|
||
INSERT INTO blocked_requests
|
||
(user_id, timestamp, prompt, reason, mode, has_person)
|
||
VALUES (%s, %s, %s, %s, %s, %s)
|
||
RETURNING id
|
||
""",
|
||
(user_id, utc_now(), prompt, reason, mode, has_person),
|
||
)
|
||
row = await cursor.fetchone()
|
||
if row is None:
|
||
raise RuntimeError("Не удалось записать блокировку")
|
||
return int(row["id"])
|
||
|
||
async def list_blocked(self, limit: int = 100) -> list[BlockedRequest]:
|
||
async with self._connection_pool().connection() as conn:
|
||
cursor = await conn.execute(
|
||
"SELECT * FROM blocked_requests ORDER BY id DESC LIMIT %s",
|
||
(limit,),
|
||
)
|
||
rows = await cursor.fetchall()
|
||
return [BlockedRequest.from_row(row) for row in rows]
|
||
|
||
async def recover(
|
||
self, reason: str = "прервано перезапуском"
|
||
) -> list[Request]:
|
||
"""Помечает running как failed и возвращает queued для повторной постановки."""
|
||
await self._update(
|
||
"""
|
||
UPDATE requests
|
||
SET status = 'failed', error_text = %s, finished_at = %s
|
||
WHERE status = 'running'
|
||
""",
|
||
(reason, utc_now()),
|
||
)
|
||
return await self.list_queued()
|
||
|
||
async def _update(self, sql: str, params: tuple[Any, ...]) -> None:
|
||
async with self._connection_pool().connection() as conn:
|
||
await conn.execute(sql, params)
|