314 lines
9.2 KiB
Python
314 lines
9.2 KiB
Python
"""
|
|
Session Management for Telegram Bot
|
|
|
|
This module handles session state for Telegram users, specifically managing
|
|
the active company selection for command handlers.
|
|
"""
|
|
|
|
import logging
|
|
import json
|
|
from typing import Dict, Any, Optional
|
|
from datetime import datetime
|
|
|
|
from backend.modules.telegram.db.operations import (
|
|
create_session,
|
|
get_user_active_session,
|
|
update_session_state,
|
|
delete_user_sessions
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class ConversationSession:
|
|
"""
|
|
Manages session state for a single user.
|
|
|
|
Attributes:
|
|
telegram_user_id: Telegram user ID
|
|
session_id: UUID of the session
|
|
active_company_id: Selected company ID
|
|
active_company_name: Selected company name
|
|
active_company_cui: Selected company CUI
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
telegram_user_id: int,
|
|
session_id: Optional[str] = None
|
|
):
|
|
"""
|
|
Initialize a session.
|
|
|
|
Args:
|
|
telegram_user_id: Telegram user ID
|
|
session_id: Existing session ID (if resuming), or None for new session
|
|
"""
|
|
self.telegram_user_id = telegram_user_id
|
|
self.session_id = session_id
|
|
self.created_at = datetime.now()
|
|
self.updated_at = datetime.now()
|
|
|
|
# Active company for this session
|
|
self.active_company_id: Optional[int] = None
|
|
self.active_company_name: Optional[str] = None
|
|
self.active_company_cui: Optional[str] = None
|
|
|
|
def set_active_company(
|
|
self,
|
|
company_id: int,
|
|
company_name: str,
|
|
company_cui: Optional[str] = None
|
|
):
|
|
"""
|
|
Set the active company for this session.
|
|
|
|
Args:
|
|
company_id: Company ID
|
|
company_name: Company name
|
|
company_cui: Company CUI (optional)
|
|
"""
|
|
self.active_company_id = company_id
|
|
self.active_company_name = company_name
|
|
self.active_company_cui = company_cui
|
|
self.updated_at = datetime.now()
|
|
logger.info(
|
|
f"Active company set for user {self.telegram_user_id}: "
|
|
f"{company_name} (ID: {company_id})"
|
|
)
|
|
|
|
def get_active_company(self) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get the active company information.
|
|
|
|
Returns:
|
|
Dict with company info (id, name, cui) or None if no company selected
|
|
"""
|
|
if self.active_company_id is not None:
|
|
return {
|
|
"id": self.active_company_id,
|
|
"name": self.active_company_name,
|
|
"cui": self.active_company_cui
|
|
}
|
|
return None
|
|
|
|
def clear_active_company(self):
|
|
"""
|
|
Clear the active company selection.
|
|
"""
|
|
logger.info(
|
|
f"Clearing active company for user {self.telegram_user_id} "
|
|
f"(was: {self.active_company_name})"
|
|
)
|
|
self.active_company_id = None
|
|
self.active_company_name = None
|
|
self.active_company_cui = None
|
|
self.updated_at = datetime.now()
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""
|
|
Serialize session to dictionary (for database storage).
|
|
|
|
Returns:
|
|
Dict representation of session
|
|
"""
|
|
return {
|
|
"telegram_user_id": self.telegram_user_id,
|
|
"session_id": self.session_id,
|
|
"created_at": self.created_at.isoformat(),
|
|
"updated_at": self.updated_at.isoformat(),
|
|
"active_company_id": self.active_company_id,
|
|
"active_company_name": self.active_company_name,
|
|
"active_company_cui": self.active_company_cui
|
|
}
|
|
|
|
@classmethod
|
|
def from_dict(cls, data: Dict[str, Any]) -> 'ConversationSession':
|
|
"""
|
|
Deserialize session from dictionary.
|
|
|
|
Args:
|
|
data: Dict representation of session
|
|
|
|
Returns:
|
|
ConversationSession instance
|
|
"""
|
|
session = cls(
|
|
telegram_user_id=data["telegram_user_id"],
|
|
session_id=data.get("session_id")
|
|
)
|
|
|
|
# Restore active company
|
|
session.active_company_id = data.get("active_company_id")
|
|
session.active_company_name = data.get("active_company_name")
|
|
session.active_company_cui = data.get("active_company_cui")
|
|
|
|
if "created_at" in data:
|
|
session.created_at = datetime.fromisoformat(data["created_at"])
|
|
if "updated_at" in data:
|
|
session.updated_at = datetime.fromisoformat(data["updated_at"])
|
|
|
|
return session
|
|
|
|
|
|
class SessionManager:
|
|
"""
|
|
Manages sessions for all users.
|
|
|
|
Provides methods to create, retrieve, update, and delete sessions.
|
|
Sessions are stored both in memory (for quick access) and in database
|
|
(for persistence).
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
Initialize the session manager.
|
|
"""
|
|
self._sessions: Dict[int, ConversationSession] = {}
|
|
logger.info("SessionManager initialized")
|
|
|
|
async def get_or_create_session(
|
|
self,
|
|
telegram_user_id: int
|
|
) -> ConversationSession:
|
|
"""
|
|
Get existing session for a user or create a new one.
|
|
|
|
Args:
|
|
telegram_user_id: Telegram user ID
|
|
|
|
Returns:
|
|
ConversationSession for the user
|
|
"""
|
|
# Check in-memory cache first
|
|
if telegram_user_id in self._sessions:
|
|
logger.debug(f"Found session in cache for user {telegram_user_id}")
|
|
return self._sessions[telegram_user_id]
|
|
|
|
# Check database for existing session
|
|
session_data = await get_user_active_session(telegram_user_id)
|
|
|
|
if session_data:
|
|
# Restore session from database
|
|
conversation_state_json = session_data.get('conversation_state')
|
|
|
|
if conversation_state_json:
|
|
try:
|
|
session_dict = json.loads(conversation_state_json)
|
|
session = ConversationSession.from_dict(session_dict)
|
|
session.session_id = session_data['session_id']
|
|
self._sessions[telegram_user_id] = session
|
|
logger.info(f"Restored session from database for user {telegram_user_id}")
|
|
return session
|
|
except json.JSONDecodeError as e:
|
|
logger.error(f"Failed to parse session state: {e}")
|
|
|
|
# Create new session
|
|
session = ConversationSession(telegram_user_id)
|
|
|
|
# Save to database
|
|
session_id = await create_session(
|
|
telegram_user_id=telegram_user_id,
|
|
conversation_state=json.dumps(session.to_dict()),
|
|
expires_in_hours=24
|
|
)
|
|
|
|
session.session_id = session_id
|
|
self._sessions[telegram_user_id] = session
|
|
|
|
logger.info(f"Created new session for user {telegram_user_id} (ID: {session_id})")
|
|
return session
|
|
|
|
async def save_session(self, telegram_user_id: int) -> bool:
|
|
"""
|
|
Save session to database.
|
|
|
|
Args:
|
|
telegram_user_id: Telegram user ID
|
|
|
|
Returns:
|
|
bool: True if saved successfully
|
|
"""
|
|
session = self._sessions.get(telegram_user_id)
|
|
|
|
if not session or not session.session_id:
|
|
logger.warning(f"No session to save for user {telegram_user_id}")
|
|
return False
|
|
|
|
try:
|
|
conversation_state = json.dumps(session.to_dict())
|
|
|
|
success = await update_session_state(
|
|
session_id=session.session_id,
|
|
conversation_state=conversation_state
|
|
)
|
|
|
|
if success:
|
|
logger.debug(f"Saved session for user {telegram_user_id}")
|
|
else:
|
|
logger.warning(f"Failed to save session for user {telegram_user_id}")
|
|
|
|
return success
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error saving session for user {telegram_user_id}: {e}")
|
|
return False
|
|
|
|
async def delete_session(self, telegram_user_id: int) -> bool:
|
|
"""
|
|
Delete session completely (from memory and database).
|
|
|
|
Args:
|
|
telegram_user_id: Telegram user ID
|
|
|
|
Returns:
|
|
bool: True if deleted successfully
|
|
"""
|
|
# Remove from memory
|
|
if telegram_user_id in self._sessions:
|
|
del self._sessions[telegram_user_id]
|
|
|
|
# Delete from database
|
|
success = await delete_user_sessions(telegram_user_id)
|
|
|
|
if success:
|
|
logger.info(f"Deleted session for user {telegram_user_id}")
|
|
else:
|
|
logger.warning(f"Failed to delete session for user {telegram_user_id}")
|
|
|
|
return success
|
|
|
|
def get_active_sessions_count(self) -> int:
|
|
"""
|
|
Get count of active sessions in memory.
|
|
|
|
Returns:
|
|
int: Number of active sessions
|
|
"""
|
|
return len(self._sessions)
|
|
|
|
|
|
# Singleton instance
|
|
_session_manager_instance: Optional[SessionManager] = None
|
|
|
|
|
|
def get_session_manager() -> SessionManager:
|
|
"""
|
|
Get or create the singleton SessionManager instance.
|
|
|
|
Returns:
|
|
SessionManager: Singleton instance
|
|
"""
|
|
global _session_manager_instance
|
|
if _session_manager_instance is None:
|
|
_session_manager_instance = SessionManager()
|
|
return _session_manager_instance
|
|
|
|
|
|
# Export main classes and functions
|
|
__all__ = [
|
|
'ConversationSession',
|
|
'SessionManager',
|
|
'get_session_manager'
|
|
]
|