feat: add log queue

Author: An0nX

userbot/__main__.py

@@ -2,7 +2,7 @@
 import gc
 import logging
 import sys
-from typing import Dict, Any
+from typing import Dict, Any, List
 
 from apscheduler.schedulers.asyncio import AsyncIOScheduler
 from rich.console import Console
@@ -15,13 +15,40 @@
     ACTIVE_CLIENTS, FAKE, TelegramClient,
     db_setup, manage_clients, GLOBAL_HELP_INFO
 )
-from userbot.src.config import GC_INTERVAL_SECONDS
+from userbot.src.config import GC_INTERVAL_SECONDS, LOG_QUEUE_INTERVAL_SECONDS, LOG_QUEUE_BATCH_SIZE
 from userbot.src.db.session import get_db
 import userbot.src.db_manager as db_manager
+from userbot.src.log_handler import log_queue
 
 console: Console = Console()
 logger: logging.Logger = logging.getLogger(__name__)
 
+# --- Log Processing Worker ---
+async def process_log_queue():
+    """
+    Periodically processes logs from the queue and writes them to the DB.
+    """
+    logs_to_process: List[Dict[str, Any]] = []
+    
+    # Drain the queue up to the batch size
+    while len(logs_to_process) < LOG_QUEUE_BATCH_SIZE and not log_queue.empty():
+        try:
+            log_item = log_queue.get_nowait()
+            logs_to_process.append(log_item)
+            log_queue.task_done()
+        except asyncio.QueueEmpty:
+            break
+            
+    if logs_to_process:
+        try:
+            async with get_db() as db_session:
+                await db_manager.add_logs_bulk(db_session, logs_to_process)
+        except Exception as e:
+            # If the DB fails, the logs are lost, but we print an error.
+            # The original handler already has a fallback for this scenario.
+            logger.error(f"Failed to process log queue batch: {e}", exc_info=True)
+
+
 # --- Helper Functions ---
 async def get_account_id_from_client(client) -> int | None:
     return next((acc_id for acc_id, c in ACTIVE_CLIENTS.items() if c == client), None)
@@ -184,11 +211,13 @@ async def main():
 
     scheduler = AsyncIOScheduler(timezone="UTC")
     scheduler.add_job(gc.collect, 'interval', seconds=GC_INTERVAL_SECONDS, id='gc_job')
+    scheduler.add_job(process_log_queue, 'interval', seconds=LOG_QUEUE_INTERVAL_SECONDS, id='log_queue_job')
     scheduler.start()
+    
     console.print(f"-> [system] - GC scheduled every {GC_INTERVAL_SECONDS} seconds.", style="blue")
+    console.print(f"-> [system] - Log queue processing scheduled every {LOG_QUEUE_INTERVAL_SECONDS} seconds.", style="blue")
 
     logger.info("Userbot is running. Press Ctrl+C to stop.")
-    # Keep the main coroutine alive indefinitely
     await asyncio.Event().wait()
 
 

userbot/src/config.py

@@ -16,14 +16,17 @@
 DB_USER: str = os.getenv("DB_USER", "userbot")
 DB_PASS: str = os.getenv("DB_PASS", "userbot_password")
 DB_NAME: str = os.getenv("DB_NAME", "userbot_db")
-# New: Connection retry settings
 DB_CONN_RETRIES: int = int(os.getenv("DB_CONN_RETRIES", 5))
 DB_CONN_RETRY_DELAY: int = int(os.getenv("DB_CONN_RETRY_DELAY", 5))
 
 
 # --- Application Settings ---
 MODULE_FOLDER: str = "userbot.modules"
 LOG_LEVEL: str = os.getenv("LOG_LEVEL", "INFO").upper()
+# New: Logging queue settings
+LOG_QUEUE_INTERVAL_SECONDS: int = int(os.getenv("LOG_QUEUE_INTERVAL_SECONDS", 5))
+LOG_QUEUE_BATCH_SIZE: int = int(os.getenv("LOG_QUEUE_BATCH_SIZE", 50))
+
 
 # --- Scheduler Settings ---
 GC_INTERVAL_SECONDS: int = int(os.getenv("GC_INTERVAL_SECONDS", 60))

