Оплата генераций звёздами
Счёт в XTR, возврат при ошибке, бесплатный доступ для админов.
This commit is contained in:
parent
e90e891c3b
commit
960672818f
|
|
@ -40,6 +40,10 @@ MODERATION_ENABLED=true
|
|||
# Дополнительные имена публичных людей через запятую
|
||||
BLOCKED_PUBLIC_FIGURES=
|
||||
|
||||
# Оплата Telegram Stars
|
||||
PAYMENTS_ENABLED=true
|
||||
STARS_PRICE=5
|
||||
|
||||
# Веб-история: только локальная сеть, без авторизации
|
||||
WEB_BIND_IP=192.168.31.158
|
||||
WEB_PAGE_SIZE=50
|
||||
|
|
|
|||
17
README.md
17
README.md
|
|
@ -10,6 +10,7 @@ Telegram-бот для генерации и редактирования изо
|
|||
- **Веб-история**: просмотр запросов с фильтрами по `user_id`, статусу и режиму, просмотр изображений.
|
||||
- **Улучшение промптов**: перевод и переписывание запроса на английский через `TextGenerate` + `qwen3vl_4b` в ComfyUI.
|
||||
- **Фильтрация промптов**: блокировка CSAM, зоофилии, NSFW с публичными людьми и обходных вариантов для не-админов.
|
||||
- **Оплата**: 5 Telegram Stars за генерацию, админам бесплатно, возврат при ошибке.
|
||||
- Глобальная очередь с одной активной задачей на пользователя.
|
||||
- Позиция в очереди и прогресс генерации через WebSocket ComfyUI.
|
||||
- SQLite: запросы, статусы, seed, пути к файлам.
|
||||
|
|
@ -121,6 +122,22 @@ BLOCKED_PUBLIC_FIGURES=
|
|||
|
||||
При блокировке пользователь получает «Запрос нарушает правила. Звёзды не списаны.», а запись с `user_id`, временем, промптом и причиной сохраняется в таблицу `blocked_requests`.
|
||||
|
||||
## Оплата
|
||||
|
||||
Генерация стоит 5 Telegram Stars:
|
||||
|
||||
```dotenv
|
||||
PAYMENTS_ENABLED=true
|
||||
STARS_PRICE=5
|
||||
```
|
||||
|
||||
- Админы из `ADMIN_USER_IDS` генерируют бесплатно.
|
||||
- Не-админу выставляется счёт через Telegram Stars (`XTR`).
|
||||
- Фильтрация промпта выполняется до выставления счёта.
|
||||
- Если поставить задачу в очередь не удалось, звёзды возвращаются.
|
||||
- При ошибке генерации звёзды также возвращаются.
|
||||
- Оплата и `telegram_payment_charge_id` сохраняются в `pending_payments` и `requests`.
|
||||
|
||||
## Веб-история
|
||||
|
||||
После `docker compose up -d` доступна страница:
|
||||
|
|
|
|||
217
bot.py
217
bot.py
|
|
@ -6,7 +6,7 @@ import time
|
|||
import uuid
|
||||
from pathlib import Path
|
||||
|
||||
from telegram import InputFile, Update
|
||||
from telegram import InputFile, LabeledPrice, Update
|
||||
from telegram.constants import ChatAction, ChatType
|
||||
from telegram.error import BadRequest, NetworkError, TelegramError, TimedOut
|
||||
from telegram.ext import (
|
||||
|
|
@ -14,13 +14,14 @@ from telegram.ext import (
|
|||
CommandHandler,
|
||||
ContextTypes,
|
||||
MessageHandler,
|
||||
PreCheckoutQueryHandler,
|
||||
filters,
|
||||
)
|
||||
|
||||
from comfy_client import ComfyClient
|
||||
from config import Settings, setup_logging
|
||||
from db import Database, Request, UserAlreadyQueued
|
||||
from moderation import PromptBlocked, detect_person
|
||||
from moderation import PromptBlocked, detect_person, is_prompt_allowed
|
||||
from prompt_enhancer import PromptEnhancer
|
||||
from queue_manager import QueueEvent, QueueFull, QueueManager
|
||||
from workflows import WorkflowBuilder
|
||||
|
|
@ -67,6 +68,12 @@ class MpakoBot:
|
|||
application.add_handler(CommandHandler("start", self._start))
|
||||
application.add_handler(CommandHandler("help", self._help))
|
||||
application.add_handler(CommandHandler("status", self._status))
|
||||
application.add_handler(
|
||||
PreCheckoutQueryHandler(self._pre_checkout)
|
||||
)
|
||||
application.add_handler(
|
||||
MessageHandler(filters.SUCCESSFUL_PAYMENT, self._successful_payment)
|
||||
)
|
||||
application.add_handler(MessageHandler(filters.PHOTO, self._photo))
|
||||
application.add_handler(
|
||||
MessageHandler(filters.TEXT & ~filters.COMMAND, self._text)
|
||||
|
|
@ -111,6 +118,15 @@ class MpakoBot:
|
|||
await self._db.close()
|
||||
logger.info("Бот остановлен")
|
||||
|
||||
def _start_text(self) -> str:
|
||||
text = START_TEXT
|
||||
if self._settings.payments_enabled:
|
||||
text += (
|
||||
f"\n\nСтоимость генерации: {self._settings.stars_price} звёзд. "
|
||||
"Админам бесплатно."
|
||||
)
|
||||
return text
|
||||
|
||||
async def _start(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
) -> None:
|
||||
|
|
@ -118,7 +134,7 @@ class MpakoBot:
|
|||
return
|
||||
message = update.effective_message
|
||||
if message is not None:
|
||||
await message.reply_text(START_TEXT)
|
||||
await message.reply_text(self._start_text())
|
||||
|
||||
async def _help(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
|
|
@ -127,7 +143,7 @@ class MpakoBot:
|
|||
return
|
||||
message = update.effective_message
|
||||
if message is not None:
|
||||
await message.reply_text(START_TEXT)
|
||||
await message.reply_text(self._start_text())
|
||||
|
||||
async def _status(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
|
|
@ -150,6 +166,76 @@ class MpakoBot:
|
|||
f"Задача в очереди. Перед тобой: {position} {self._people(position)}."
|
||||
)
|
||||
|
||||
async def _pre_checkout(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
) -> None:
|
||||
query = update.pre_checkout_query
|
||||
if query is None:
|
||||
return
|
||||
pending = await self._db.get_pending_payment(query.invoice_payload)
|
||||
allowed = (
|
||||
pending is not None
|
||||
and pending.user_id == query.from_user.id
|
||||
and pending.status == "pending"
|
||||
and query.currency == "XTR"
|
||||
and query.total_amount == pending.amount
|
||||
)
|
||||
try:
|
||||
await query.answer(
|
||||
ok=allowed,
|
||||
error_message=None
|
||||
if allowed
|
||||
else "Счёт устарел. Отправь запрос заново.",
|
||||
)
|
||||
except TelegramError:
|
||||
logger.exception("Не удалось ответить на pre_checkout_query")
|
||||
|
||||
async def _successful_payment(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
) -> None:
|
||||
message = update.effective_message
|
||||
user = update.effective_user
|
||||
payment = message.successful_payment if message is not None else None
|
||||
if message is None or user is None or payment is None:
|
||||
return
|
||||
|
||||
payload = payment.invoice_payload
|
||||
pending = await self._db.get_pending_payment(payload)
|
||||
if pending is None or pending.user_id != user.id:
|
||||
logger.error(
|
||||
"Платёж %s не найден или принадлежит другому пользователю",
|
||||
payload,
|
||||
)
|
||||
return
|
||||
if pending.status == "paid":
|
||||
await message.reply_text("Этот платёж уже обработан.")
|
||||
return
|
||||
if pending.status != "pending":
|
||||
return
|
||||
|
||||
charge_id = payment.telegram_payment_charge_id
|
||||
if payment.currency != "XTR" or payment.total_amount != pending.amount:
|
||||
refunded = await self._refund_stars(user.id, charge_id)
|
||||
if refunded:
|
||||
await self._db.mark_pending_refunded(payload)
|
||||
note = " Звёзды возвращены." if refunded else ""
|
||||
await message.reply_text("Платёж не принят." + note)
|
||||
return
|
||||
|
||||
if not await self._db.mark_pending_paid(payload, charge_id):
|
||||
await message.reply_text("Этот платёж уже обработан.")
|
||||
return
|
||||
|
||||
await self._submit(
|
||||
update,
|
||||
mode=pending.mode,
|
||||
prompt=pending.prompt,
|
||||
input_path=pending.input_file_path,
|
||||
skip_payment=True,
|
||||
payment_charge_id=charge_id,
|
||||
payment_payload=payload,
|
||||
)
|
||||
|
||||
async def _text(
|
||||
self, update: Update, context: ContextTypes.DEFAULT_TYPE
|
||||
) -> None:
|
||||
|
|
@ -290,6 +376,9 @@ class MpakoBot:
|
|||
mode: str,
|
||||
prompt: str,
|
||||
input_path: str | None,
|
||||
skip_payment: bool = False,
|
||||
payment_charge_id: str | None = None,
|
||||
payment_payload: str | None = None,
|
||||
) -> bool:
|
||||
message = update.effective_message
|
||||
user = update.effective_user
|
||||
|
|
@ -306,6 +395,59 @@ class MpakoBot:
|
|||
):
|
||||
has_person = await asyncio.to_thread(detect_person, input_path)
|
||||
|
||||
if (
|
||||
not skip_payment
|
||||
and not is_admin
|
||||
and self._settings.payments_enabled
|
||||
):
|
||||
allowed, reason = is_prompt_allowed(
|
||||
prompt,
|
||||
mode=mode,
|
||||
has_person=has_person,
|
||||
extra_public_figures=self._settings.blocked_public_figures,
|
||||
)
|
||||
if not allowed:
|
||||
await self._db.log_blocked(
|
||||
user_id=user.id,
|
||||
prompt=prompt,
|
||||
reason=reason or "unknown",
|
||||
mode=mode,
|
||||
has_person=has_person,
|
||||
)
|
||||
await message.reply_text(BLOCKED_TEXT)
|
||||
return False
|
||||
|
||||
payload = uuid.uuid4().hex
|
||||
amount = self._settings.stars_price
|
||||
await self._db.create_pending_payment(
|
||||
payload=payload,
|
||||
user_id=user.id,
|
||||
mode=mode,
|
||||
prompt=prompt,
|
||||
username=user.username,
|
||||
input_file_path=input_path,
|
||||
amount=amount,
|
||||
)
|
||||
try:
|
||||
await self._retry(
|
||||
self._bot.send_invoice,
|
||||
chat_id=user.id,
|
||||
title="Генерация изображения",
|
||||
description="Одна генерация или редактирование изображения.",
|
||||
payload=payload,
|
||||
provider_token="",
|
||||
currency="XTR",
|
||||
prices=[LabeledPrice(label="Генерация", amount=amount)],
|
||||
)
|
||||
except TelegramError:
|
||||
logger.exception("Не удалось выставить счёт пользователю %s", user.id)
|
||||
await self._db.mark_pending_cancelled(payload)
|
||||
await message.reply_text(
|
||||
"Не удалось выставить счёт. Попробуй позже."
|
||||
)
|
||||
return False
|
||||
return True
|
||||
|
||||
try:
|
||||
request, position = await self._queue.enqueue(
|
||||
user_id=user.id,
|
||||
|
|
@ -315,11 +457,18 @@ class MpakoBot:
|
|||
input_file_path=input_path,
|
||||
has_person=has_person,
|
||||
is_admin=is_admin,
|
||||
payment_charge_id=payment_charge_id,
|
||||
)
|
||||
except PromptBlocked:
|
||||
await message.reply_text(BLOCKED_TEXT)
|
||||
note = await self._refund_if_paid(
|
||||
user.id, payment_charge_id, payment_payload
|
||||
)
|
||||
await message.reply_text(BLOCKED_TEXT + note)
|
||||
return False
|
||||
except UserAlreadyQueued:
|
||||
note = await self._refund_if_paid(
|
||||
user.id, payment_charge_id, payment_payload
|
||||
)
|
||||
active = await self._db.get_active_by_user(user.id)
|
||||
text = "Ты уже в очереди."
|
||||
if active is not None:
|
||||
|
|
@ -328,16 +477,27 @@ class MpakoBot:
|
|||
else:
|
||||
ahead = await self._queue.position(active.id)
|
||||
text += f" Перед тобой: {ahead} {self._people(ahead)}."
|
||||
await message.reply_text(text)
|
||||
await message.reply_text(text + note)
|
||||
return False
|
||||
except QueueFull as exc:
|
||||
await message.reply_text(exc.user_message)
|
||||
note = await self._refund_if_paid(
|
||||
user.id, payment_charge_id, payment_payload
|
||||
)
|
||||
await message.reply_text(exc.user_message + note)
|
||||
return False
|
||||
except Exception:
|
||||
logger.exception("Не удалось поставить задачу пользователю %s", user.id)
|
||||
await message.reply_text("Не удалось поставить задачу. Попробуй позже.")
|
||||
note = await self._refund_if_paid(
|
||||
user.id, payment_charge_id, payment_payload
|
||||
)
|
||||
await message.reply_text(
|
||||
"Не удалось поставить задачу. Попробуй позже." + note
|
||||
)
|
||||
return False
|
||||
|
||||
if payment_payload:
|
||||
await self._db.attach_pending_request(payment_payload, request.id)
|
||||
|
||||
try:
|
||||
status_message = await self._retry(
|
||||
message.reply_text, self._ack_text(position)
|
||||
|
|
@ -367,6 +527,35 @@ class MpakoBot:
|
|||
)
|
||||
return True
|
||||
|
||||
async def _refund_if_paid(
|
||||
self,
|
||||
user_id: int,
|
||||
charge_id: str | None,
|
||||
payload: str | None,
|
||||
) -> str:
|
||||
if not charge_id:
|
||||
return ""
|
||||
refunded = await self._refund_stars(user_id, charge_id)
|
||||
if refunded and payload:
|
||||
await self._db.mark_pending_refunded(payload)
|
||||
if refunded:
|
||||
return " Звёзды возвращены."
|
||||
return " Не удалось вернуть звёзды, напиши админу."
|
||||
|
||||
async def _refund_stars(self, user_id: int, charge_id: str) -> bool:
|
||||
if self._bot is None:
|
||||
return False
|
||||
try:
|
||||
return await self._retry(
|
||||
self._bot.refund_star_payment,
|
||||
user_id=user_id,
|
||||
telegram_payment_charge_id=charge_id,
|
||||
attempts=2,
|
||||
)
|
||||
except TelegramError:
|
||||
logger.exception("Не удалось вернуть звёзды пользователю %s", user_id)
|
||||
return False
|
||||
|
||||
def _ack_text(self, position: int) -> str:
|
||||
if position <= 0:
|
||||
return "Принято. Ты первый, начинаю…"
|
||||
|
|
@ -513,7 +702,17 @@ class MpakoBot:
|
|||
async def _finish_failed(self, request: Request, error: str) -> None:
|
||||
target = self._status_messages.get(request.id)
|
||||
chat_id = target[0] if target is not None else request.user_id
|
||||
text = f"Не получилось: {error}"
|
||||
refund_note = ""
|
||||
if request.payment_charge_id and not request.payment_refunded:
|
||||
refunded = await self._refund_stars(
|
||||
request.user_id, request.payment_charge_id
|
||||
)
|
||||
if refunded:
|
||||
await self._db.set_refunded(request.id)
|
||||
refund_note = " Звёзды возвращены."
|
||||
else:
|
||||
refund_note = " Не удалось вернуть звёзды, напиши админу."
|
||||
text = f"Не получилось: {error}{refund_note}"
|
||||
if target is not None and self._bot is not None:
|
||||
try:
|
||||
await self._retry(
|
||||
|
|
|
|||
17
checklist.md
17
checklist.md
|
|
@ -6,10 +6,10 @@
|
|||
|
||||
## Текущий статус
|
||||
|
||||
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14) и модерация (15).
|
||||
- Реализованы блоки 2–7, базовое хранилище (9), файлы Docker (10), тесты (11), README (12), веб-история (13), улучшение промптов (14), модерация (15) и оплата Stars (16).
|
||||
- Вживую проверены: T2I (Z-Image), I2I (Krea2), WebSocket-прогресс, улучшение промптов, Docker, CI и веб-история.
|
||||
- 45 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, вспомогательные функции ComfyUI.
|
||||
- Не проверено: очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
|
||||
- 46 тестов проходят: БД, workflows, очередь, веб-история, enhancer, модерация, платежи, вспомогательные функции ComfyUI.
|
||||
- Не проверено: реальная оплата Stars в Telegram, очередь из 3+ пользователей, падение ComfyUI, рестарт с `running`-задачей, детекция лица на реальном фото.
|
||||
|
||||
## 0. Допущения и решения по умолчанию
|
||||
|
||||
|
|
@ -329,6 +329,17 @@ ON requests(user_id) WHERE status IN ('queued', 'running');
|
|||
|
||||
**DoD:** небезопасный промпт не попадает в очередь, причина пишется в БД, админ может отправлять без фильтра.
|
||||
|
||||
## 16. Оплата Stars
|
||||
|
||||
- [x] `PAYMENTS_ENABLED` / `STARS_PRICE`, админы бесплатно.
|
||||
- [x] Invoice `XTR` + `pre_checkout_query` + `successful_payment`.
|
||||
- [x] `pending_payments`: payload, сумма, статус, charge_id, request_id.
|
||||
- [x] Возврат звёзд при ошибке постановки в очередь и генерации.
|
||||
- [x] `payment_charge_id` и `payment_refunded` в `requests`.
|
||||
- [x] Тесты жизненного цикла платежа.
|
||||
|
||||
**DoD:** не-админ получает счёт на 5 звёзд, после оплаты задача ставится в очередь; при сбое звёзды возвращаются.
|
||||
|
||||
## Порядок работ
|
||||
|
||||
1. Блок 1 (подготовка) — без него нельзя валидировать клиент и workflow.
|
||||
|
|
|
|||
|
|
@ -94,6 +94,8 @@ class Settings:
|
|||
max_photo_bytes: int = 15 * 1024 * 1024
|
||||
moderation_enabled: bool = True
|
||||
blocked_public_figures: frozenset[str] = field(default_factory=frozenset)
|
||||
payments_enabled: bool = True
|
||||
stars_price: int = 5
|
||||
web_page_size: int = 50
|
||||
web_bind_ip: str = "127.0.0.1"
|
||||
|
||||
|
|
@ -152,6 +154,8 @@ class Settings:
|
|||
max_photo_bytes=max(1, _int("MAX_PHOTO_BYTES", 15 * 1024 * 1024)),
|
||||
moderation_enabled=_bool("MODERATION_ENABLED", True),
|
||||
blocked_public_figures=_csv("BLOCKED_PUBLIC_FIGURES"),
|
||||
payments_enabled=_bool("PAYMENTS_ENABLED", True),
|
||||
stars_price=max(1, _int("STARS_PRICE", 5)),
|
||||
web_page_size=max(10, min(200, _int("WEB_PAGE_SIZE", 50))),
|
||||
web_bind_ip=os.getenv("WEB_BIND_IP", "127.0.0.1").strip()
|
||||
or "127.0.0.1",
|
||||
|
|
|
|||
153
db.py
153
db.py
|
|
@ -37,6 +37,8 @@ class Request:
|
|||
output_file_path: str | None
|
||||
status: str
|
||||
comfy_prompt_id: str | None
|
||||
payment_charge_id: str | None
|
||||
payment_refunded: bool
|
||||
mode: str
|
||||
seed: int | None
|
||||
error_text: str | None
|
||||
|
|
@ -59,6 +61,8 @@ class Request:
|
|||
output_file_path=row["output_file_path"],
|
||||
status=row["status"],
|
||||
comfy_prompt_id=row["comfy_prompt_id"],
|
||||
payment_charge_id=row["payment_charge_id"],
|
||||
payment_refunded=bool(row["payment_refunded"]),
|
||||
mode=row["mode"],
|
||||
seed=row["seed"],
|
||||
error_text=row["error_text"],
|
||||
|
|
@ -101,6 +105,37 @@ class BlockedRequest:
|
|||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PendingPayment:
|
||||
payload: str
|
||||
user_id: int
|
||||
timestamp: str
|
||||
mode: str
|
||||
prompt: str
|
||||
username: str | None
|
||||
input_file_path: str | None
|
||||
amount: int
|
||||
status: str
|
||||
telegram_payment_charge_id: str | None
|
||||
request_id: int | None
|
||||
|
||||
@classmethod
|
||||
def from_row(cls, row: aiosqlite.Row) -> PendingPayment:
|
||||
return cls(
|
||||
payload=row["payload"],
|
||||
user_id=row["user_id"],
|
||||
timestamp=row["timestamp"],
|
||||
mode=row["mode"],
|
||||
prompt=row["prompt"],
|
||||
username=row["username"],
|
||||
input_file_path=row["input_file_path"],
|
||||
amount=int(row["amount"]),
|
||||
status=row["status"],
|
||||
telegram_payment_charge_id=row["telegram_payment_charge_id"],
|
||||
request_id=row["request_id"],
|
||||
)
|
||||
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS requests (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
|
|
@ -114,6 +149,8 @@ CREATE TABLE IF NOT EXISTS requests (
|
|||
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,
|
||||
mode TEXT NOT NULL CHECK (mode IN ('t2i', 'i2i')),
|
||||
seed INTEGER,
|
||||
error_text TEXT,
|
||||
|
|
@ -141,6 +178,24 @@ CREATE TABLE IF NOT EXISTS blocked_requests (
|
|||
|
||||
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')),
|
||||
telegram_payment_charge_id TEXT,
|
||||
request_id INTEGER
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_pending_user_time
|
||||
ON pending_payments(user_id, timestamp);
|
||||
"""
|
||||
|
||||
|
||||
|
|
@ -173,6 +228,12 @@ class Database:
|
|||
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"
|
||||
)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._conn is not None:
|
||||
|
|
@ -193,6 +254,7 @@ class Database:
|
|||
prompt_text: str,
|
||||
input_file_path: str | None,
|
||||
seed: int,
|
||||
payment_charge_id: str | None = None,
|
||||
) -> Request:
|
||||
conn = self._connection()
|
||||
async with self._write_lock:
|
||||
|
|
@ -201,8 +263,8 @@ class Database:
|
|||
"""
|
||||
INSERT INTO requests
|
||||
(user_id, username, timestamp, prompt_text,
|
||||
input_file_path, status, mode, seed)
|
||||
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?)
|
||||
input_file_path, status, mode, seed, payment_charge_id)
|
||||
VALUES (?, ?, ?, ?, ?, 'queued', ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
user_id,
|
||||
|
|
@ -212,6 +274,7 @@ class Database:
|
|||
input_file_path,
|
||||
mode,
|
||||
seed,
|
||||
payment_charge_id,
|
||||
),
|
||||
)
|
||||
request_id = cursor.lastrowid
|
||||
|
|
@ -389,6 +452,92 @@ class Database:
|
|||
(prompt, request_id),
|
||||
)
|
||||
|
||||
async def set_refunded(self, request_id: int) -> None:
|
||||
await self._update(
|
||||
"UPDATE requests SET payment_refunded = 1 WHERE id = ?",
|
||||
(request_id,),
|
||||
)
|
||||
|
||||
async def create_pending_payment(
|
||||
self,
|
||||
*,
|
||||
payload: str,
|
||||
user_id: int,
|
||||
mode: str,
|
||||
prompt: str,
|
||||
username: str | None,
|
||||
input_file_path: str | None,
|
||||
amount: int,
|
||||
) -> PendingPayment:
|
||||
conn = self._connection()
|
||||
async with self._write_lock:
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO pending_payments
|
||||
(payload, user_id, timestamp, mode, prompt, username,
|
||||
input_file_path, amount, status)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, 'pending')
|
||||
""",
|
||||
(
|
||||
payload,
|
||||
user_id,
|
||||
utc_now(),
|
||||
mode,
|
||||
prompt,
|
||||
username,
|
||||
input_file_path,
|
||||
amount,
|
||||
),
|
||||
)
|
||||
await conn.commit()
|
||||
pending = await self.get_pending_payment(payload)
|
||||
if pending is None:
|
||||
raise RuntimeError("Не удалось прочитать счёт")
|
||||
return pending
|
||||
|
||||
async def get_pending_payment(self, payload: str) -> PendingPayment | None:
|
||||
conn = self._connection()
|
||||
async with conn.execute(
|
||||
"SELECT * FROM pending_payments WHERE payload = ?", (payload,)
|
||||
) as cursor:
|
||||
row = await cursor.fetchone()
|
||||
return PendingPayment.from_row(row) if row is not None else None
|
||||
|
||||
async def mark_pending_paid(self, payload: str, charge_id: str) -> bool:
|
||||
conn = self._connection()
|
||||
async with self._write_lock:
|
||||
cursor = await conn.execute(
|
||||
"""
|
||||
UPDATE pending_payments
|
||||
SET status = 'paid', telegram_payment_charge_id = ?
|
||||
WHERE payload = ? AND status = 'pending'
|
||||
""",
|
||||
(charge_id, payload),
|
||||
)
|
||||
await conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def mark_pending_refunded(self, payload: str) -> None:
|
||||
await self._update(
|
||||
"UPDATE pending_payments SET status = 'refunded' WHERE payload = ?",
|
||||
(payload,),
|
||||
)
|
||||
|
||||
async def mark_pending_cancelled(self, payload: str) -> None:
|
||||
await self._update(
|
||||
"""
|
||||
UPDATE pending_payments SET status = 'cancelled'
|
||||
WHERE payload = ? AND status = 'pending'
|
||||
""",
|
||||
(payload,),
|
||||
)
|
||||
|
||||
async def attach_pending_request(self, payload: str, request_id: int) -> None:
|
||||
await self._update(
|
||||
"UPDATE pending_payments SET request_id = ? WHERE payload = ?",
|
||||
(request_id, payload),
|
||||
)
|
||||
|
||||
async def set_done(self, request_id: int, output_file_path: str) -> None:
|
||||
await self._update(
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -107,6 +107,7 @@ class QueueManager:
|
|||
username: str | None = None,
|
||||
has_person: bool = False,
|
||||
is_admin: bool = False,
|
||||
payment_charge_id: str | None = None,
|
||||
) -> tuple[Request, int]:
|
||||
async with self._lock:
|
||||
if self._stopping:
|
||||
|
|
@ -141,6 +142,7 @@ class QueueManager:
|
|||
prompt_text=prompt,
|
||||
input_file_path=input_file_path,
|
||||
seed=seed,
|
||||
payment_charge_id=payment_charge_id,
|
||||
)
|
||||
self._queue.append(request.id)
|
||||
self._active_users.add(user_id)
|
||||
|
|
|
|||
|
|
@ -165,6 +165,51 @@ async def test_history_filters_and_stats(db: Database) -> None:
|
|||
assert alice.display_user == "@alice"
|
||||
|
||||
|
||||
async def test_pending_payment_lifecycle(db: Database) -> None:
|
||||
pending = await db.create_pending_payment(
|
||||
payload="payload-1",
|
||||
user_id=1,
|
||||
mode="t2i",
|
||||
prompt="apple",
|
||||
username="alice",
|
||||
input_file_path=None,
|
||||
amount=5,
|
||||
)
|
||||
assert pending.status == "pending"
|
||||
assert pending.amount == 5
|
||||
|
||||
assert await db.mark_pending_paid("payload-1", "charge-1") is True
|
||||
loaded = await db.get_pending_payment("payload-1")
|
||||
assert loaded is not None
|
||||
assert loaded.status == "paid"
|
||||
assert loaded.telegram_payment_charge_id == "charge-1"
|
||||
|
||||
await db.attach_pending_request("payload-1", 42)
|
||||
loaded = await db.get_pending_payment("payload-1")
|
||||
assert loaded is not None
|
||||
assert loaded.request_id == 42
|
||||
|
||||
await db.mark_pending_refunded("payload-1")
|
||||
loaded = await db.get_pending_payment("payload-1")
|
||||
assert loaded is not None
|
||||
assert loaded.status == "refunded"
|
||||
|
||||
request = await db.create_request(
|
||||
user_id=2,
|
||||
mode="t2i",
|
||||
prompt_text="paid apple",
|
||||
input_file_path=None,
|
||||
seed=1,
|
||||
payment_charge_id="charge-2",
|
||||
)
|
||||
assert request.payment_charge_id == "charge-2"
|
||||
assert request.payment_refunded is False
|
||||
await db.set_refunded(request.id)
|
||||
loaded_request = await db.get_request(request.id)
|
||||
assert loaded_request is not None
|
||||
assert loaded_request.payment_refunded is True
|
||||
|
||||
|
||||
async def test_migrates_legacy_schema_without_username(tmp_path) -> None:
|
||||
path = tmp_path / "legacy.db"
|
||||
connection = sqlite3.connect(path)
|
||||
|
|
|
|||
Loading…
Reference in New Issue