from __future__ import annotations import uuid as uuid_lib from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.exc import IntegrityError from yuxi.channels.models import ChannelMessage from yuxi.storage.postgres.models_channels import ChannelUserMapping, ChannelThreadMapping from yuxi.storage.postgres.models_business import User from yuxi.utils.datetime_utils import utc_now_naive from yuxi.utils.logging_config import logger VIRTUAL_DEPARTMENT_ID = -1 USER_SOURCE_PREFIX = "channel:" class SessionMapper: def __init__(self, db: AsyncSession, department_id: int = VIRTUAL_DEPARTMENT_ID): self.db = db self.department_id = department_id async def resolve_user(self, message: ChannelMessage) -> str: identity = message.identity mapping = await self._get_user_mapping(identity.channel_id, identity.channel_user_id) if mapping: return mapping.internal_user_id internal_user_id = f"ch_{identity.channel_id}_{uuid_lib.uuid4().hex[:8]}" source = f"{USER_SOURCE_PREFIX}{identity.channel_id}" try: user = User( username=f"{identity.channel_id}_{identity.channel_user_id}", user_id=internal_user_id, password_hash="", role="user", department_id=self.department_id, source=source, ) self.db.add(user) await self.db.flush() mapping = ChannelUserMapping( channel_id=identity.channel_id, channel_user_id=identity.channel_user_id, internal_user_id=internal_user_id, ) self.db.add(mapping) await self.db.commit() logger.info( f"Created channel user mapping: {identity.channel_id}/{identity.channel_user_id} -> {internal_user_id}" ) return internal_user_id except IntegrityError: await self.db.rollback() mapping = await self._get_user_mapping(identity.channel_id, identity.channel_user_id) if mapping: return mapping.internal_user_id raise async def resolve_thread(self, message: ChannelMessage, internal_user_id: str) -> str: identity = message.identity mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id) if mapping: mapping.last_active_at = utc_now_naive() await self.db.flush() return mapping.thread_id thread_id = str(uuid_lib.uuid4()) try: mapping = ChannelThreadMapping( channel_id=identity.channel_id, channel_chat_id=identity.channel_chat_id, internal_user_id=internal_user_id, thread_id=thread_id, last_active_at=utc_now_naive(), ) self.db.add(mapping) await self.db.commit() logger.info( f"Created channel thread mapping: " f"{identity.channel_id}/{identity.channel_chat_id}/{internal_user_id} -> {thread_id}" ) return thread_id except IntegrityError: await self.db.rollback() mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id) if mapping: return mapping.thread_id raise async def reset_thread(self, message: ChannelMessage, internal_user_id: str) -> str: identity = message.identity mapping = await self._get_thread_mapping(identity.channel_id, identity.channel_chat_id, internal_user_id) if mapping: new_thread_id = str(uuid_lib.uuid4()) mapping.thread_id = new_thread_id mapping.last_active_at = utc_now_naive() await self.db.commit() logger.info(f"Reset thread: {identity.channel_id}/{identity.channel_chat_id} -> {new_thread_id}") return new_thread_id return await self.resolve_thread(message, internal_user_id) async def _get_user_mapping(self, channel_id: str, channel_user_id: str) -> ChannelUserMapping | None: result = await self.db.execute( select(ChannelUserMapping).where( ChannelUserMapping.channel_id == channel_id, ChannelUserMapping.channel_user_id == channel_user_id, ) ) return result.scalar_one_or_none() async def _get_thread_mapping( self, channel_id: str, channel_chat_id: str, internal_user_id: str ) -> ChannelThreadMapping | None: result = await self.db.execute( select(ChannelThreadMapping).where( ChannelThreadMapping.channel_id == channel_id, ChannelThreadMapping.channel_chat_id == channel_chat_id, ChannelThreadMapping.internal_user_id == internal_user_id, ) ) return result.scalar_one_or_none()