Переход на PostgreSQL
Build and push Docker image / tests (push) Failing after 50s Details
Build and push Docker image / docker (push) Has been skipped Details

Заменил SQLite на PostgreSQL 17: psycopg и пул соединений, BOOLEAN/TIMESTAMPTZ/BIGSERIAL, отдельный сервис db в Compose. Добавил скрипт переноса старых данных и тесты БД в CI.
This commit is contained in:
MpakobecEx 2026-10-19 02:41:04 +05:00
parent 7fc148a661
commit 44e77e83b1
19 changed files with 1236 additions and 580 deletions

View File

@ -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

View File

@ -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:

View File

@ -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:<sha>`

6
bot.py
View File

@ -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

View File

@ -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` и `:<sha>` в 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.
- Админ-панель.

View File

@ -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"),

778
db.py

File diff suppressed because it is too large Load Diff

View File

@ -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:

View File

@ -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` — результат, может быть пустым;

View File

@ -1,3 +1,4 @@
-r requirements.txt
pytest>=8,<9
pytest-asyncio>=0.24,<1
ruff>=0.8,<1

View File

@ -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

0
scripts/__init__.py Normal file
View File

View File

@ -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()

62
tests/conftest.py Normal file
View File

@ -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)

View File

@ -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

102
tests/test_migration.py Normal file
View File

@ -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

View File

@ -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",

View File

@ -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

6
web.py
View File

@ -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