userbot/src/db_manager.py

@@ -25,7 +25,6 @@ async def add_account(
     app_version: str,
     user_telegram_id: Optional[int] = None
 ) -> Optional[Account]:
-    """Adds a new account to the database with encrypted credentials."""
     try:
         new_account = Account(
             account_name=account_name,
@@ -43,7 +42,7 @@ async def add_account(
         logger.info(f"Added account '{account_name}' with ID: {new_account.account_id}")
         return new_account
     except IntegrityError:
-        logger.warning(f"Account with name '{account_name}' already exists.")
+        logger.warning(f"Account with name '{account_name}' or user_id '{user_telegram_id}' already exists.")
         await db.rollback()
         return None
     except Exception as e:
@@ -52,27 +51,22 @@ async def add_account(
         return None
 
 async def get_account(db: AsyncSession, account_name: str) -> Optional[Account]:
-    """Retrieves an account by its name."""
     result = await db.execute(select(Account).where(Account.account_name == account_name))
     return result.scalars().first()
 
 async def get_account_by_user_id(db: AsyncSession, user_id: int) -> Optional[Account]:
-    """Retrieves an account by its Telegram User ID."""
     result = await db.execute(select(Account).where(Account.user_telegram_id == user_id))
     return result.scalars().first()
 
 async def get_all_accounts(db: AsyncSession) -> List[Account]:
-    """Retrieves all accounts from the database, with their sessions for status display."""
     result = await db.execute(select(Account).options(selectinload(Account.session)).order_by(Account.account_id))
     return result.scalars().all()
 
 async def get_all_active_accounts(db: AsyncSession) -> List[Account]:
-    """Retrieves all enabled accounts from the database."""
     result = await db.execute(select(Account).where(Account.is_enabled == True))
     return result.scalars().all()
 
 async def delete_account(db: AsyncSession, account_name: str) -> bool:
-    """Deletes an account by its name."""
     account = await get_account(db, account_name)
     if not account:
         return False
@@ -81,7 +75,6 @@ async def delete_account(db: AsyncSession, account_name: str) -> bool:
     return True
 
 async def toggle_account_status(db: AsyncSession, account_name: str) -> Optional[bool]:
-    """Toggles the is_enabled status of an account."""
     account = await get_account(db, account_name)
     if not account:
         return None
@@ -90,14 +83,12 @@ async def toggle_account_status(db: AsyncSession, account_name: str) -> Optional
     return account.is_enabled
 
 async def update_account_lang(db: AsyncSession, account_id: int, lang_code: str) -> bool:
-    """Updates the language for a specific account."""
     stmt = update(Account).where(Account.account_id == account_id).values(lang_code=lang_code)
     result = await db.execute(stmt)
     return result.rowcount > 0
 
 # --- Session CRUD ---
 async def get_session(db: AsyncSession, account_id: int) -> Optional[Session]:
-    """Retrieves a session for a given account and decrypts its auth key."""
     result = await db.execute(select(Session).where(Session.account_id == account_id))
     session = result.scalars().first()
     if session and session.auth_key_data:
@@ -109,7 +100,6 @@ async def get_session(db: AsyncSession, account_id: int) -> Optional[Session]:
     return session
 
 async def add_or_update_session(db: AsyncSession, **kwargs) -> Optional[Session]:
-    """Adds or updates a session in the database, encrypting the auth key."""
     account_id = kwargs.get("account_id")
     if not account_id: return None
 
@@ -129,37 +119,17 @@ async def add_or_update_session(db: AsyncSession, **kwargs) -> Optional[Session]
     await db.flush()
     return session
 
-async def delete_session(db: AsyncSession, account_id: int) -> bool:
-    """Deletes a session from the database."""
-    stmt = delete(Session).where(Session.account_id == account_id)
-    result = await db.execute(stmt)
-    await db.flush()
-    return result.rowcount > 0
-
 # --- Log CRUD ---
-async def add_log(db: AsyncSession, level: str, message: str, account_id: Optional[int] = None, module_name: Optional[str] = None) -> Optional[Log]:
-    """Adds a log entry to the database."""
+async def add_logs_bulk(db: AsyncSession, logs: List[Dict[str, Any]]) -> None:
+    """Adds a batch of log entries to the database."""
+    if not logs:
+        return
     try:
-        new_log = Log(
-            level=level.upper(),
-            message=message,
-            account_id=account_id,
-            module_name=module_name
-        )
-        db.add(new_log)
+        # SQLAlchemy 2.0 style bulk insert
+        db.add_all([Log(**log_data) for log_data in logs])
         await db.flush()
-        return new_log
     except Exception as e:
-        print(f"CRITICAL: Error adding log to database: {e}")
-        return None
-
-async def get_logs(db: AsyncSession, account_id: Optional[int] = None, limit: int = 100, level: Optional[str] = None) -> List[Log]:
-    """Retrieves logs from the database with optional filters."""
-    stmt = select(Log).order_by(Log.timestamp.desc()).limit(limit)
-    if account_id is not None:
-        stmt = stmt.where(Log.account_id == account_id)
-    if level is not None:
-        stmt = stmt.where(Log.level == level.upper())
-    
-    result = await db.execute(stmt)
-    return result.scalars().all()
+        # Fallback to printing if bulk DB logging fails
+        print(f"CRITICAL: Error during bulk log insert: {e}")
+        for log in logs:
+            print(f"Fallback Log: [{log.get('level')}] {log.get('message')}")

userbot/src/log_handler.py

@@ -1,12 +1,14 @@
 import logging
 import asyncio
-from typing import Optional
+from typing import Optional, Dict, Any
 
-from userbot.src.db_manager import add_log
+# A simple, thread-safe in-memory queue for log records.
+# asyncio.Queue is used because the consumer (writer) is async.
+log_queue: asyncio.Queue = asyncio.Queue()
 
 class DBLogHandler(logging.Handler):
     """
-    A custom logging handler that writes log records to the database.
+    A custom logging handler that puts log records into an async queue.
     """
     def __init__(self, level: int = logging.NOTSET) -> None:
         """
@@ -19,34 +21,26 @@ def __init__(self, level: int = logging.NOTSET) -> None:
 
     def emit(self, record: logging.LogRecord) -> None:
         """
-        Writes a log record to the database.
+        Puts a log record into the in-memory queue.
 
-        This method is called by the logging framework for each log message.
-        It schedules the database write operation on the running event loop.
+        This operation is non-blocking and safe to call from synchronous code.
 
         Args:
             record (logging.LogRecord): The log record to be processed.
         """
-        # We need to get the running event loop to schedule our async db call
         try:
-            loop: asyncio.AbstractEventLoop = asyncio.get_running_loop()
-            # Default values for account_id and module_name
-            account_id: Optional[int] = getattr(record, 'account_id', None)
-            module_name: Optional[str] = getattr(record, 'module_name', None)
-
-            # Fire and forget: schedule the coroutine to run on the loop
-            asyncio.ensure_future(add_log(
-                level=record.levelname,
-                message=self.format(record),
-                account_id=account_id,
-                module_name=module_name
-            ), loop=loop)
-        except RuntimeError:
-            # This can happen if logging occurs when no event loop is running.
-            # In this case, we can't log to the DB.
-            # We can print to stderr as a fallback for critical logs.
-            if record.levelno >= logging.ERROR:
-                print(f"CRITICAL [DBLogHandler-Fallback]: Could not log to DB (no running loop): {self.format(record)}")
+            log_entry: Dict[str, Any] = {
+                'level': record.levelname,
+                'message': self.format(record),
+                'account_id': getattr(record, 'account_id', None),
+                'module_name': getattr(record, 'module_name', None)
+            }
+            # put_nowait is a non-blocking call, safe for sync contexts.
+            log_queue.put_nowait(log_entry)
+        except asyncio.QueueFull:
+            # This should rarely happen with a default-sized queue.
+            # A fallback print is better than losing the log.
+            print(f"CRITICAL [DBLogHandler-Fallback]: Log queue is full. Dropping log: {self.format(record)}")
         except Exception:
-            # Fallback for any other error during logging to prevent crashes.
+            # Broad exception to prevent any logging-related crash.
             print(f"CRITICAL [DBLogHandler-Fallback]: An unexpected error occurred in emit: {self.format(record)}")