fix: some more fixes with cache

Author: An0nX

userbot/__init__.py

@@ -5,7 +5,6 @@
 
 from faker import Faker
 from telethon import TelegramClient as TelethonTelegramClient, events
-from telethon.sessions import StringSession
 from telethon.errors.rpcerrorlist import UserAlreadyParticipantError
 from python_socks import ProxyType
 
@@ -18,31 +17,22 @@
 from userbot.src.log_handler import DBLogHandler
 from userbot.src.locales import translator
 
-# --- Basic Setup ---
 logger: logging.Logger = logging.getLogger("userbot")
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
 
-# --- Globals ---
 ACTIVE_CLIENTS: Dict[int, "TelegramClient"] = {}
 FAKE: Faker = Faker()
 GLOBAL_HELP_INFO: Dict[int, Dict[str, str]] = {}
 
-# --- Helper Functions ---
 def _generate_random_device() -> Dict[str, str]:
-    """Generates a dictionary with random device information."""
     return {
         "device_model": FAKE.user_agent(),
         "system_version": f"SDK {FAKE.random_int(min=28, max=33)}",
         "app_version": f"{FAKE.random_int(min=9, max=10)}.{FAKE.random_int(min=0, max=9)}.{FAKE.random_int(min=0, max=9)}"
     }
 
-# --- Core TelegramClient Class ---
 class TelegramClient(TelethonTelegramClient):
-    """
-    Custom TelegramClient class with database interaction and localization methods.
-    """
     def __init__(self, *args, **kwargs):
-        """Initializes the custom Telegram client."""
         super().__init__(*args, **kwargs)
         self.lang_code: str = 'ru'
 
@@ -55,7 +45,6 @@ def current_account_id(self) -> Optional[int]:
     async def get_string(self, key: str, module_name: Optional[str] = None, **kwargs) -> str:
         return translator.get_string(self.lang_code, key, module_name, **kwargs)
 
-# --- Startup Logic ---
 async def db_setup() -> None:
     if not API_ID or not API_HASH:
         logger.critical("API_ID or API_HASH is not set. Please run 'python3 -m scripts.setup'. Exiting.")
@@ -146,16 +135,16 @@ async def manage_clients() -> None:
              proxy_type_enum = proxy_map.get(account.proxy_type.lower())
              if proxy_type_enum:
                  proxy_details = (
-                     proxy_type_enum,
-                     account.proxy_ip,
-                     account.proxy_port,
-                     True,
+                     proxy_type_enum, account.proxy_ip, account.proxy_port, True,
                      encryption_manager.decrypt(account.proxy_username).decode() if account.proxy_username else None,
                      encryption_manager.decrypt(account.proxy_password).decode() if account.proxy_password else None
                  )
