From d5994e732d7b1dfa469cf400132ba49c8f75315e Mon Sep 17 00:00:00 2001 From: Tolmachev Igor Date: Mon, 23 Mar 2026 18:40:40 +0300 Subject: Add new_invoice command --- alembic/versions/1627487324fd_init_database.py | 72 ++++++++++++++++ alembic/versions/940afc736a0f_init_database.py | 72 ---------------- handlers/admin/__init__.py | 2 + handlers/admin/new_announcement.py | 9 +- handlers/admin/new_invoice.py | 109 +++++++++++++++++++++++++ libs/msg.py | 37 ++++++++- models/__init__.py | 2 + models/callback_data.py | 5 ++ models/invoce.py | 13 ++- models/rich_text.py | 10 ++- 10 files changed, 250 insertions(+), 81 deletions(-) create mode 100644 alembic/versions/1627487324fd_init_database.py delete mode 100644 alembic/versions/940afc736a0f_init_database.py create mode 100644 handlers/admin/new_invoice.py create mode 100644 models/callback_data.py diff --git a/alembic/versions/1627487324fd_init_database.py b/alembic/versions/1627487324fd_init_database.py new file mode 100644 index 0000000..f25277f --- /dev/null +++ b/alembic/versions/1627487324fd_init_database.py @@ -0,0 +1,72 @@ +"""init database + +Revision ID: 1627487324fd +Revises: +Create Date: 2026-03-23 18:33:08.493629 + +""" + +from typing import Sequence, Union + +import sqlalchemy as sa + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "1627487324fd" +down_revision: Union[str, Sequence[str], None] = None +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Upgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "announcement", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("message", sa.JSON(), nullable=False), + sa.Column("datetime", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_announcement")), + ) + op.create_table( + "invoice", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("message", sa.JSON(), nullable=False), + sa.Column("datetime", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_invoice")), + ) + op.create_table( + "user", + sa.Column("id", sa.Integer(), autoincrement=False, nullable=False), + sa.Column("role", sa.Enum("REGULAR", "ADMIN", name="userrole"), nullable=False), + sa.Column("vpn_link", sa.String(), nullable=False), + sa.Column("datetime", sa.DateTime(), nullable=False), + sa.PrimaryKeyConstraint("id", name=op.f("pk_user")), + ) + op.create_table( + "payment", + sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), + sa.Column("user_id", sa.Integer(), nullable=False), + sa.Column("invoice_id", sa.Integer(), nullable=False), + sa.Column("receipt_file_id", sa.String(), nullable=False), + sa.Column("datetime", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint( + ["invoice_id"], ["invoice.id"], name=op.f("fk_payment_invoice_id_invoice") + ), + sa.ForeignKeyConstraint( + ["user_id"], ["user.id"], name=op.f("fk_payment_user_id_user") + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_payment")), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade schema.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("payment") + op.drop_table("user") + op.drop_table("invoice") + op.drop_table("announcement") + # ### end Alembic commands ### diff --git a/alembic/versions/940afc736a0f_init_database.py b/alembic/versions/940afc736a0f_init_database.py deleted file mode 100644 index bfaa771..0000000 --- a/alembic/versions/940afc736a0f_init_database.py +++ /dev/null @@ -1,72 +0,0 @@ -"""init database - -Revision ID: 940afc736a0f -Revises: -Create Date: 2026-03-23 02:21:35.245634 - -""" - -from typing import Sequence, Union - -import sqlalchemy as sa - -from alembic import op - -# revision identifiers, used by Alembic. -revision: str = "940afc736a0f" -down_revision: Union[str, Sequence[str], None] = None -branch_labels: Union[str, Sequence[str], None] = None -depends_on: Union[str, Sequence[str], None] = None - - -def upgrade() -> None: - """Upgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.create_table( - "announcement", - sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("message", sa.JSON(), nullable=False), - sa.Column("datetime", sa.DateTime(), nullable=False), - sa.PrimaryKeyConstraint("id", name=op.f("pk_announcement")), - ) - op.create_table( - "invoice", - sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("amount", sa.Float(), nullable=False), - sa.Column("datetime", sa.DateTime(), nullable=False), - sa.PrimaryKeyConstraint("id", name=op.f("pk_invoice")), - ) - op.create_table( - "user", - sa.Column("id", sa.Integer(), autoincrement=False, nullable=False), - sa.Column("role", sa.Enum("REGULAR", "ADMIN", name="userrole"), nullable=False), - sa.Column("vpn_link", sa.String(), nullable=False), - sa.Column("datetime", sa.DateTime(), nullable=False), - sa.PrimaryKeyConstraint("id", name=op.f("pk_user")), - ) - op.create_table( - "payment", - sa.Column("id", sa.Integer(), autoincrement=True, nullable=False), - sa.Column("user_id", sa.Integer(), nullable=False), - sa.Column("invoice_id", sa.Integer(), nullable=False), - sa.Column("receipt_file_id", sa.String(), nullable=False), - sa.Column("datetime", sa.DateTime(), nullable=False), - sa.ForeignKeyConstraint( - ["invoice_id"], ["invoice.id"], name=op.f("fk_payment_invoice_id_invoice") - ), - sa.ForeignKeyConstraint( - ["user_id"], ["user.id"], name=op.f("fk_payment_user_id_user") - ), - sa.PrimaryKeyConstraint("id", name=op.f("pk_payment")), - ) - # ### end Alembic commands ### - - -def downgrade() -> None: - """Downgrade schema.""" - # ### commands auto generated by Alembic - please adjust! ### - op.drop_table("payment") - op.drop_table("user") - op.drop_table("invoice") - op.drop_table("announcement") - # ### end Alembic commands ### diff --git a/handlers/admin/__init__.py b/handlers/admin/__init__.py index 2f7c74f..98b127f 100644 --- a/handlers/admin/__init__.py +++ b/handlers/admin/__init__.py @@ -3,6 +3,7 @@ from aiogram.filters import MagicData # isort: off from . import new_announcement +from . import new_invoice # isort: on router = Router(name="admin") @@ -11,4 +12,5 @@ router.callback_query.filter(MagicData(F.user.is_admin())) router.include_routers( new_announcement.router, + new_invoice.router, ) diff --git a/handlers/admin/new_announcement.py b/handlers/admin/new_announcement.py index 79cf8d4..0920c47 100644 --- a/handlers/admin/new_announcement.py +++ b/handlers/admin/new_announcement.py @@ -1,6 +1,7 @@ from datetime import UTC, datetime from aiogram import Bot, F, Router +from aiogram.enums import ButtonStyle from aiogram.exceptions import TelegramAPIError from aiogram.filters import Command from aiogram.fsm.context import FSMContext @@ -17,7 +18,7 @@ from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from libs.fsm import get_data, set_data -from libs.msg import publish +from libs.msg import publish_announcement from models import Announcement, RichText, User router = Router(name="new_announcement") @@ -42,8 +43,8 @@ async def new_announcement_command(msg: Message, state: FSMContext) -> None: reply_markup=ReplyKeyboardMarkup( keyboard=[ [ - KeyboardButton(text=SEND_BUTTON), - KeyboardButton(text=CANCEL_BUTTON), + KeyboardButton(text=SEND_BUTTON, style=ButtonStyle.SUCCESS), + KeyboardButton(text=CANCEL_BUTTON, style=ButtonStyle.DANGER), ] ], resize_keyboard=True, @@ -69,7 +70,7 @@ async def announcement_send( status_template = "Публикация анонса...\nОпубликовано: {}" status_msg = await msg.answer(status_template.format(0)) - async for n in publish(bot, users, data.rich_text): + async for n in publish_announcement(bot, users, data.rich_text): try: await status_msg.edit_text(status_template.format(n)) except TelegramAPIError: diff --git a/handlers/admin/new_invoice.py b/handlers/admin/new_invoice.py new file mode 100644 index 0000000..7e1a64d --- /dev/null +++ b/handlers/admin/new_invoice.py @@ -0,0 +1,109 @@ +from datetime import UTC, datetime + +from aiogram import Bot, F, Router +from aiogram.enums import ButtonStyle, ContentType +from aiogram.exceptions import TelegramAPIError +from aiogram.filters import Command +from aiogram.fsm.context import FSMContext +from aiogram.fsm.state import State, StatesGroup +from aiogram.types import ( + KeyboardButton, + Message, + ReplyKeyboardMarkup, + ReplyKeyboardRemove, +) +from pydantic import BaseModel +from sqlalchemy import select +from sqlalchemy.ext.asyncio import AsyncSession + +from libs.fsm import get_data, set_data +from libs.msg import send_invoice +from models import Invoice, RichText, User + +router = Router(name="new_invoice") + + +class NewInvoiceStates(StatesGroup): + message = State() + + +class NewInvoiceData(BaseModel): + rich_text: RichText | None = None + + +CREATE_BUTTON = "Создать" +CANCEL_BUTTON = "Отменить создание" + + +@router.message(Command("new_invoice")) +async def new_invoice_command(msg: Message, state: FSMContext) -> None: + await msg.answer( + "Укажите сообщение для создания счёта", + reply_markup=ReplyKeyboardMarkup( + keyboard=[ + [ + KeyboardButton(text=CREATE_BUTTON, style=ButtonStyle.SUCCESS), + KeyboardButton(text=CANCEL_BUTTON, style=ButtonStyle.DANGER), + ] + ], + resize_keyboard=True, + ), + ) + await state.set_state(NewInvoiceStates.message) + + +@router.message(NewInvoiceStates.message, F.text == CREATE_BUTTON) +async def invoice_send( + msg: Message, + bot: Bot, + state: FSMContext, + session: AsyncSession, +) -> None: + users = await session.scalars(select(User.id).where(User.id != msg.chat.id)) + data = await get_data(state, NewInvoiceData) + + if data.rich_text is None: + await msg.answer("Для создания счёта укажите сообщение.") + return + + status_template = "Рассылка счёта...\nОтправлено: {}" + status_msg = await msg.answer(status_template.format(0)) + + invoice = Invoice(message=data, datetime=datetime.now(UTC)) + session.add(invoice) + await session.flush() + + async for n in send_invoice(bot, users, data.rich_text, invoice.id): + try: + await status_msg.edit_text(status_template.format(n)) + except TelegramAPIError: + pass + + await status_msg.delete() + await msg.answer( + "Счёт отправлен всем пользователям", + # reply_markup=ReplyKeyboardRemove(), + ) + # await state.clear() + + +@router.message(NewInvoiceStates.message, F.text == CANCEL_BUTTON) +async def invoice_cancel(msg: Message, state: FSMContext) -> None: + await msg.answer("Создание счёта отменено", reply_markup=ReplyKeyboardRemove()) + await state.clear() + + +@router.message(NewInvoiceStates.message) +async def invoice_message(msg: Message, bot: Bot, state: FSMContext) -> None: + if msg.content_type != ContentType.TEXT or msg.text is None: + await msg.answer( + "Неверный тип сообщения.\n" + "Бот поддерживает отправку только текстовых счетов." + ) + return + + rich_text = RichText.from_message(msg) + await set_data(state, NewInvoiceData(rich_text=rich_text)) + + msg_rich_text = RichText.from_text("Сообщение вашего счёта:\n", rich_text) + await msg_rich_text.send(bot, msg.chat.id) diff --git a/libs/msg.py b/libs/msg.py index 9bcc52a..05bddfc 100644 --- a/libs/msg.py +++ b/libs/msg.py @@ -2,12 +2,15 @@ import asyncio from typing import AsyncGenerator, Iterable from aiogram import Bot +from aiogram.enums import ButtonStyle from aiogram.exceptions import TelegramAPIError, TelegramRetryAfter +from aiogram.types import InlineKeyboardButton, InlineKeyboardMarkup from models import RichText +from models.callback_data import PayInvoiceData -async def publish( +async def publish_announcement( bot: Bot, users: Iterable[int], rich_text: RichText, @@ -23,3 +26,35 @@ async def publish( await asyncio.sleep(5) yield n + + +async def send_invoice( + bot: Bot, + users: Iterable[int], + rich_text: RichText, + invoice_id: int, +) -> AsyncGenerator[int]: + callback_data = PayInvoiceData(invoice_id=invoice_id).pack() + reply_markup = InlineKeyboardMarkup( + inline_keyboard=[ + [ + InlineKeyboardButton( + text="Оплатить", + style=ButtonStyle.PRIMARY, + callback_data=callback_data, + ) + ] + ] + ) + + for n, user_id in enumerate(users, start=1): + for _ in range(5): + try: + await rich_text.send(bot, user_id, reply_markup=reply_markup) + break + except TelegramRetryAfter as e: + await asyncio.sleep(e.retry_after + 1) + except TelegramAPIError: + await asyncio.sleep(5) + + yield n diff --git a/models/__init__.py b/models/__init__.py index 9d56b33..0547429 100644 --- a/models/__init__.py +++ b/models/__init__.py @@ -5,6 +5,7 @@ from .user import User from .invoce import Invoice from .payment import Payment from .announcement import Announcement +from . import callback_data # isort: on __all__ = [ @@ -14,4 +15,5 @@ __all__ = [ "Payment", "RichText", "Announcement", + "callback_data", ] diff --git a/models/callback_data.py b/models/callback_data.py new file mode 100644 index 0000000..d3e6d61 --- /dev/null +++ b/models/callback_data.py @@ -0,0 +1,5 @@ +from aiogram.filters.callback_data import CallbackData + + +class PayInvoiceData(CallbackData, prefix="pay_invoice"): + invoice_id: int diff --git a/models/invoce.py b/models/invoce.py index 54a445e..66ff019 100644 --- a/models/invoce.py +++ b/models/invoce.py @@ -1,13 +1,22 @@ from datetime import datetime +from sqlalchemy import JSON from sqlalchemy.orm import Mapped, mapped_column -from models import BaseTable +from models import BaseTable, RichText class Invoice(BaseTable): __tablename__ = "invoice" id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - amount: Mapped[float] + __message: Mapped[str] = mapped_column("message", JSON()) datetime: Mapped[datetime] + + @property + def message(self) -> RichText: + return RichText.model_validate_json(self.__message) + + @message.setter + def message(self, value: RichText) -> None: + self.__message = value.model_dump_json() diff --git a/models/rich_text.py b/models/rich_text.py index e5441b6..6d1c1cb 100644 --- a/models/rich_text.py +++ b/models/rich_text.py @@ -1,7 +1,7 @@ from typing import Self from aiogram import Bot -from aiogram.types import Message, MessageEntity +from aiogram.types import Message, MessageEntity, ReplyMarkupUnion from pydantic import BaseModel @@ -34,10 +34,16 @@ class RichText(BaseModel): return result - async def send(self, bot: Bot, chat_id: int) -> Message: + async def send( + self, + bot: Bot, + chat_id: int, + reply_markup: ReplyMarkupUnion | None = None, + ) -> Message: return await bot.send_message( chat_id=chat_id, text=self.text, entities=self.entities, parse_mode=None, + reply_markup=reply_markup, ) -- cgit v1.3