139 lines
4.6 KiB
Python
139 lines
4.6 KiB
Python
"""Long-polling inbox (§3.8).
|
|
|
|
Every update lands in ``telegram_updates`` before the offset moves past it;
|
|
a worker handles rows from the table, oldest first, and finishes whatever a
|
|
previous process left unprocessed at startup.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import logging
|
|
from datetime import UTC, datetime
|
|
from typing import TYPE_CHECKING
|
|
|
|
from aiogram.exceptions import TelegramConflictError, TelegramNetworkError
|
|
from aiogram.types import Update
|
|
from sqlalchemy import func
|
|
from sqlalchemy.exc import IntegrityError
|
|
from sqlmodel import col, select
|
|
|
|
from beaver_gateway.storage.models import TelegramUpdate
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
from aiogram import Bot
|
|
|
|
from beaver_gateway.storage.db import Database
|
|
|
|
__all__ = ["Inbox"]
|
|
|
|
_log = logging.getLogger("beaver_gateway.frontends.telegram.inbox")
|
|
|
|
ALLOWED_UPDATES = ("message", "callback_query")
|
|
|
|
|
|
class Inbox:
|
|
def __init__(
|
|
self,
|
|
db: Database,
|
|
bot: Bot,
|
|
*,
|
|
handler: Callable[[Update], Awaitable[None]],
|
|
poll_timeout: int = 30,
|
|
) -> None:
|
|
self._db = db
|
|
self._bot = bot
|
|
self._handler = handler
|
|
self._poll_timeout = poll_timeout
|
|
self._wake = asyncio.Event()
|
|
|
|
async def run(self) -> None:
|
|
async with asyncio.TaskGroup() as tg:
|
|
tg.create_task(self._poll())
|
|
tg.create_task(self._work())
|
|
|
|
async def _poll(self) -> None:
|
|
offset = await self._next_offset()
|
|
backoff = 1.0
|
|
while True:
|
|
try:
|
|
updates = await self._bot.get_updates(
|
|
offset=offset,
|
|
timeout=self._poll_timeout,
|
|
allowed_updates=list(ALLOWED_UPDATES),
|
|
request_timeout=self._poll_timeout + 10,
|
|
)
|
|
except TelegramConflictError:
|
|
_log.error("another poller holds this bot token; retrying in 10s")
|
|
await asyncio.sleep(10)
|
|
continue
|
|
except (TimeoutError, TelegramNetworkError, OSError) as exc:
|
|
_log.warning("getUpdates failed (%s); retrying in %.0fs", exc, backoff)
|
|
await asyncio.sleep(backoff)
|
|
backoff = min(backoff * 2, 60.0)
|
|
continue
|
|
backoff = 1.0
|
|
for update in updates:
|
|
await self._store(update)
|
|
offset = update.update_id + 1
|
|
if updates:
|
|
self._wake.set()
|
|
|
|
async def _store(self, update: Update) -> None:
|
|
row = TelegramUpdate(
|
|
update_id=update.update_id,
|
|
payload=update.model_dump(mode="json", by_alias=True, exclude_none=True),
|
|
)
|
|
async with self._db.session() as session:
|
|
session.add(row)
|
|
try:
|
|
await session.commit()
|
|
except IntegrityError:
|
|
await session.rollback()
|
|
|
|
async def _next_offset(self) -> int | None:
|
|
async with self._db.session() as session:
|
|
latest = (
|
|
await session.exec(select(func.max(col(TelegramUpdate.update_id))))
|
|
).one()
|
|
return int(latest) + 1 if latest is not None else None
|
|
|
|
async def _work(self) -> None:
|
|
while True:
|
|
rows = await self._pending()
|
|
if not rows:
|
|
self._wake.clear()
|
|
with contextlib.suppress(TimeoutError):
|
|
await asyncio.wait_for(self._wake.wait(), timeout=5.0)
|
|
continue
|
|
for row in rows:
|
|
await self._handle(row)
|
|
|
|
async def _pending(self, limit: int = 50) -> list[TelegramUpdate]:
|
|
async with self._db.session() as session:
|
|
result = await session.exec(
|
|
select(TelegramUpdate)
|
|
.where(col(TelegramUpdate.processed_at).is_(None))
|
|
.order_by(col(TelegramUpdate.update_id))
|
|
.limit(limit)
|
|
)
|
|
return list(result.all())
|
|
|
|
async def _handle(self, row: TelegramUpdate) -> None:
|
|
error: str | None = None
|
|
try:
|
|
await self._handler(Update.model_validate(row.payload))
|
|
except Exception as exc: # noqa: BLE001
|
|
error = f"{type(exc).__name__}: {exc}"[:500]
|
|
_log.exception("update %s failed", row.update_id)
|
|
async with self._db.session() as session:
|
|
stored = await session.get(TelegramUpdate, row.update_id)
|
|
if stored is not None:
|
|
stored.processed_at = datetime.now(UTC)
|
|
stored.error = error
|
|
session.add(stored)
|
|
await session.commit()
|