Skip to content

🔃 Callback functions

The BaseModule class provides several methods that act as lifecycle callbacks, allowing you to execute custom code at specific points during the module's existence. You can override these methods in your module class.

  • def on_init(self):
    • Called after Stage 2 initialization (self.stage2()) is complete, meaning all standard handlers have been registered and any ModuleExtensions have been loaded and initialized.
    • This is a good place for setting up internal state, loading data from non-database sources, or starting background tasks that don't depend directly on the database being ready.

    INFO

    Note:** Database availability (self.db) is not guaranteed at this point, as database setup happens asynchronously. If you need the DB, use on_db_ready.

  • async def on_db_ready(self):
    • Called after the module's database (self.db) has been successfully initialized and the tables defined in self.db_meta have been created (if they didn't exist).
    • This is the guaranteed point where you can safely interact with your module's database immediately after startup.
    • Use this for tasks like loading initial state from the database, performing startup data validation, or initializing database-dependent components.
  • def on_unload(self):
    • Called just before the module is fully unloaded by the ModuleLoader. This happens during operations like module reloading, updating, or uninstallation.
    • Crucial: Use this method to perform any necessary cleanup. This includes:
      • Stopping background tasks or asyncio loops started by the module.
      • Saving any volatile state to disk or database.
      • Closing network connections or releasing other resources.

      WARNING

      Failure to properly clean up in on_unload can lead to resource leaks (memory, file handles, network sockets) or zombie tasks.**

Example usage:

python
# main.py (inside your BaseModule class)
import asyncio

class MyBackgroundTaskModule(BaseModule):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._background_task = None
        self._some_state = {} # Example state

    def on_init(self):
        """Load state from a file, maybe start task if DB not needed."""
        self.logger.info("Module initialized (on_init). Loading state...")
        try:
            # Example: Load state from a hypothetical file
            # with open("module_state.json", "r") as f:
            #     self._some_state = json.load(f)
            pass # Replace with actual loading logic
        except FileNotFoundError:
            self.logger.warning("State file not found, starting fresh.")
            self._some_state = {'runs': 0}

        # If task doesn't need DB immediately, could start here
        # self._start_task()

    async def on_db_ready(self):
        """DB is ready, load DB state, start DB-dependent tasks."""
        self.logger.info("Database is ready (on_db_ready).")
        # Example: Load something from DB
        # async with self.db.session_maker() as session:
        #     config_value = await session.scalar(...)
        #     self._some_state['db_config'] = config_value

        # Start the task now if it depends on the DB
        self._start_task()

    def _start_task(self):
        """Helper to start the background task."""
        if self._background_task is None or self._background_task.done():
            self.logger.info("Starting background task...")
            self._background_task = asyncio.create_task(self._run_background_job())
        else:
             self.logger.warning("Task already running.")

    async def _run_background_job(self):
        """Example async task."""
        try:
            while True:
                self.logger.info(f"Background task running... Run count: {self._some_state.get('runs', 0)}")
                self._some_state['runs'] = self._some_state.get('runs', 0) + 1
                # Access self.db here if needed (ensure it's ready via on_db_ready)
                await asyncio.sleep(60) # Run every minute
        except asyncio.CancelledError:
            self.logger.info("Background task cancelled.")
        except Exception as e:
            self.logger.error(f"Background task failed: {e}", exc_info=True)
        finally:
            self.logger.info("Background task finished.")


    def on_unload(self):
        """Cleanup: Cancel task, save state."""
        # 1. Cancel background task
        if self._background_task and not self._background_task.done():
            self.logger.info("Cancelling background task...")
            self._background_task.cancel()
            # It's good practice to await briefly or handle task completion
            # but keep on_unload synchronous if possible for simplicity during shutdown.
            # For complex cleanup, consider patterns involving asyncio.gather/wait_for
            # within an async helper called from here if strictly needed.

        # 2. Save state
        # try:
        #     self.logger.info("Saving state to file...")
        #     # with open("module_state.json", "w") as f:
        #     #     json.dump(self._some_state, f)
        # except Exception as e:
        #      self.logger.error(f"Failed to save state: {e}")

        self.logger.info("Cleanup complete.")