From 44e77e83b17fa606f60245a9718592ffe54071b5 Mon Sep 17 00:00:00 2001 From: MpakobecEx Date: Mon, 19 Oct 2026 02:41:04 +0500 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=B5=D1=80=D0=B5=D1=85=D0=BE=D0=B4=20?= =?UTF-8?q?=D0=BD=D0=B0=20PostgreSQL?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Заменил SQLite на PostgreSQL 17: psycopg и пул соединений, BOOLEAN/TIMESTAMPTZ/BIGSERIAL, отдельный сервис db в Compose. Добавил скрипт переноса старых данных и тесты БД в CI. --- .env.example | 10 +- .gitea/workflows/docker-build.yml | 38 ++ README.md | 32 +- bot.py | 6 +- checklist.md | 36 +- config.py | 10 +- db.py | 906 ++++++++++++++------------ docker-compose.yml | 21 + plan.md | 6 +- requirements-dev.txt | 1 + requirements.txt | 3 +- scripts/__init__.py | 0 scripts/migrate_sqlite_to_postgres.py | 344 ++++++++++ tests/conftest.py | 62 ++ tests/test_db.py | 71 +- tests/test_migration.py | 102 +++ tests/test_queue.py | 12 +- tests/test_web.py | 150 +++-- web.py | 6 +- 19 files changed, 1236 insertions(+), 580 deletions(-) create mode 100644 scripts/__init__.py create mode 100644 scripts/migrate_sqlite_to_postgres.py create mode 100644 tests/conftest.py create mode 100644 tests/test_migration.py diff --git a/.env.example b/.env.example index bec13e5..758ec3b 100644 --- a/.env.example +++ b/.env.example @@ -22,7 +22,15 @@ PROMPT_ENHANCE_TIMEOUT=120 # Хранилище STORAGE_DIR=storage -DB_PATH=data/bot.db + +# PostgreSQL +# В Docker хост базы — db. При локальном запуске замени db на localhost. +POSTGRES_DB=mpakoimageaibot +POSTGRES_USER=mpako +POSTGRES_PASSWORD=change-me +DATABASE_URL=postgresql://mpako:change-me@db:5432/mpakoimageaibot +DB_POOL_MIN_SIZE=1 +DB_POOL_MAX_SIZE=10 # Очередь WORKER_CONCURRENCY=1 diff --git a/.gitea/workflows/docker-build.yml b/.gitea/workflows/docker-build.yml index 2d9712a..1c8424d 100644 --- a/.gitea/workflows/docker-build.yml +++ b/.gitea/workflows/docker-build.yml @@ -12,7 +12,45 @@ env: IMAGE: git.mpakosrv.ru/mpakobecex/mpakoimageaibot jobs: + tests: + runs-on: ubuntu-latest + timeout-minutes: 10 + services: + postgres: + image: postgres:17-alpine + env: + POSTGRES_DB: mpako_test + POSTGRES_USER: mpako + POSTGRES_PASSWORD: mpako + ports: + - 55432:5432 + options: >- + --health-cmd "pg_isready -U mpako -d mpako_test" + --health-interval 5s + --health-timeout 5s + --health-retries 20 + steps: + - name: Checkout + uses: actions/checkout@v4 + - name: Setup Python + uses: actions/setup-python@v5 + with: + python-version: "3.13" + - name: Install dependencies + run: pip install -r requirements-dev.txt + - name: Lint + run: ruff check . --exclude .venv + - name: Tests + run: | + if getent hosts postgres >/dev/null 2>&1; then + export TEST_DATABASE_URL=postgresql://mpako:mpako@postgres:5432/mpako_test + else + export TEST_DATABASE_URL=postgresql://mpako:mpako@localhost:55432/mpako_test + fi + pytest -q + docker: + needs: tests runs-on: ubuntu-latest timeout-minutes: 20 steps: diff --git a/README.md b/README.md index e771a24..5379139 100644 --- a/README.md +++ b/README.md @@ -17,7 +17,7 @@ Telegram-бот для генерации и редактирования изо - **Подтверждение 18+**: без подтверждения возраста бот не принимает запросы. - Глобальная очередь с одной активной задачей на пользователя. - Позиция в очереди и прогресс генерации через WebSocket ComfyUI. -- SQLite: запросы, статусы, seed, пути к файлам. +- PostgreSQL: запросы, статусы, seed, пути к файлам. ## Требования @@ -40,6 +40,7 @@ Copy-Item .env.example .env - `BOT_TOKEN` — токен от `@BotFather`. - `TELEGRAM_PROXY` — прокси для Telegram API (пусто — напрямую). - `COMFY_BASE_URL` — например `http://192.168.31.113:8188`. +- `DATABASE_URL` — строка подключения PostgreSQL. В Docker хост `db`, при локальном запуске — `localhost`. - Остальные параметры можно оставить по умолчанию. ### Прокси для Telegram @@ -67,17 +68,42 @@ python bot.py Бот работает в приватных чатах. Пользователь отправляет текст или фото с подписью. +Для локального запуска нужен PostgreSQL и `DATABASE_URL` с хостом `localhost`. + ## Docker ```powershell docker compose up -d --build ``` -Каталоги `./storage` и `./data` монтируются как volume. Для ComfyUI на другой машине используйте её LAN-IP, а не `localhost`. +Compose поднимает PostgreSQL, бота и веб. Каталог `./storage` монтируется в контейнеры, данные PostgreSQL лежат в `./data/postgres`. Для ComfyUI на другой машине используйте её LAN-IP, а не `localhost`. + +## PostgreSQL + +```dotenv +POSTGRES_DB=mpakoimageaibot +POSTGRES_USER=mpako +POSTGRES_PASSWORD=change-me +DATABASE_URL=postgresql://mpako:change-me@db:5432/mpakoimageaibot +DB_POOL_MIN_SIZE=1 +DB_POOL_MAX_SIZE=10 +``` + +Пароль в `DATABASE_URL` должен совпадать с `POSTGRES_PASSWORD`; спецсимволы в пароле нужно URL-кодировать. + +### Перенос старой SQLite-базы + +```powershell +python scripts/migrate_sqlite_to_postgres.py ` + --sqlite data/bot.db ` + --database-url "postgresql://mpako:change-me@localhost:5432/mpakoimageaibot" +``` + +Перед миграцией останови старый бот, чтобы SQLite закрылась и данные были консистентны. Скрипт создаёт схему, переносит данные и пропускает уже существующие записи. ## CI -При push в `main` или `master` Gitea Actions собирает образ на сервере, где запущен runner, и пушит его в Gitea Registry: +При push в `main` или `master` Gitea Actions сначала запускает тесты на PostgreSQL, затем собирает образ на сервере, где запущен runner, и пушит его в Gitea Registry: - `git.mpakosrv.ru/mpakobecex/mpakoimageaibot:latest` - `git.mpakosrv.ru/mpakobecex/mpakoimageaibot:` diff --git a/bot.py b/bot.py index 9862484..fc1e7e6 100644 --- a/bot.py +++ b/bot.py @@ -58,7 +58,11 @@ AGE_DENIED = "Доступ к боту разрешён только с 18 ле class MpakoBot: def __init__(self, settings: Settings) -> None: self._settings = settings - self._db = Database(settings.db_path) + self._db = Database( + settings.database_url, + min_size=settings.db_pool_min_size, + max_size=settings.db_pool_max_size, + ) self._workflows: WorkflowBuilder | None = None self._comfy: ComfyClient | None = None self._queue: QueueManager | None = None diff --git a/checklist.md b/checklist.md index c5d84be..e7999ea 100644 --- a/checklist.md +++ b/checklist.md @@ -1,6 +1,6 @@ # Checklist реализации MpakoImageAIbot -Источник: `plan.md`. Бот на `python-telegram-bot` v21+ (async) генерирует изображения через удалённый ComfyUI API, ведёт глобальную очередь, показывает позицию/прогресс и хранит запросы в SQLite. +Источник: `plan.md`. Бот на `python-telegram-bot` v21+ (async) генерирует изображения через удалённый ComfyUI API, ведёт глобальную очередь, показывает позицию/прогресс и хранит запросы в PostgreSQL. Блоки идут в порядке зависимостей. В конце каждого блока — критерий готовности (DoD). @@ -8,8 +8,8 @@ - Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15), оплата Stars (16), баланс/рефералы (17), управление пользователями (18), возрастной gate (19) и надёжность Telegram (20). - Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история. -- 54 теста проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry. -- Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. +- 56 тестов: 34 проходят локально без БД, 22 теста БД выполняются в CI на PostgreSQL через `TEST_DATABASE_URL`. Покрыты workflows, очередь, веб-история, enhancer, модерация, платежи, премиум-очередь, баланс/рефералы, управление пользователями, retry, миграция. +- Не проверено: реальная оплата Stars и реферальная выплата в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото. PostgreSQL-миграция и тесты БД ещё не запускались вживую — это сделает CI и первый деплой. ## 0. Допущения и решения по умолчанию @@ -18,7 +18,7 @@ | HTTP-клиент | `httpx.AsyncClient`, без блокирующего `requests` | | Очередь | Собственная очередь в приложении, по умолчанию 1 воркер; в ComfyUI уходит не более 1 задачи одновременно | | Прогресс | WebSocket `/ws` + fallback-опрос `/history/{prompt_id}` | -| Один активный запрос на пользователя | Проверка в памяти + частичный уникальный индекс в SQLite | +| Один активный запрос на пользователя | Проверка в памяти + частичный уникальный индекс в PostgreSQL | | Задача в статусе `running` при рестарте | Пометить `failed` (не перезапускать автоматически) | | Задачи `queued` при рестарте | Восстановить из БД и поставить в очередь заново | | i2i без подписи | Отклонить и попросить промпт (не использовать дефолтный) | @@ -93,7 +93,7 @@ │ ├── img2img.json │ └── mapping.json ├── storage/ # создаётся автоматически, в .gitignore -├── data/ # SQLite, в .gitignore +├── data/ # PostgreSQL и старый SQLite, в .gitignore ├── tests/ ├── requirements.txt ├── requirements-dev.txt @@ -104,10 +104,10 @@ └── README.md ``` -- [ ] `requirements.txt`: `python-telegram-bot`, `httpx`, `aiosqlite`, `websockets`, `python-dotenv` (версии — диапазонами, например `>=21,<23`). +- [ ] `requirements.txt`: `python-telegram-bot`, `httpx`, `psycopg[binary]`, `psycopg-pool`, `websockets`, `python-dotenv` (версии — диапазонами, например `>=21,<23`). - [ ] `requirements-dev.txt`: `pytest`, `pytest-asyncio`, `respx`. - [ ] `config.py`: dataclass/`Settings`, загрузка `.env`, валидация обязательных полей, вычисление WS-URL из base URL. -- [ ] Переменные окружения: `BOT_TOKEN`, `COMFY_BASE_URL`, `WORKFLOW_T2I`, `WORKFLOW_I2I`, `WORKFLOW_MAPPING`, `STORAGE_DIR`, `DB_PATH`, `WORKER_CONCURRENCY=1`, `MAX_QUEUE_SIZE`, `ALLOWED_USER_IDS`, `PROGRESS_INTERVAL`, `GENERATION_TIMEOUT`, `LOG_LEVEL`. +- [ ] Переменные окружения: `BOT_TOKEN`, `COMFY_BASE_URL`, `WORKFLOW_T2I`, `WORKFLOW_I2I`, `WORKFLOW_MAPPING`, `STORAGE_DIR`, `DATABASE_URL`, `DB_POOL_MIN_SIZE`, `DB_POOL_MAX_SIZE`, `WORKER_CONCURRENCY=1`, `MAX_QUEUE_SIZE`, `ALLOWED_USER_IDS`, `PROGRESS_INTERVAL`, `GENERATION_TIMEOUT`, `LOG_LEVEL`. - [ ] Настроить логирование в stdout: уровень из env, формат с timestamp/level/logger, никогда не логировать токен и содержимое `.env`. - [ ] `.env.example` со всеми переменными и комментариями. - [ ] `.gitignore`: `.env`, `storage/`, `data/`, `__pycache__/`, `.pytest_cache/`. @@ -117,14 +117,14 @@ ## 3. Слой хранения (`db.py`) -- [ ] Реализовать `init_db()`: подключение `aiosqlite`, `PRAGMA journal_mode=WAL`, `PRAGMA busy_timeout`, `PRAGMA foreign_keys=ON`, `row_factory`. +- [ ] Реализовать `connect()`: пул `psycopg_pool.AsyncConnectionPool`, `row_factory=dict_row`, схема и миграции. - [ ] Создать таблицу `requests` по плану: ```sql CREATE TABLE IF NOT EXISTS requests ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id INTEGER NOT NULL, - timestamp TEXT NOT NULL, + id BIGSERIAL PRIMARY KEY, + user_id BIGINT NOT NULL, + timestamp TIMESTAMPTZ NOT NULL DEFAULT now(), prompt_text TEXT NOT NULL, input_file_path TEXT, output_file_path TEXT, @@ -134,7 +134,7 @@ CREATE TABLE IF NOT EXISTS requests ( ); ``` -- [ ] Добавить расширения схемы (согласовать): `mode TEXT NOT NULL CHECK (mode IN ('t2i','i2i'))`, `error_text TEXT`, `finished_at TEXT`, `seed INTEGER`. +- [ ] Добавить расширения схемы (согласовать): `mode TEXT NOT NULL CHECK (mode IN ('t2i','i2i'))`, `error_text TEXT`, `finished_at TIMESTAMPTZ`, `seed BIGINT`. - [ ] Индексы: `(user_id, status)`, `(status, id)`. - [ ] Частичный уникальный индекс «один активный запрос на пользователя»: @@ -144,12 +144,12 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); ``` - [ ] CRUD: `create_request`, `get_request`, `get_active_by_user`, `set_status`, `set_comfy_prompt_id`, `set_output_path`, `set_error`, `list_queued`, `count_ahead`, `mark_running_as_failed`. -- [ ] Все временные метки — ISO 8601 UTC (`datetime.now(timezone.utc).isoformat()`). +- [ ] Все временные метки — `TIMESTAMPTZ`, в Python — `datetime.now(UTC)`. - [ ] Атомарность: создание запроса и проверка активного — в одной транзакции; при конфликте unique-индекса возвращать понятную ошибку `UserAlreadyQueued`. - [ ] Восстановление при старте: `queued` читаются по `id`, `running` → `failed` с текстом «прервано перезапуском». - [ ] Тесты: создание, смена статусов, запрет второго активного запроса, выборка очереди, восстановление. -**DoD:** тесты БД зелёные; повторный `init_db()` не ломает существующую БД. +**DoD:** тесты БД зелёные на PostgreSQL; повторный `connect()` не ломает существующую БД. ## 4. Клиент ComfyUI (`comfy_client.py`) @@ -251,7 +251,7 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); - [ ] Создание директорий при старте/записи. - [ ] (Опционально) очистка файлов старше N дней по расписанию; настройка `FILE_RETENTION_DAYS`. - [ ] (Опционально) проверка свободного места и понятная ошибка при переполнении. -- [ ] Не хранить бинарные данные в SQLite — только пути. +- [ ] Не хранить бинарные данные в PostgreSQL — только пути. **DoD:** входные и выходные файлы лежат по ожидаемым путям; повторный запуск не конфликтует по именам. @@ -260,11 +260,11 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); - [x] `Dockerfile` на `python:3.13-slim`, непривилегированный пользователь, `PYTHONUNBUFFERED=1`. - [x] Копирование кода и установка `requirements.txt`; без лишних build-зависимостей. - [x] `.dockerignore`: `.env`, `storage/`, `data/`, `__pycache__/`, `tests/`, `.git/`. -- [x] `docker-compose.yml`: `env_file: .env`, volume для `./storage` и `./data`, `restart: unless-stopped`. +- [x] `docker-compose.yml`: сервис `db` (PostgreSQL 17), healthcheck, volume `./data/postgres`, `env_file: .env`, volume `./storage`, `restart: unless-stopped`. - [x] `docker-compose.yml` использует готовый образ `mpakoimageaibot:latest`. - [x] Gitea Actions: `.gitea/workflows/docker-build.yml` собирает и пушит `mpakoimageaibot:latest` и `:` в Registry при push в `main`/`master`. - [x] Сетевой доступ к ComfyUI: `host.docker.internal` для Windows/macOS или LAN-IP для Linux; зафиксировать в README. -- [x] Проверить, что SQLite и storage переживают пересоздание контейнера. +- [x] Проверить, что PostgreSQL и storage переживают пересоздание контейнера. - [ ] (Опционально) healthcheck. - [x] Логи только в stdout/stderr. @@ -406,5 +406,5 @@ ON requests(user_id) WHERE status IN ('queued', 'running'); - Несколько моделей/чекпоинтов и выбор модели пользователем. - Галерея, история генераций в боте, повтор генерации. - LoRA/ControlNet и другие продвинутые workflow. -- Горизонтальное масштабирование и несколько экземпляров бота на одной SQLite. +- Горизонтальное масштабирование и несколько экземпляров бота на одной PostgreSQL. - Админ-панель. diff --git a/config.py b/config.py index 68be34e..b8907fa 100644 --- a/config.py +++ b/config.py @@ -84,7 +84,9 @@ class Settings: prompt_enhance_enabled: bool = False prompt_enhance_timeout: float = 120.0 storage_dir: Path = PROJECT_ROOT / "storage" - db_path: Path = PROJECT_ROOT / "data" / "bot.db" + database_url: str = "" + db_pool_min_size: int = 1 + db_pool_max_size: int = 10 worker_concurrency: int = 1 max_queue_size: int = 20 allowed_user_ids: frozenset[int] = field(default_factory=frozenset) @@ -129,6 +131,8 @@ class Settings: "socks5:// или socks5h://" ) + database_url = os.getenv("DATABASE_URL", "").strip() + settings = cls( bot_token=token, telegram_proxy=proxy, @@ -151,7 +155,9 @@ class Settings: 10.0, _float("PROMPT_ENHANCE_TIMEOUT", 120.0) ), storage_dir=_path("STORAGE_DIR", "storage"), - db_path=_path("DB_PATH", "data/bot.db"), + database_url=database_url, + db_pool_min_size=max(1, _int("DB_POOL_MIN_SIZE", 1)), + db_pool_max_size=max(1, _int("DB_POOL_MAX_SIZE", 10)), worker_concurrency=max(1, _int("WORKER_CONCURRENCY", 1)), max_queue_size=max(1, _int("MAX_QUEUE_SIZE", 20)), allowed_user_ids=_user_ids("ALLOWED_USER_IDS"), diff --git a/db.py b/db.py index 46170ff..4e1819f 100644 --- a/db.py +++ b/db.py @@ -1,29 +1,52 @@ from __future__ import annotations -import asyncio import logging -import sqlite3 +from collections.abc import Mapping from dataclasses import dataclass from datetime import UTC, date, datetime, timedelta, timezone -from pathlib import Path from typing import Any -import aiosqlite +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() -> str: - return datetime.now(UTC).isoformat() +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} уже есть активный запрос") @@ -55,12 +78,12 @@ class Request: return f"@{self.username}" if self.username else str(self.user_id) @classmethod - def from_row(cls, row: aiosqlite.Row) -> Request: + def from_row(cls, row: Mapping[str, Any]) -> Request: return cls( - id=row["id"], - user_id=row["user_id"], + id=int(row["id"]), + user_id=int(row["user_id"]), username=row["username"], - timestamp=row["timestamp"], + timestamp=_as_iso(row["timestamp"]) or "", prompt_text=row["prompt_text"], enhanced_prompt=row["enhanced_prompt"], input_file_path=row["input_file_path"], @@ -73,7 +96,7 @@ class Request: mode=row["mode"], seed=row["seed"], error_text=row["error_text"], - finished_at=row["finished_at"], + finished_at=_as_iso(row["finished_at"]), ) @@ -100,11 +123,11 @@ class BlockedRequest: has_person: bool @classmethod - def from_row(cls, row: aiosqlite.Row) -> BlockedRequest: + def from_row(cls, row: Mapping[str, Any]) -> BlockedRequest: return cls( - id=row["id"], - user_id=row["user_id"], - timestamp=row["timestamp"], + 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"], @@ -131,15 +154,15 @@ class User: return f"@{self.username}" if self.username else str(self.user_id) @classmethod - def from_row(cls, row: aiosqlite.Row) -> User: + 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=row["last_daily_credit"], + last_daily_credit=_as_iso(row["last_daily_credit"]) or "", referrer_id=row["referrer_id"], - first_paid_at=row["first_paid_at"], - created_at=row["created_at"], + 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"]), @@ -163,11 +186,11 @@ class PendingPayment: request_id: int | None @classmethod - def from_row(cls, row: aiosqlite.Row) -> PendingPayment: + def from_row(cls, row: Mapping[str, Any]) -> PendingPayment: return cls( payload=row["payload"], - user_id=row["user_id"], - timestamp=row["timestamp"], + user_id=int(row["user_id"]), + timestamp=_as_iso(row["timestamp"]) or "", mode=row["mode"], prompt=row["prompt"], username=row["username"], @@ -182,24 +205,24 @@ class PendingPayment: 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, + 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, + 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 + payment_refunded BOOLEAN NOT NULL DEFAULT FALSE, + is_premium 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 @@ -212,49 +235,49 @@ 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, + 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 INTEGER NOT NULL DEFAULT 0 + 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 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 + 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, + 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 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, - age_confirmed INTEGER NOT NULL DEFAULT 0, - age_declined INTEGER NOT NULL DEFAULT 0 + 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 @@ -266,90 +289,107 @@ CREATE TABLE IF NOT EXISTS app_settings ( ); """ +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 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 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, 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 + 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: - 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() + 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) -> 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" - ) - if "age_confirmed" not in user_columns: - await conn.execute( - "ALTER TABLE users ADD COLUMN age_confirmed INTEGER NOT NULL DEFAULT 0" - ) - if "age_declined" not in user_columns: - await conn.execute( - "ALTER TABLE users ADD COLUMN age_declined INTEGER NOT NULL DEFAULT 0" - ) + async def _migrate(self, conn: AsyncConnection) -> None: + for statement in MIGRATIONS: + await conn.execute(statement) 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) + 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().isoformat(), utc_now()), + (msk_today(), utc_now()), ) async def close(self) -> None: - if self._conn is not None: - await self._conn.close() - self._conn = None + if self._pool is not None: + await self._pool.close() + self._pool = None - def _connection(self) -> aiosqlite.Connection: - if self._conn is None: + def _connection_pool(self) -> AsyncConnectionPool: + if self._pool is None: raise RuntimeError("База данных не подключена") - return self._conn + return self._pool async def create_request( self, @@ -363,16 +403,16 @@ class Database: payment_charge_id: str | None = None, is_premium: bool = False, ) -> Request: - conn = self._connection() - async with self._write_lock: - try: + try: + async with self._connection_pool().connection() as conn: 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', ?, ?, ?, ?) + VALUES (%s, %s, %s, %s, %s, 'queued', %s, %s, %s, %s) + RETURNING * """, ( user_id, @@ -383,248 +423,256 @@ class Database: mode, seed, payment_charge_id, - int(is_premium), + 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 + 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: - conn = self._connection() - async with conn.execute("SELECT * FROM requests WHERE id = ?", (request_id,)) as cursor: + 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: - 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: + 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: - conn = self._connection() - async with conn.execute( - "SELECT * FROM users WHERE user_id = ?", (user_id,) - ) as cursor: + 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 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( + async with self._connection_pool().connection() as conn: + cursor = await conn.execute( """ - INSERT OR IGNORE INTO users + INSERT INTO users (user_id, username, credits, last_daily_credit, created_at) - VALUES (?, ?, 3, ?, ?) + VALUES (%s, %s, 3, %s, %s) + ON CONFLICT (user_id) DO UPDATE + SET username = COALESCE(EXCLUDED.username, users.username) + RETURNING * """, - (user_id, username, today, utc_now()), + (user_id, username, msk_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: + row = await cursor.fetchone() + if row is None: raise RuntimeError("Не удалось прочитать пользователя") - return user + return User.from_row(row) 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( + async with ( + self._connection_pool().connection() as conn, + conn.transaction(), + ): + cursor = await 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() + ) + if await cursor.fetchone() is not None: + return + today = msk_today() + now = utc_now() 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, ?, ?) + 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, utc_now()), + (user_id, today, now), ) await conn.execute( - "UPDATE users SET is_admin = 1 WHERE user_id = ?", + "UPDATE users SET is_admin = TRUE WHERE user_id = %s", (user_id,), ) await conn.execute( - "INSERT OR REPLACE INTO app_settings (key, value)" - " VALUES ('admins_seeded', '1')" + """ + INSERT INTO app_settings (key, value) + VALUES ('admins_seeded', '1') + ON CONFLICT (key) DO UPDATE SET value = EXCLUDED.value + """ ) - 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: + 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 - ) + 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 - ) + return await self._update_user_flag("is_blocked", user_id, enabled) async def set_age_confirmed(self, user_id: int, confirmed: bool) -> bool: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( - "UPDATE users SET age_confirmed = ?, age_declined = 0" - " WHERE user_id = ?", - (int(confirmed), user_id), + """ + UPDATE users + SET age_confirmed = %s, age_declined = FALSE + WHERE user_id = %s + RETURNING user_id + """, + (confirmed, user_id), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None async def set_age_declined(self, user_id: int) -> bool: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( - "UPDATE users SET age_declined = 1, age_confirmed = 0" - " WHERE user_id = ?", + """ + UPDATE users + SET age_declined = TRUE, age_confirmed = FALSE + WHERE user_id = %s + RETURNING user_id + """, (user_id,), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None async def set_credits(self, user_id: int, credits: int) -> bool: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( - "UPDATE users SET credits = ? WHERE user_id = ?", + """ + UPDATE users SET credits = %s + WHERE user_id = %s + RETURNING user_id + """, (max(0, credits), user_id), ) - await conn.commit() - return cursor.rowcount > 0 + 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("Недопустимое поле") - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( - f"UPDATE users SET {column} = ? WHERE user_id = ?", - (int(enabled), user_id), + f""" + UPDATE users SET {column} = %s + WHERE user_id = %s + RETURNING user_id + """, + (enabled, user_id), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None 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: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( """ - UPDATE users SET referrer_id = ? - WHERE user_id = ? AND referrer_id IS NULL AND user_id != ? + 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), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None async def mark_first_payment(self, user_id: int) -> bool: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( - "UPDATE users SET first_paid_at = ?" - " WHERE user_id = ? AND first_paid_at IS NULL", + """ + UPDATE users SET first_paid_at = %s + WHERE user_id = %s AND first_paid_at IS NULL + RETURNING user_id + """, (utc_now(), user_id), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None 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() + 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: - try: - last = date.fromisoformat(row["last_daily_credit"]) - days = max(1, (current - last).days) - except ValueError: - days = 1 + 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 + ?," - " last_daily_credit = ? WHERE user_id = ?", - (days, current_text, int(row["user_id"])), + """ + UPDATE users + SET credits = credits + %s, last_daily_credit = %s + WHERE user_id = %s + """, + (days, current, 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: + 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] @@ -637,13 +685,13 @@ class Database: clauses: list[str] = [] params: list[Any] = [] if user_id is not None: - clauses.append("user_id = ?") + clauses.append("user_id = %s") params.append(user_id) if status: - clauses.append("status = ?") + clauses.append("status = %s") params.append(status) if mode: - clauses.append("mode = ?") + clauses.append("mode = %s") params.append(mode) where = f" WHERE {' AND '.join(clauses)}" if clauses else "" return where, params @@ -657,12 +705,12 @@ class Database: 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 ?" + f"SELECT * FROM requests{where} ORDER BY id DESC LIMIT %s OFFSET %s" ) - async with conn.execute(query, (*params, limit, offset)) as cursor: + 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] @@ -673,23 +721,26 @@ class Database: 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: + 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) -> 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: + async with self._connection_pool().connection() as conn: + cursor = await conn.execute( + "SELECT status, COUNT(*) AS total FROM requests 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" + ) row = await cursor.fetchone() return { "total": sum(by_status.values()), @@ -698,101 +749,105 @@ class Database: } 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: + async with self._connection_pool().connection() as conn: + cursor = await 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 %s + """, + (limit,), + ) 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"]), + last_seen=_as_iso(row["last_seen"]) or "", ) 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 < ? + WHERE is_premium = TRUE AND status IN ('queued', 'running') + AND id < %s """ - 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 < ?)) + AND (is_premium = TRUE OR (is_premium = FALSE AND id < %s)) """ - params = (request_id,) - async with conn.execute(query, params) as cursor: + 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]: - 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"]) + 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: + 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), + comfy_prompt_id = COALESCE(%s, comfy_prompt_id), error_text = NULL - WHERE id = ? + WHERE id = %s """, (comfy_prompt_id, request_id), ) - async def set_comfy_prompt_id(self, request_id: int, comfy_prompt_id: str) -> None: + 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 = ?", + "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 = ? WHERE id = ?", + "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 = 1 WHERE id = ?", + "UPDATE requests SET payment_refunded = TRUE WHERE id = %s", (request_id,), ) @@ -808,14 +863,14 @@ class Database: amount: int, is_premium: bool = False, ) -> PendingPayment: - conn = self._connection() - async with self._write_lock: - await conn.execute( + 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) - VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending', ?) + VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 'pending', %s) + RETURNING * """, ( payload, @@ -826,56 +881,59 @@ class Database: username, input_file_path, amount, - int(is_premium), + is_premium, ), ) - await conn.commit() - pending = await self.get_pending_payment(payload) - if pending is None: - raise RuntimeError("Не удалось прочитать счёт") - return pending + 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: - conn = self._connection() - async with conn.execute( - "SELECT * FROM pending_payments WHERE payload = ?", (payload,) - ) as cursor: + 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: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( """ UPDATE pending_payments - SET is_premium = ?, amount = ? - WHERE payload = ? AND status = 'pending' AND amount = 0 + SET is_premium = %s, amount = %s + WHERE payload = %s AND status = 'pending' AND amount = 0 + RETURNING payload """, - (int(is_premium), amount, payload), + (is_premium, amount, payload), ) - await conn.commit() - return cursor.rowcount > 0 + return await cursor.fetchone() is not None async def mark_pending_paid(self, payload: str, charge_id: str) -> bool: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: cursor = await conn.execute( """ UPDATE pending_payments - SET status = 'paid', telegram_payment_charge_id = ? - WHERE payload = ? AND status = 'pending' + SET status = 'paid', telegram_payment_charge_id = %s + WHERE payload = %s AND status = 'pending' + RETURNING payload """, (charge_id, payload), ) - await conn.commit() - return cursor.rowcount > 0 + 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 = ?", + """ + UPDATE pending_payments SET status = 'refunded' + WHERE payload = %s + """, (payload,), ) @@ -883,14 +941,16 @@ class Database: await self._update( """ UPDATE pending_payments SET status = 'cancelled' - WHERE payload = ? AND status = 'pending' + WHERE payload = %s AND status = 'pending' """, (payload,), ) - async def attach_pending_request(self, payload: str, request_id: int) -> None: + async def attach_pending_request( + self, payload: str, request_id: int + ) -> None: await self._update( - "UPDATE pending_payments SET request_id = ? WHERE payload = ?", + "UPDATE pending_payments SET request_id = %s WHERE payload = %s", (request_id, payload), ) @@ -899,10 +959,10 @@ class Database: """ UPDATE requests SET status = 'done', - output_file_path = ?, + output_file_path = %s, error_text = NULL, - finished_at = ? - WHERE id = ? + finished_at = %s + WHERE id = %s """, (output_file_path, utc_now(), request_id), ) @@ -912,9 +972,9 @@ class Database: """ UPDATE requests SET status = 'failed', - error_text = ?, - finished_at = ? - WHERE id = ? + error_text = %s, + finished_at = %s + WHERE id = %s """, (error_text[:1000], utc_now(), request_id), ) @@ -928,44 +988,44 @@ class Database: mode: str | None = None, has_person: bool = False, ) -> int: - conn = self._connection() - async with self._write_lock: + 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 (?, ?, ?, ?, ?, ?) + VALUES (%s, %s, %s, %s, %s, %s) + RETURNING id """, - (user_id, utc_now(), prompt, reason, mode, int(has_person)), + (user_id, utc_now(), prompt, reason, mode, has_person), ) - await conn.commit() - return int(cursor.lastrowid) + row = await cursor.fetchone() + if row is None: + raise RuntimeError("Не удалось записать блокировку") + return int(row["id"]) 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: + 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]: + 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() + 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: - conn = self._connection() - async with self._write_lock: + async with self._connection_pool().connection() as conn: await conn.execute(sql, params) - await conn.commit() diff --git a/docker-compose.yml b/docker-compose.yml index d920a4f..441b5a0 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -1,8 +1,26 @@ services: + db: + image: postgres:17-alpine + environment: + POSTGRES_DB: ${POSTGRES_DB:-mpakoimageaibot} + POSTGRES_USER: ${POSTGRES_USER:-mpako} + POSTGRES_PASSWORD: ${POSTGRES_PASSWORD:?POSTGRES_PASSWORD должен быть задан в .env} + volumes: + - ./data/postgres:/var/lib/postgresql/data + healthcheck: + test: ["CMD-SHELL", "pg_isready -U $${POSTGRES_USER} -d $${POSTGRES_DB}"] + interval: 5s + timeout: 5s + retries: 10 + restart: unless-stopped + bot: image: mpakoimageaibot:latest build: . env_file: .env + depends_on: + db: + condition: service_healthy extra_hosts: - "host.docker.internal:host-gateway" volumes: @@ -15,6 +33,9 @@ services: build: . command: ["uvicorn", "web:app", "--host", "0.0.0.0", "--port", "8080"] env_file: .env + depends_on: + db: + condition: service_healthy ports: - "${WEB_BIND_IP:-127.0.0.1}:8099:8080" volumes: diff --git a/plan.md b/plan.md index d0f6dba..bfe8376 100644 --- a/plan.md +++ b/plan.md @@ -6,7 +6,7 @@ Telegram-бот для генерации изображений через уд - Python 3.10+ - python-telegram-bot 21+ (async) -- SQLite (aiosqlite) +- PostgreSQL (psycopg 3 + psycopg-pool) - httpx для вызова ComfyUI API - Docker для упаковки @@ -23,11 +23,11 @@ Telegram-бот для генерации изображений через уд ## Хранение -SQLite, таблица `requests`: +PostgreSQL, таблица `requests`: - `id` — первичный ключ; - `user_id` — Telegram ID; -- `timestamp` — время создания в ISO; +- `timestamp` — время создания, `TIMESTAMPTZ` (UTC); - `prompt_text` — текст запроса; - `input_file_path` — входное изображение, может быть пустым; - `output_file_path` — результат, может быть пустым; diff --git a/requirements-dev.txt b/requirements-dev.txt index 2968596..0a2a937 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,3 +1,4 @@ -r requirements.txt pytest>=8,<9 pytest-asyncio>=0.24,<1 +ruff>=0.8,<1 diff --git a/requirements.txt b/requirements.txt index 3e349ba..a93f530 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,7 @@ python-telegram-bot>=21.6,<23 httpx[socks]>=0.27,<1 -aiosqlite>=0.20,<1 +psycopg[binary]>=3.2,<4 +psycopg-pool>=3.2,<4 websockets>=13,<16 python-dotenv>=1.0,<2 fastapi>=0.115,<1 diff --git a/scripts/__init__.py b/scripts/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/scripts/migrate_sqlite_to_postgres.py b/scripts/migrate_sqlite_to_postgres.py new file mode 100644 index 0000000..2843ffb --- /dev/null +++ b/scripts/migrate_sqlite_to_postgres.py @@ -0,0 +1,344 @@ +from __future__ import annotations + +import argparse +import os +import sqlite3 +import sys +from datetime import UTC, date, datetime +from pathlib import Path +from typing import Any + +import psycopg +from psycopg.rows import dict_row + +PROJECT_ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(PROJECT_ROOT)) + +from db import MIGRATIONS, SCHEMA_STATEMENTS, _normalize_dsn + +REQUEST_COLUMNS = ( + "id", + "user_id", + "username", + "timestamp", + "prompt_text", + "enhanced_prompt", + "input_file_path", + "output_file_path", + "status", + "comfy_prompt_id", + "payment_charge_id", + "payment_refunded", + "is_premium", + "mode", + "seed", + "error_text", + "finished_at", +) + +USER_COLUMNS = ( + "user_id", + "username", + "credits", + "last_daily_credit", + "referrer_id", + "first_paid_at", + "created_at", + "is_admin", + "is_blocked", + "age_confirmed", + "age_declined", +) + +BLOCKED_COLUMNS = ( + "id", + "user_id", + "timestamp", + "prompt", + "reason", + "mode", + "has_person", +) + +PAYMENT_COLUMNS = ( + "payload", + "user_id", + "timestamp", + "mode", + "prompt", + "username", + "input_file_path", + "amount", + "status", + "is_premium", + "telegram_payment_charge_id", + "request_id", +) + +SETTING_COLUMNS = ("key", "value") + + +def parse_dt(value: Any) -> datetime | None: + if value in (None, ""): + return None + if isinstance(value, datetime): + parsed = value + else: + parsed = datetime.fromisoformat(str(value)) + if parsed.tzinfo is None: + parsed = parsed.replace(tzinfo=UTC) + return parsed + + +def parse_day(value: Any) -> date | None: + if value in (None, ""): + return None + if isinstance(value, datetime): + return value.date() + if isinstance(value, date): + return value + return date.fromisoformat(str(value)) + + +def parse_bool(value: Any) -> bool: + return bool(int(value or 0)) + + +def table_names(connection: sqlite3.Connection) -> set[str]: + rows = connection.execute( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ).fetchall() + return {str(row["name"]) for row in rows} + + +def read_rows( + connection: sqlite3.Connection, + table: str, + columns: tuple[str, ...], + defaults: dict[str, Any], +) -> list[dict[str, Any]]: + if table not in table_names(connection): + print(f" {table}: в SQLite нет таблицы, пропуск") + return [] + existing = { + str(row["name"]) + for row in connection.execute(f"PRAGMA table_info({table})") + } + select = ", ".join( + column if column in existing else f":{column} AS {column}" + for column in columns + ) + params = { + column: defaults[column] + for column in columns + if column not in existing + } + query = f"SELECT {select} FROM {table}" + cursor = ( + connection.execute(query, params) + if params + else connection.execute(query) + ) + return [dict(row) for row in cursor.fetchall()] + + +def insert_rows( + connection: psycopg.Connection, + table: str, + columns: tuple[str, ...], + rows: list[dict[str, Any]], + conflict: str | None, +) -> int: + if not rows: + return 0 + placeholders = ", ".join(["%s"] * len(columns)) + conflict_clause = ( + f"ON CONFLICT ({conflict}) DO NOTHING" + if conflict + else "ON CONFLICT DO NOTHING" + ) + query = ( + f"INSERT INTO {table} ({', '.join(columns)}) " + f"VALUES ({placeholders}) {conflict_clause}" + ) + with connection.cursor() as cursor: + cursor.executemany( + query, + [tuple(row[column] for column in columns) for row in rows], + ) + return cursor.rowcount + + +def ensure_schema(connection: psycopg.Connection) -> None: + with connection.cursor() as cursor: + for statement in (*SCHEMA_STATEMENTS, *MIGRATIONS): + cursor.execute(statement) + + +def reset_sequence(connection: psycopg.Connection, table: str) -> None: + with connection.cursor() as cursor: + cursor.execute( + f""" + SELECT setval( + pg_get_serial_sequence('{table}', 'id'), + COALESCE((SELECT MAX(id) FROM {table}), 1), + (SELECT MAX(id) IS NOT NULL FROM {table}) + ) + """ + ) + + +def migrate(sqlite_path: Path, database_url: str) -> None: + if not sqlite_path.exists(): + raise SystemExit(f"SQLite-база не найдена: {sqlite_path}") + + source = sqlite3.connect( + f"{sqlite_path.resolve().as_uri()}?mode=ro", uri=True + ) + source.row_factory = sqlite3.Row + try: + users = read_rows( + source, + "users", + USER_COLUMNS, + { + "username": None, + "credits": 3, + "last_daily_credit": date.today().isoformat(), + "referrer_id": None, + "first_paid_at": None, + "created_at": datetime.now().astimezone().isoformat(), + "is_admin": 0, + "is_blocked": 0, + "age_confirmed": 0, + "age_declined": 0, + }, + ) + requests = read_rows( + source, + "requests", + REQUEST_COLUMNS, + { + "username": None, + "enhanced_prompt": None, + "payment_charge_id": None, + "payment_refunded": 0, + "is_premium": 0, + "seed": None, + "error_text": None, + "finished_at": None, + }, + ) + blocked = read_rows( + source, + "blocked_requests", + BLOCKED_COLUMNS, + {"mode": None, "has_person": 0}, + ) + payments = read_rows( + source, + "pending_payments", + PAYMENT_COLUMNS, + {"username": None, "input_file_path": None, "is_premium": 0}, + ) + settings = read_rows(source, "app_settings", SETTING_COLUMNS, {}) + finally: + source.close() + + for row in users: + row["last_daily_credit"] = parse_day(row["last_daily_credit"]) + row["first_paid_at"] = parse_dt(row["first_paid_at"]) + row["created_at"] = parse_dt(row["created_at"]) + row["is_admin"] = parse_bool(row["is_admin"]) + row["is_blocked"] = parse_bool(row["is_blocked"]) + row["age_confirmed"] = parse_bool(row["age_confirmed"]) + row["age_declined"] = parse_bool(row["age_declined"]) + + for row in requests: + row["timestamp"] = parse_dt(row["timestamp"]) + row["finished_at"] = parse_dt(row["finished_at"]) + row["payment_refunded"] = parse_bool(row["payment_refunded"]) + row["is_premium"] = parse_bool(row["is_premium"]) + + for row in blocked: + row["timestamp"] = parse_dt(row["timestamp"]) + row["has_person"] = parse_bool(row["has_person"]) + + for row in payments: + row["timestamp"] = parse_dt(row["timestamp"]) + row["is_premium"] = parse_bool(row["is_premium"]) + + dsn = _normalize_dsn(database_url) + with psycopg.connect( + dsn, row_factory=dict_row, autocommit=True + ) as target: + ensure_schema(target) + with target.transaction(): + inserted_users = insert_rows( + target, "users", USER_COLUMNS, users, "user_id" + ) + inserted_requests = insert_rows( + target, "requests", REQUEST_COLUMNS, requests, None + ) + inserted_blocked = insert_rows( + target, + "blocked_requests", + BLOCKED_COLUMNS, + blocked, + None, + ) + inserted_payments = insert_rows( + target, + "pending_payments", + PAYMENT_COLUMNS, + payments, + "payload", + ) + inserted_settings = insert_rows( + target, + "app_settings", + SETTING_COLUMNS, + settings, + "key", + ) + reset_sequence(target, "requests") + reset_sequence(target, "blocked_requests") + + print(f"SQLite: {sqlite_path}") + print("PostgreSQL: подключение выполнено") + print(f"users: {inserted_users} из {len(users)}") + print(f"requests: {inserted_requests} из {len(requests)}") + print(f"blocked_requests: {inserted_blocked} из {len(blocked)}") + print(f"pending_payments: {inserted_payments} из {len(payments)}") + print(f"app_settings: {inserted_settings} из {len(settings)}") + print("Готово.") + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Перенос данных из SQLite в PostgreSQL" + ) + parser.add_argument( + "--sqlite", + default=os.getenv("SQLITE_PATH", str(PROJECT_ROOT / "data" / "bot.db")), + help="путь к старой SQLite-базе", + ) + parser.add_argument( + "--database-url", + default=os.getenv("DATABASE_URL", ""), + help="строка подключения PostgreSQL", + ) + args = parser.parse_args() + + if not args.database_url: + print( + "DATABASE_URL не задан. Передай --database-url или переменную окружения.", + file=sys.stderr, + ) + raise SystemExit(1) + + migrate(Path(args.sqlite), args.database_url) + + +if __name__ == "__main__": + main() diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..c313288 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,62 @@ +from __future__ import annotations + +import os +from collections.abc import AsyncIterator, Iterator +from pathlib import Path + +import psycopg +import pytest +import pytest_asyncio + +from config import Settings +from db import Database + +TEST_DATABASE_URL = os.getenv("TEST_DATABASE_URL", "").strip() + +TABLES = ( + "requests", + "blocked_requests", + "pending_payments", + "users", + "app_settings", +) + + +def database_url() -> str: + if not TEST_DATABASE_URL: + pytest.skip("TEST_DATABASE_URL не задан") + return TEST_DATABASE_URL + + +def drop_tables(url: str) -> None: + with psycopg.connect(url, autocommit=True) as connection: + for table in TABLES: + connection.execute(f"DROP TABLE IF EXISTS {table} CASCADE") + + +@pytest_asyncio.fixture +async def db() -> AsyncIterator[Database]: + url = database_url() + drop_tables(url) + database = Database(url) + await database.connect() + try: + yield database + finally: + await database.close() + drop_tables(url) + + +@pytest.fixture +def web_settings(tmp_path: Path) -> Iterator[Settings]: + url = database_url() + drop_tables(url) + settings = Settings( + bot_token="", + database_url=url, + storage_dir=tmp_path / "storage", + ) + try: + yield settings + finally: + drop_tables(url) diff --git a/tests/test_db.py b/tests/test_db.py index 7b13182..eda13dd 100644 --- a/tests/test_db.py +++ b/tests/test_db.py @@ -1,20 +1,10 @@ from __future__ import annotations -import sqlite3 -from datetime import date, timedelta +from datetime import timedelta import pytest -import pytest_asyncio -from db import Database, UserAlreadyQueued - - -@pytest_asyncio.fixture -async def db(tmp_path): - database = Database(tmp_path / "bot.db") - await database.connect() - yield database - await database.close() +from db import Database, UserAlreadyQueued, msk_today async def test_create_and_read(db: Database) -> None: @@ -275,7 +265,7 @@ async def test_user_credits_referral_and_daily_accrual(db: Database) -> None: assert friend.referrer_id == 1 assert friend.first_paid_at is not None - future = date.today() + timedelta(days=2) + future = msk_today() + timedelta(days=2) assert await db.accrue_daily_credits(future) >= 2 user = await db.get_user(1) assert user is not None @@ -283,46 +273,17 @@ async def test_user_credits_referral_and_daily_accrual(db: Database) -> None: assert await db.accrue_daily_credits(future) == 0 -async def test_migrates_legacy_schema_without_username(tmp_path) -> None: - path = tmp_path / "legacy.db" - connection = sqlite3.connect(path) - connection.execute( - """ - CREATE TABLE requests ( - id INTEGER PRIMARY KEY AUTOINCREMENT, - user_id INTEGER NOT NULL, - 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 - ) - """ - ) - connection.commit() - connection.close() +async def test_ensure_user_updates_username_without_resetting_credits( + db: Database, +) -> None: + user = await db.ensure_user(1, "alice") + assert user.credits == 3 - database = Database(path) - await database.connect() - try: - request = await database.create_request( - user_id=1, - username="alice", - mode="t2i", - prompt_text="apple", - input_file_path=None, - seed=1, - ) - assert request.username == "alice" - await database.set_enhanced_prompt(request.id, "a red apple") - loaded = await database.get_request(request.id) - assert loaded is not None - assert loaded.enhanced_prompt == "a red apple" - finally: - await database.close() + assert await db.add_credits(1, 2) == 5 + user = await db.ensure_user(1, "alice_new") + assert user.username == "alice_new" + assert user.credits == 5 + + user = await db.ensure_user(1) + assert user.username == "alice_new" + assert user.credits == 5 diff --git a/tests/test_migration.py b/tests/test_migration.py new file mode 100644 index 0000000..a7e4d89 --- /dev/null +++ b/tests/test_migration.py @@ -0,0 +1,102 @@ +from __future__ import annotations + +import sqlite3 +from datetime import UTC +from pathlib import Path + +from scripts.migrate_sqlite_to_postgres import ( + REQUEST_COLUMNS, + USER_COLUMNS, + parse_bool, + parse_dt, + read_rows, +) + +REQUEST_DEFAULTS = { + "username": None, + "enhanced_prompt": None, + "payment_charge_id": None, + "payment_refunded": 0, + "is_premium": 0, + "seed": None, + "error_text": None, + "finished_at": None, +} + +USER_DEFAULTS = { + "username": None, + "credits": 3, + "last_daily_credit": "2026-10-18", + "referrer_id": None, + "first_paid_at": None, + "created_at": "2026-10-18T12:00:00+00:00", + "is_admin": 0, + "is_blocked": 0, + "age_confirmed": 0, + "age_declined": 0, +} + + +def test_read_rows_fills_legacy_defaults(tmp_path: Path) -> None: + path = tmp_path / "legacy.db" + connection = sqlite3.connect(path) + connection.row_factory = sqlite3.Row + connection.executescript( + """ + CREATE TABLE requests ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + user_id INTEGER NOT NULL, + timestamp TEXT NOT NULL, + prompt_text TEXT NOT NULL, + input_file_path TEXT, + output_file_path TEXT, + status TEXT NOT NULL DEFAULT 'queued', + comfy_prompt_id TEXT, + mode TEXT NOT NULL, + seed INTEGER, + error_text TEXT, + finished_at TEXT + ); + CREATE TABLE 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 + ); + """ + ) + connection.execute( + "INSERT INTO requests (user_id, timestamp, prompt_text, mode, seed)" + " VALUES (1, '2026-10-18T12:00:00+00:00', 'apple', 't2i', 42)" + ) + connection.execute( + "INSERT INTO users" + " (user_id, username, credits, last_daily_credit, created_at)" + " VALUES (1, 'alice', 3, '2026-10-18'," + " '2026-10-18T12:00:00+00:00')" + ) + connection.commit() + + requests = read_rows( + connection, "requests", REQUEST_COLUMNS, REQUEST_DEFAULTS + ) + users = read_rows(connection, "users", USER_COLUMNS, USER_DEFAULTS) + connection.close() + + assert requests[0]["username"] is None + assert requests[0]["payment_refunded"] == 0 + assert requests[0]["is_premium"] == 0 + assert users[0]["is_admin"] == 0 + assert users[0]["age_declined"] == 0 + assert users[0]["username"] == "alice" + + +def test_parse_helpers() -> None: + parsed = parse_dt("2026-10-18T12:00:00") + assert parsed is not None + assert parsed.tzinfo is UTC + assert parse_bool("1") is True + assert parse_bool(0) is False diff --git a/tests/test_queue.py b/tests/test_queue.py index 1a53984..273b569 100644 --- a/tests/test_queue.py +++ b/tests/test_queue.py @@ -1,11 +1,11 @@ from __future__ import annotations import asyncio +import os import time from pathlib import Path import pytest -import pytest_asyncio from config import Settings from db import Database, UserAlreadyQueued @@ -56,19 +56,11 @@ class FakeEnhancer: return f"enhanced: {prompt}" -@pytest_asyncio.fixture -async def db(tmp_path): - database = Database(tmp_path / "bot.db") - await database.connect() - yield database - await database.close() - - def make_settings(tmp_path: Path, **overrides) -> Settings: values = { "bot_token": "test", "storage_dir": tmp_path / "storage", - "db_path": tmp_path / "bot.db", + "database_url": os.getenv("TEST_DATABASE_URL", ""), "workflow_t2i": BASE / "txt2img.json", "workflow_i2i": BASE / "img2img.json", "workflow_mapping": BASE / "mapping.json", diff --git a/tests/test_web.py b/tests/test_web.py index 083153a..00c6d5a 100644 --- a/tests/test_web.py +++ b/tests/test_web.py @@ -1,27 +1,18 @@ from __future__ import annotations -import sqlite3 -from pathlib import Path +from datetime import UTC, date, datetime +import psycopg from fastapi.testclient import TestClient from PIL import Image +from psycopg.rows import dict_row from config import Settings from web import create_app -def make_settings(tmp_path: Path, **overrides) -> Settings: - values = { - "bot_token": "", - "db_path": tmp_path / "bot.db", - "storage_dir": tmp_path / "storage", - } - values.update(overrides) - return Settings(**values) - - def insert_request( - path: Path, + database_url: str, *, user_id: int, prompt: str, @@ -31,56 +22,75 @@ def insert_request( username: str | None = None, enhanced: str | None = None, ) -> int: - connection = sqlite3.connect(path) - try: - cursor = connection.execute( + with psycopg.connect(database_url, row_factory=dict_row) as connection: + row = connection.execute( """ INSERT INTO requests (user_id, username, timestamp, prompt_text, enhanced_prompt, status, mode, seed, output_file_path) - VALUES (?, ?, '2026-10-18T12:00:00+00:00', ?, ?, ?, ?, 1, ?) + VALUES (%s, %s, %s, %s, %s, %s, %s, 1, %s) + RETURNING id """, - (user_id, username, prompt, enhanced, status, mode, output), - ) - connection.commit() - return int(cursor.lastrowid) - finally: - connection.close() + ( + user_id, + username, + datetime(2026, 10, 18, 12, tzinfo=UTC), + prompt, + enhanced, + status, + mode, + output, + ), + ).fetchone() + if row is None: + raise RuntimeError("Не удалось вставить запрос") + return int(row["id"]) def insert_user( - path: Path, + database_url: str, *, user_id: int, username: str | None = "alice", credits: int = 3, - is_admin: int = 0, - is_blocked: int = 0, + is_admin: bool = False, + is_blocked: bool = False, ) -> None: - connection = sqlite3.connect(path) - try: + with psycopg.connect(database_url) as connection: connection.execute( """ INSERT INTO users (user_id, username, credits, last_daily_credit, created_at, is_admin, is_blocked) - VALUES (?, ?, ?, '2026-10-18', '2026-10-18T12:00:00+00:00', ?, ?) + VALUES (%s, %s, %s, %s, %s, %s, %s) + ON CONFLICT (user_id) DO NOTHING """, - (user_id, username, credits, is_admin, is_blocked), + ( + user_id, + username, + credits, + date(2026, 10, 18), + datetime(2026, 10, 18, 12, tzinfo=UTC), + is_admin, + is_blocked, + ), ) - connection.commit() - finally: - connection.close() -def test_history_page_and_user_filter(tmp_path: Path) -> None: - settings = make_settings(tmp_path) - app = create_app(settings) +def test_history_page_and_user_filter(web_settings: Settings) -> None: + app = create_app(web_settings) with TestClient(app) as client: insert_request( - settings.db_path, user_id=1, username="alice", prompt="apple prompt" + web_settings.database_url, + user_id=1, + username="alice", + prompt="apple prompt", + ) + insert_request( + web_settings.database_url, + user_id=2, + prompt="banana prompt", ) - insert_request(settings.db_path, user_id=2, prompt="banana prompt") response = client.get("/") assert response.status_code == 200 @@ -99,17 +109,16 @@ def test_history_page_and_user_filter(tmp_path: Path) -> None: assert "user_id должен состоять только из цифр" in response.text -def test_detail_and_media(tmp_path: Path) -> None: - settings = make_settings(tmp_path) - output = settings.storage_dir / "1" / "1_output.png" +def test_detail_and_media(web_settings: Settings) -> None: + output = web_settings.storage_dir / "1" / "1_output.png" output.parent.mkdir(parents=True) with Image.new("RGB", (320, 200), (180, 40, 40)) as image: image.save(output) - app = create_app(settings) + app = create_app(web_settings) with TestClient(app) as client: request_id = insert_request( - settings.db_path, + web_settings.database_url, user_id=1, prompt="detail prompt", enhanced="An enhanced English prompt", @@ -139,42 +148,59 @@ def test_detail_and_media(tmp_path: Path) -> None: assert missing.status_code == 404 -def test_user_management(tmp_path: Path) -> None: - settings = make_settings(tmp_path) - app = create_app(settings) +def test_user_management(web_settings: Settings) -> None: + app = create_app(web_settings) with TestClient(app) as client: - insert_user(settings.db_path, user_id=1, username="alice") - insert_user(settings.db_path, user_id=2, username="bob", credits=0) + insert_user( + web_settings.database_url, user_id=1, username="alice" + ) + insert_user( + web_settings.database_url, + user_id=2, + username="bob", + credits=0, + ) page = client.get("/users") assert page.status_code == 200 assert "@alice" in page.text assert "@bob" in page.text - assert client.post("/users/1/credits", data={"credits": "9"}).status_code == 200 + assert ( + client.post("/users/1/credits", data={"credits": "9"}).status_code + == 200 + ) assert client.post("/users/1/admin").status_code == 200 assert client.post("/users/2/block").status_code == 200 - connection = sqlite3.connect(settings.db_path) - try: + with psycopg.connect( + web_settings.database_url, row_factory=dict_row + ) as connection: first = connection.execute( - "SELECT credits, is_admin, is_blocked FROM users WHERE user_id = 1" + "SELECT credits, is_admin, is_blocked FROM users" + " WHERE user_id = %s", + (1,), ).fetchone() second = connection.execute( - "SELECT is_blocked FROM users WHERE user_id = 2" + "SELECT is_blocked FROM users WHERE user_id = %s", (2,) ).fetchone() - finally: - connection.close() - assert first == (9, 1, 0) - assert second == (1,) + assert first is not None + assert second is not None + assert first["credits"] == 9 + assert first["is_admin"] is True + assert first["is_blocked"] is False + assert second["is_blocked"] is True -def test_user_search(tmp_path: Path) -> None: - settings = make_settings(tmp_path) - app = create_app(settings) +def test_user_search(web_settings: Settings) -> None: + app = create_app(web_settings) with TestClient(app) as client: - insert_user(settings.db_path, user_id=1, username="alice") - insert_user(settings.db_path, user_id=2, username="bob") + insert_user( + web_settings.database_url, user_id=1, username="alice" + ) + insert_user( + web_settings.database_url, user_id=2, username="bob" + ) page = client.get("/users?q=alice") assert page.status_code == 200 diff --git a/web.py b/web.py index 2de967c..413eae1 100644 --- a/web.py +++ b/web.py @@ -77,7 +77,11 @@ def create_app(settings: Settings | None = None) -> FastAPI: @asynccontextmanager async def lifespan(app: FastAPI): - db = Database(settings.db_path) + db = Database( + settings.database_url, + min_size=settings.db_pool_min_size, + max_size=settings.db_pool_max_size, + ) await db.connect() app.state.db = db yield