Переход на PostgreSQL
Заменил SQLite на PostgreSQL 17: psycopg и пул соединений, BOOLEAN/TIMESTAMPTZ/BIGSERIAL, отдельный сервис db в Compose. Добавил скрипт переноса старых данных и тесты БД в CI.
This commit is contained in:
parent
7fc148a661
commit
44e77e83b1
10
.env.example
10
.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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
32
README.md
32
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:<sha>`
|
||||
|
|
|
|||
6
bot.py
6
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
|
||||
|
|
|
|||
36
checklist.md
36
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` и `:<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.
|
||||
- Админ-панель.
|
||||
|
|
|
|||
10
config.py
10
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"),
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
6
plan.md
6
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` — результат, может быть пустым;
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
-r requirements.txt
|
||||
pytest>=8,<9
|
||||
pytest-asyncio>=0.24,<1
|
||||
ruff>=0.8,<1
|
||||
|
|
|
|||
|
|
@ -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,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()
|
||||
|
|
@ -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)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
6
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
|
||||
|
|
|
|||
Loading…
Reference in New Issue