+        
+        # Use the asynchronous factory method to create the session
+        session = await DbSession.create(account_id=account.account_id)
 
         new_client: TelegramClient = TelegramClient(
-            session=DbSession(account_id=account.account_id), 
+            session=session, 
             api_id=int(acc_api_id), 
             api_hash=acc_api_hash, 
             device_model=account.device_model, 

userbot/src/db_session.py

@@ -4,7 +4,7 @@
 
 from telethon.sessions.abstract import Session
 from telethon.crypto import AuthKey
-from telethon.tl.types import InputPhoto, InputDocument, PeerUser, PeerChat, PeerChannel
+from telethon.tl.types import PeerUser
 
 import userbot.src.db_manager as db_manager
 from userbot.src.db.session import get_db
@@ -14,22 +14,23 @@
 class DbSession(Session):
     """
     A Telethon session class that stores session data in a PostgreSQL database.
-    This implementation correctly provides all the abstract methods and properties
-    required by Telethon's base Session class.
     """
 
-    def __init__(self, account_id: int):
+    def __init__(self, account_id: int, self_user_id: Optional[int]):
         """
         Initializes the database-backed session.
 
         Args:
             account_id (int): The unique identifier for the account this session belongs to.
+            self_user_id (Optional[int]): The Telegram user ID of the account holder.
         """
         super().__init__()
         if not isinstance(account_id, int):
             raise ValueError("DbSession requires a valid integer account_id.")
 
         self.account_id: int = account_id
+        self._self_user_id: Optional[int] = self_user_id
+        
         self._auth_key: Optional[AuthKey] = None
         self._dc_id: int = 0
         self._server_address: Optional[str] = None
@@ -40,11 +41,29 @@ def __init__(self, account_id: int):
         self._qts: Optional[int] = None
         self._date: Optional[int] = None
         self._seq: Optional[int] = None
-        
-    async def load(self) -> None:
+
+    @classmethod
+    async def create(cls, account_id: int) -> "DbSession":
         """
-        Loads the session data for the current account_id from the database.
+        Asynchronously creates and pre-loads a DbSession instance.
+        This factory method is used to fetch necessary data like the user_telegram_id
+        before the synchronous parts of the session are accessed by Telethon.
+
+        Args:
+            account_id (int): The unique identifier for the account.
+
+        Returns:
+            DbSession: A new instance of DbSession.
         """
+        self_user_id: Optional[int] = None
+        async with get_db() as db:
+            account = await db_manager.get_account_by_id(db, account_id)
+            if account:
+                self_user_id = account.user_telegram_id
+        return cls(account_id, self_user_id)
+        
+    async def load(self) -> None:
+        """Loads the session data for the current account_id from the database."""
         logger.debug(f"Attempting to load session for account_id: {self.account_id}")
         async with get_db() as db:
             session_data = await db_manager.get_session(db, self.account_id)
@@ -54,157 +73,81 @@ async def load(self) -> None:
             self._dc_id = session_data.dc_id
             self._server_address = session_data.server_address
             self._port = session_data.port
-            
-            auth_key_bytes = session_data.auth_key_data
-            if auth_key_bytes:
-                self._auth_key = AuthKey(data=bytes(auth_key_bytes))
-            else:
-                self._auth_key = None
-            
+            self._auth_key = AuthKey(data=bytes(session_data.auth_key_data)) if session_data.auth_key_data else None
             self._takeout_id = session_data.takeout_id
             self._pts = session_data.pts
             self._qts = session_data.qts
             self._date = session_data.date
             self._seq = session_data.seq
         else:
-            logger.info(f"No session data found in DB for account_id: {self.account_id}. New login required.")
+            logger.info(f"No session data in DB for account_id: {self.account_id}. New login required.")
 
     async def save(self) -> None:
-        """
-        Saves the current session data to the database.
-        This is called automatically by Telethon.
-        """
-        logger.debug(f"Attempting to save session for account_id: {self.account_id}")
+        """Saves the current session data to the database."""
         session_data = {
-            "account_id": self.account_id,
-            "dc_id": self._dc_id,
-            "server_address": self._server_address,
-            "port": self._port,
+            "account_id": self.account_id, "dc_id": self._dc_id,
+            "server_address": self._server_address, "port": self._port,
             "auth_key_data": self._auth_key.key if self._auth_key else None,
-            "pts": self._pts,
-            "qts": self._qts,
-            "date": self._date,
-            "seq": self._seq,
-            "takeout_id": self._takeout_id,
+            "pts": self._pts, "qts": self._qts, "date": self._date,
+            "seq": self._seq, "takeout_id": self._takeout_id,
         }
         async with get_db() as db:
             await db_manager.add_or_update_session(db, **session_data)
         logger.info(f"Session saved for account_id: {self.account_id}")
 
     async def delete(self) -> None:
         """Deletes the session for the current account from the database."""
-        logger.info(f"Deleting session from DB for account_id: {self.account_id}")
         async with get_db() as db:
             await db_manager.delete_session(db, self.account_id)
-        # Clear in-memory data
-        self._auth_key = None
-        self._dc_id = 0
-    
-    # --- Abstract Properties Implementation ---
+        self._auth_key = None; self._dc_id = 0
 
     @property
-    def dc_id(self) -> int:
-        return self._dc_id
-
+    def dc_id(self) -> int: return self._dc_id
     @property
-    def server_address(self) -> Optional[str]:
-        return self._server_address
-
+    def server_address(self) -> Optional[str]: return self._server_address
     @property
-    def port(self) -> int:
-        return self._port
-
+    def port(self) -> int: return self._port
     @property
-    def auth_key(self) -> Optional[AuthKey]:
-        return self._auth_key
-
+    def auth_key(self) -> Optional[AuthKey]: return self._auth_key
     @auth_key.setter
-    def auth_key(self, value: Optional[AuthKey]):
-        self._auth_key = value
-
+    def auth_key(self, value: Optional[AuthKey]): self._auth_key = value
     @property
-    def takeout_id(self) -> Optional[int]:
-        return self._takeout_id
-
+    def takeout_id(self) -> Optional[int]: return self._takeout_id
     @takeout_id.setter
-    def takeout_id(self, value: Optional[int]):
-        self._takeout_id = value
-
-    # --- Abstract Methods Implementation ---
+    def takeout_id(self, value: Optional[int]): self._takeout_id = value
 
     def set_dc(self, dc_id: int, server_address: str, port: int):
-        self._dc_id = dc_id
-        self._server_address = server_address
-        self._port = port
+        self._dc_id, self._server_address, self._port = dc_id, server_address, port
 
     def get_update_state(self, entity_id: int) -> Optional[Tuple[int, int, int, int, int]]:
-        if self._pts is None:
-            return None
+        if self._pts is None: return None
         return self._pts, self._qts, self._date, self._seq, 0
 
     def set_update_state(self, entity_id: int, state: Any):
         if isinstance(state.date, datetime):
             date_ts = int(state.date.replace(tzinfo=timezone.utc).timestamp())
         else:
             date_ts = int(state.date)
-
-        self._pts = state.pts
-        self._qts = state.qts
-        self._date = date_ts
-        self._seq = state.seq
+        self._pts, self._qts, self._date, self._seq = state.pts, state.qts, date_ts, state.seq
 
-    async def close(self) -> None:
-        pass
+    async def close(self) -> None: pass
 
     def get_update_states(self) -> List[Tuple[int, int, int, int, int, int]]:
-        if self._pts is None:
-            return []
+        if self._pts is None: return []
         return [(0, self._pts, self._qts, self._date, self._seq, 0)]
 
-    # --- New Stubs for Entity and File Caching ---
-
-    def process_entities(self, tlo: object) -> None:
-        """
-        This session does not cache entities, so this method does nothing.
-        
-        Args:
-            tlo (object): A TLObject containing entities.
-        """
-        pass
+    def process_entities(self, tlo: object) -> None: pass
 
     def get_input_entity(self, key: Any) -> Any:
         """
-        This session does not cache entities, so this method always fails.
-
-        Args:
-            key (Any): The key to look up an entity.
-
-        Raises:
-            KeyError: Always, as no entities are cached.
+        Returns the input entity for the current user if key is 0.
+        This is crucial for the client to know "who it is" upon startup.
         """
-        raise KeyError("Entity not found in DbSession cache (caching is not implemented).")
+        if key == 0 and self._self_user_id is not None:
+            return PeerUser(self._self_user_id)
+        raise KeyError("Entity not found in DbSession cache (caching is not implemented for other entities).")
 
-    def cache_file(self, md5_digest: bytes, file_size: int, instance: Any) -> None:
-        """
-        This session does not cache files, so this method does nothing.
-
-        Args:
-            md5_digest (bytes): The MD5 digest of the file.
-            file_size (int): The size of the file.
-            instance (Any): The InputFile or InputPhoto instance.
-        """
-        pass
+    def cache_file(self, md5_digest: bytes, file_size: int, instance: Any) -> None: pass
 
     def get_file(self, md5_digest: bytes, file_size: int, exact: bool = True) -> Optional[Any]:
-        """
-        This session does not cache files, so this method always returns None.
-
-        Args:
-            md5_digest (bytes): The MD5 digest of the file.
-            file_size (int): The size of the file.
-            exact (bool): Whether the file size must be exact.
-        
-        Returns:
-            None: Always, as no files are cached.
-        """
         return None