invix_test_task/bot/app/main.py

102 lines
3.5 KiB
Python

import asyncio
import logging
import asyncpg
from aiogram import Bot, Dispatcher
from aiogram.types import BotCommand
from aiogram.webhook.aiohttp_server import SimpleRequestHandler, setup_application
from aiohttp import web
from app.api import ShopApi
from app.config import settings
from app.dedup import DedupMiddleware
from app.handlers import router
log = logging.getLogger("shop-bot")
LEADER_LOCK = (1, 42001)
async def _wait_for_db() -> asyncpg.Connection:
delay = 1
while True:
try:
return await asyncpg.connect(settings.database_url)
except (OSError, asyncpg.PostgresError) as exc:
log.warning("postgres is not ready: %s", exc)
await asyncio.sleep(delay)
delay = min(delay * 2, 10)
async def _become_leader(conn: asyncpg.Connection) -> None:
while True:
locked = await conn.fetchval("SELECT pg_try_advisory_lock($1, $2)", *LEADER_LOCK)
if locked:
log.info("this replica is the polling leader")
return
log.info("another replica holds getUpdates, waiting")
await asyncio.sleep(2)
async def _run_bot(bot: Bot, dp: Dispatcher, api: ShopApi) -> None:
pool = await asyncpg.create_pool(settings.database_url, min_size=1, max_size=4)
dp.update.middleware(DedupMiddleware(pool))
dp.include_router(router)
await bot.set_my_commands(
[
BotCommand(command="catalog", description="Каталог"),
BotCommand(command="cart", description="Корзина"),
BotCommand(command="checkout", description="Оформить заказ"),
]
)
try:
if settings.bot_mode == "webhook":
if not settings.webhook_url:
raise RuntimeError("WEBHOOK_URL is required when BOT_MODE=webhook")
await bot.set_webhook(
settings.webhook_url,
secret_token=settings.bot_secret[:256],
drop_pending_updates=False,
)
application = web.Application()
SimpleRequestHandler(dispatcher=dp, bot=bot, secret_token=settings.bot_secret[:256]).register(
application, path=settings.webhook_path
)
setup_application(application, dp, bot=bot)
runner = web.AppRunner(application)
await runner.setup()
await web.TCPSite(runner, "0.0.0.0", 8081).start()
log.info("webhook listening on :8081%s", settings.webhook_path)
await asyncio.Event().wait()
else:
await bot.delete_webhook(drop_pending_updates=False)
await dp.start_polling(bot, api=api)
finally:
await api.close()
await pool.close()
async def main() -> None:
logging.basicConfig(level=logging.INFO)
if not settings.bot_token:
log.warning("BOT_TOKEN is empty; bot replicas stay idle. The website works without it.")
await asyncio.Event().wait()
return
api = ShopApi(settings.api_url, settings.bot_secret)
bot = Bot(settings.bot_token)
dp = Dispatcher()
leader: asyncpg.Connection | None = None
try:
if settings.bot_mode != "webhook":
leader = await _wait_for_db()
await _become_leader(leader)
await _run_bot(bot, dp, api)
finally:
await bot.session.close()
if leader is not None:
await leader.close()
if __name__ == "__main__":
asyncio.run(main())