134 lines
5.7 KiB
Python
134 lines
5.7 KiB
Python
from app.core.logger import make_log, logger
|
|
from app.core.models._telegram import Wrapped_CBotChat
|
|
from app.core.models.user import User
|
|
|
|
from aiogram import BaseMiddleware, types
|
|
from app.core.models.messages import KnownTelegramMessage
|
|
from datetime import datetime
|
|
|
|
# Bot handlers historically use synchronous SQLAlchemy patterns.
|
|
# Keep a dedicated sync engine/session for bot middleware to preserve legacy behavior.
|
|
import re
|
|
from typing import Optional
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.orm import sessionmaker, Session
|
|
from app.core.config import get_settings
|
|
|
|
|
|
_SYNC_ENGINE = None
|
|
_SYNC_FACTORY: Optional[sessionmaker] = None
|
|
|
|
|
|
def _to_sync_dsn(async_dsn: str) -> str:
|
|
# Convert postgresql+asyncpg:// to postgresql+psycopg2:// for synchronous engine
|
|
return re.sub(r"\+asyncpg", "+psycopg2", async_dsn)
|
|
|
|
|
|
def _ensure_sync_session() -> Session:
|
|
global _SYNC_ENGINE, _SYNC_FACTORY
|
|
if _SYNC_ENGINE is None or _SYNC_FACTORY is None:
|
|
settings = get_settings()
|
|
dsn = _to_sync_dsn(settings.DATABASE_URL)
|
|
_SYNC_ENGINE = create_engine(dsn, pool_pre_ping=True, future=True)
|
|
_SYNC_FACTORY = sessionmaker(bind=_SYNC_ENGINE, autocommit=False, autoflush=True)
|
|
return _SYNC_FACTORY()
|
|
|
|
|
|
class UserDataMiddleware(BaseMiddleware):
|
|
|
|
async def __call__(self, handler, event, data):
|
|
update_body = event.message or event.callback_query or getattr(event, 'inline_query', None) or getattr(event, 'pre_checkout_query', None)
|
|
if not update_body or getattr(update_body, 'from_user', None) is None:
|
|
return
|
|
if getattr(update_body.from_user, 'is_bot', False):
|
|
return
|
|
|
|
user_id = update_body.from_user.id
|
|
assert user_id >= 1
|
|
|
|
# Use sync session for bot handlers compatibility
|
|
from app.core.models.user.user import User as DbUser
|
|
from app.core.models.messages import KnownTelegramMessage as DbKnownMsg
|
|
from app.core.logging import logger as app_logger
|
|
|
|
session = _ensure_sync_session()
|
|
try:
|
|
# Load or create user
|
|
try:
|
|
user = session.query(DbUser).filter(DbUser.telegram_id == user_id).first()
|
|
except Exception as e:
|
|
await app_logger.aerror("Middleware get user failed", error=str(e))
|
|
user = None
|
|
|
|
if user is None:
|
|
await app_logger.adebug("Creating new user", telegram_id=user_id)
|
|
user = DbUser(
|
|
telegram_id=user_id,
|
|
username=getattr(update_body.from_user, 'username', None),
|
|
language_code=getattr(update_body.from_user, 'language_code', 'en'),
|
|
meta={
|
|
'first_name': getattr(update_body.from_user, 'first_name', '') or '',
|
|
'last_name': getattr(update_body.from_user, 'last_name', '') or '',
|
|
'username': getattr(update_body.from_user, 'username', None),
|
|
'language_code': getattr(update_body.from_user, 'language_code', None),
|
|
'is_premium': getattr(update_body.from_user, 'is_premium', False),
|
|
}
|
|
)
|
|
session.add(user)
|
|
session.commit()
|
|
else:
|
|
# Update username/metadata
|
|
changed = False
|
|
if user.username != getattr(update_body.from_user, 'username', None):
|
|
user.username = getattr(update_body.from_user, 'username', None)
|
|
changed = True
|
|
meta = dict(user.meta or {})
|
|
if meta.get('first_name') != getattr(update_body.from_user, 'first_name', None):
|
|
meta['first_name'] = getattr(update_body.from_user, 'first_name', None)
|
|
changed = True
|
|
if meta.get('last_name') != getattr(update_body.from_user, 'last_name', None):
|
|
meta['last_name'] = getattr(update_body.from_user, 'last_name', None)
|
|
changed = True
|
|
user.meta = meta
|
|
user.last_activity = datetime.utcnow()
|
|
if changed:
|
|
session.commit()
|
|
|
|
data['user'] = user
|
|
# Pass sync session for routers expecting .query()
|
|
data['db_session'] = session
|
|
# chat_wrap can work with sync sessions too
|
|
data['chat_wrap'] = Wrapped_CBotChat(data['bot'], chat_id=user_id, db_session=session, user=user)
|
|
|
|
# De-duplicate known messages
|
|
if getattr(update_body, 'text', None):
|
|
existed = session.query(DbKnownMsg).filter(
|
|
(DbKnownMsg.chat_id == update_body.chat.id) &
|
|
(DbKnownMsg.message_id == update_body.message_id) &
|
|
(DbKnownMsg.from_user == True)
|
|
).first()
|
|
if existed:
|
|
await app_logger.adebug("Message already processed", message_id=update_body.message_id)
|
|
return
|
|
|
|
new_message = DbKnownMsg(
|
|
type='start_command' if str(update_body.text).startswith('/start') else 'common',
|
|
bot_id=data['chat_wrap'].bot_id,
|
|
chat_id=update_body.chat.id,
|
|
message_id=update_body.message_id,
|
|
from_user=True,
|
|
from_telegram_id=user_id,
|
|
created=datetime.utcnow(),
|
|
meta={}
|
|
)
|
|
session.add(new_message)
|
|
session.commit()
|
|
|
|
result = await handler(event, data)
|
|
return result
|
|
finally:
|
|
try:
|
|
session.close()
|
|
except Exception:
|
|
pass
|