98 lines
3.9 KiB
Python
98 lines
3.9 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 app.core.storage import db_session
|
|
from aiogram import BaseMiddleware, types
|
|
from app.core.models.messages import KnownTelegramMessage
|
|
from datetime import datetime
|
|
|
|
|
|
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, db_session=session, user=user)
|
|
data['memory'] = data['dispatcher']._s_memory
|
|
|
|
if isinstance(update_body, types.Message):
|
|
message_type = 'common'
|
|
if update_body.text.startswith('/start'):
|
|
message_type = 'start_command'
|
|
|
|
if session.query(KnownTelegramMessage).filter_by(
|
|
chat_id=update_body.chat.id,
|
|
message_id=update_body.message_id,
|
|
from_user=True
|
|
).first():
|
|
make_log("UserDataMiddleware", f"Message {update_body.message_id} already processed", level='debug')
|
|
return
|
|
|
|
new_message = KnownTelegramMessage(
|
|
type=message_type,
|
|
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.now(),
|
|
meta={}
|
|
)
|
|
session.add(new_message)
|
|
session.commit()
|
|
|
|
result = await handler(event, data)
|
|
return result
|