uploader-bot/app/client_bot/__init__.py

81 lines
2.9 KiB
Python

# @MY_UploaderRobot
from datetime import datetime
from aiogram import BaseMiddleware, Dispatcher
from aiogram.fsm.storage.memory import MemoryStorage
from app.client_bot.routers.index import main_router
from app.core.logger import logger
from app.core.models._telegram import Wrapped_CBotChat
from app.core.models.user import User
from app.core.storage import db_session
dp = Dispatcher(storage=MemoryStorage())
class UserDataMiddleware(BaseMiddleware):
async def __call__(self, handler, event, data):
update_body = event.message or event.callback_query
if not update_body:
return
if update_body.from_user.is_bot is True:
return
user_id = update_body.from_user.id
assert user_id >= 1
# TODO: maybe make users cache
with db_session(auto_commit=False) as session:
try:
user = session.query(User).filter_by(telegram_id=user_id).first()
except BaseException as e:
logger.error(f"Error when middleware getting user: {e}")
user = None
if user is None:
logger.debug(f"User {user_id} not found. Creating new user")
user = User(
telegram_id=user_id,
username=update_body.from_user.username,
lang_code='en',
last_use=datetime.now(),
meta=dict(first_name=update_body.from_user.first_name or '',
last_name=update_body.from_user.last_name or '', username=update_body.from_user.username,
language_code=update_body.from_user.language_code,
is_premium=update_body.from_user.is_premium),
created=datetime.now()
)
session.add(user)
session.commit()
else:
if user.username != update_body.from_user.username:
user.username = update_body.from_user.username
updated_meta_fields = {}
if user.meta.get('first_name') != update_body.from_user.first_name:
updated_meta_fields['first_name'] = update_body.from_user.first_name
if user.meta.get('last_name') != update_body.from_user.last_name:
updated_meta_fields['last_name'] = update_body.from_user.last_name
user.meta = {
**user.meta,
**updated_meta_fields
}
user.last_use = datetime.now()
session.commit()
data['user'] = user
data['db_session'] = session
data['chat_wrap'] = Wrapped_CBotChat(data['bot'], chat_id=user_id)
data['memory'] = dp._s_memory
result = await handler(event, data)
return result
dp.update.outer_middleware(UserDataMiddleware())
dp.include_router(main_router)