aa
This commit is contained in:
parent
cbf6d763fb
commit
5ca9881df8
@ -536,13 +536,9 @@ class SettingsCog(commands.Cog, name="Settings"):
|
|||||||
|
|
||||||
async def setup(bot: commands.Bot):
|
async def setup(bot: commands.Bot):
|
||||||
# Ensure pools are initialized before adding the cog
|
# Ensure pools are initialized before adding the cog
|
||||||
if settings_manager.pg_pool is None or settings_manager.redis_pool is None:
|
if getattr(bot, "pg_pool", None) is None or getattr(bot, "redis", None) is None:
|
||||||
log.warning("Settings Manager pools not initialized before loading SettingsCog. Attempting initialization.")
|
log.warning("Bot pools not initialized before loading SettingsCog. Cog will not load.")
|
||||||
try:
|
return # Prevent loading if pools are missing
|
||||||
await settings_manager.initialize_pools()
|
|
||||||
except Exception as e:
|
|
||||||
log.exception("Failed to initialize Settings Manager pools during SettingsCog setup. Cog will not load.")
|
|
||||||
return # Prevent loading if pools fail
|
|
||||||
|
|
||||||
await bot.add_cog(SettingsCog(bot))
|
await bot.add_cog(SettingsCog(bot))
|
||||||
log.info("SettingsCog loaded.")
|
log.info("SettingsCog loaded.")
|
||||||
|
86
main.py
86
main.py
@ -14,6 +14,8 @@ import subprocess
|
|||||||
import importlib.util
|
import importlib.util
|
||||||
import argparse
|
import argparse
|
||||||
import logging # Add logging
|
import logging # Add logging
|
||||||
|
import asyncpg
|
||||||
|
import redis.asyncio as aioredis
|
||||||
from commands import load_all_cogs, reload_all_cogs
|
from commands import load_all_cogs, reload_all_cogs
|
||||||
from error_handler import handle_error, patch_discord_methods, store_interaction_content
|
from error_handler import handle_error, patch_discord_methods, store_interaction_content
|
||||||
from utils import reload_script
|
from utils import reload_script
|
||||||
@ -66,31 +68,52 @@ class MyBot(commands.Bot):
|
|||||||
self.owner_id = int(os.getenv('OWNER_USER_ID'))
|
self.owner_id = int(os.getenv('OWNER_USER_ID'))
|
||||||
self.core_cogs = CORE_COGS # Attach core cogs list to bot instance
|
self.core_cogs = CORE_COGS # Attach core cogs list to bot instance
|
||||||
self.settings_manager = settings_manager # Attach settings manager instance
|
self.settings_manager = settings_manager # Attach settings manager instance
|
||||||
self.pool = None # Will be initialized in setup_hook
|
self.pg_pool = None # Will be initialized in setup_hook
|
||||||
|
self.redis = None # Will be initialized in setup_hook
|
||||||
self.ai_cogs_to_skip = [] # For --disable-ai flag
|
self.ai_cogs_to_skip = [] # For --disable-ai flag
|
||||||
|
|
||||||
async def setup_hook(self):
|
async def setup_hook(self):
|
||||||
# This method is called before the bot logs in but after it's ready.
|
|
||||||
# Ideal place for async initialization.
|
|
||||||
log.info("Running setup_hook...")
|
log.info("Running setup_hook...")
|
||||||
|
|
||||||
# Initialize pools
|
# Create Postgres pool on this loop
|
||||||
log.info("Initializing database/cache pools from setup_hook...")
|
self.pg_pool = await asyncpg.create_pool(
|
||||||
# Pass the bot's event loop to initialize_pools
|
dsn=os.environ["DATABASE_URL"],
|
||||||
pools_initialized = await self.settings_manager.initialize_pools(event_loop=self.loop)
|
min_size=1,
|
||||||
if not pools_initialized:
|
max_size=10,
|
||||||
log.critical("Failed to initialize database/cache pools in setup_hook. Bot may not function correctly.")
|
loop=self.loop # Explicitly use the bot's event loop
|
||||||
# Depending on severity, you might want to prevent the bot from starting.
|
)
|
||||||
# For now, it will continue, but operations requiring pools will fail.
|
log.info("Postgres pool initialized and attached to bot.pg_pool.")
|
||||||
return
|
|
||||||
|
# Create Redis client on this loop
|
||||||
|
self.redis = await aioredis.from_url(
|
||||||
|
os.environ["REDIS_URL"],
|
||||||
|
max_connections=10,
|
||||||
|
decode_responses=True,
|
||||||
|
)
|
||||||
|
log.info("Redis client initialized and attached to bot.redis.")
|
||||||
|
|
||||||
|
# Pass the initialized pools to the settings_manager
|
||||||
|
if self.pg_pool and self.redis:
|
||||||
|
settings_manager.set_bot_pools(self.pg_pool, self.redis)
|
||||||
|
log.info("Bot's pg_pool and redis_client passed to settings_manager.")
|
||||||
|
|
||||||
|
# Initialize database schema and run migrations using settings_manager
|
||||||
|
# These functions will now use the pools provided by the bot.
|
||||||
|
try:
|
||||||
|
await settings_manager.initialize_database() # Uses the pool set via set_bot_pools
|
||||||
|
log.info("Database schema initialization called via settings_manager.")
|
||||||
|
await settings_manager.run_migrations() # Uses the pool set via set_bot_pools
|
||||||
|
log.info("Database migrations called via settings_manager.")
|
||||||
|
except Exception as e:
|
||||||
|
log.exception("CRITICAL: Failed during settings_manager database setup (init/migrations).")
|
||||||
|
else:
|
||||||
|
log.error("CRITICAL: pg_pool or redis_client not initialized in setup_hook. Cannot proceed with settings_manager setup.")
|
||||||
|
|
||||||
self.pool = self.settings_manager.pg_pool # Attach the pool to the bot instance
|
|
||||||
log.info("Database/cache pools initialized and pg_pool attached to bot instance.")
|
|
||||||
|
|
||||||
# Setup the moderation log table *after* pool initialization
|
# Setup the moderation log table *after* pool initialization
|
||||||
if self.pool:
|
if self.pg_pool:
|
||||||
try:
|
try:
|
||||||
await mod_log_db.setup_moderation_log_table(self.pool)
|
await mod_log_db.setup_moderation_log_table(self.pg_pool)
|
||||||
log.info("Moderation log table setup complete via setup_hook.")
|
log.info("Moderation log table setup complete via setup_hook.")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.exception("CRITICAL: Failed to setup moderation log table in setup_hook.")
|
log.exception("CRITICAL: Failed to setup moderation log table in setup_hook.")
|
||||||
@ -98,8 +121,6 @@ class MyBot(commands.Bot):
|
|||||||
log.warning("pg_pool not available in setup_hook, skipping mod_log_db setup.")
|
log.warning("pg_pool not available in setup_hook, skipping mod_log_db setup.")
|
||||||
|
|
||||||
# Load all cogs from the 'cogs' directory, skipping AI if requested
|
# Load all cogs from the 'cogs' directory, skipping AI if requested
|
||||||
# ai_cogs_to_skip needs to be set based on args before bot.start is called.
|
|
||||||
# This is handled by setting bot.ai_cogs_to_skip in main() before bot.start().
|
|
||||||
await load_all_cogs(self, skip_cogs=self.ai_cogs_to_skip)
|
await load_all_cogs(self, skip_cogs=self.ai_cogs_to_skip)
|
||||||
log.info(f"Cogs loaded in setup_hook. Skipped: {self.ai_cogs_to_skip or 'None'}")
|
log.info(f"Cogs loaded in setup_hook. Skipped: {self.ai_cogs_to_skip or 'None'}")
|
||||||
|
|
||||||
@ -156,10 +177,10 @@ async def on_ready():
|
|||||||
log.info("Bot status set to 'Listening to !help'")
|
log.info("Bot status set to 'Listening to !help'")
|
||||||
|
|
||||||
# --- Add current guilds to DB ---
|
# --- Add current guilds to DB ---
|
||||||
if settings_manager and settings_manager.pg_pool:
|
if bot.pg_pool:
|
||||||
log.info("Syncing guilds with database...")
|
log.info("Syncing guilds with database...")
|
||||||
try:
|
try:
|
||||||
async with settings_manager.pg_pool.acquire() as conn:
|
async with bot.pg_pool.acquire() as conn:
|
||||||
# Get guilds bot is currently in
|
# Get guilds bot is currently in
|
||||||
current_guild_ids = {guild.id for guild in bot.guilds}
|
current_guild_ids = {guild.id for guild in bot.guilds}
|
||||||
log.debug(f"Bot is currently in {len(current_guild_ids)} guilds.")
|
log.debug(f"Bot is currently in {len(current_guild_ids)} guilds.")
|
||||||
@ -186,7 +207,7 @@ async def on_ready():
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.exception("Error syncing guilds with database on ready.")
|
log.exception("Error syncing guilds with database on ready.")
|
||||||
else:
|
else:
|
||||||
log.warning("Settings manager not available or pool not initialized, skipping guild sync.")
|
log.warning("Bot Postgres pool not initialized, skipping guild sync.")
|
||||||
# -----------------------------
|
# -----------------------------
|
||||||
|
|
||||||
# Patch Discord methods to store message content
|
# Patch Discord methods to store message content
|
||||||
@ -240,9 +261,9 @@ async def on_shard_disconnect(shard_id):
|
|||||||
async def on_guild_join(guild: discord.Guild):
|
async def on_guild_join(guild: discord.Guild):
|
||||||
"""Adds guild to database when bot joins and syncs commands."""
|
"""Adds guild to database when bot joins and syncs commands."""
|
||||||
log.info(f"Joined guild: {guild.name} ({guild.id})")
|
log.info(f"Joined guild: {guild.name} ({guild.id})")
|
||||||
if settings_manager and settings_manager.pg_pool:
|
if bot.pg_pool:
|
||||||
try:
|
try:
|
||||||
async with settings_manager.pg_pool.acquire() as conn:
|
async with bot.pg_pool.acquire() as conn:
|
||||||
await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT DO NOTHING;", guild.id)
|
await conn.execute("INSERT INTO guilds (guild_id) VALUES ($1) ON CONFLICT DO NOTHING;", guild.id)
|
||||||
log.info(f"Added guild {guild.id} to database.")
|
log.info(f"Added guild {guild.id} to database.")
|
||||||
|
|
||||||
@ -256,22 +277,22 @@ async def on_guild_join(guild: discord.Guild):
|
|||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.exception(f"Failed to add guild {guild.id} to database on join.")
|
log.exception(f"Failed to add guild {guild.id} to database on join.")
|
||||||
else:
|
else:
|
||||||
log.warning("Settings manager not available or pool not initialized, cannot add guild on join.")
|
log.warning("Bot Postgres pool not initialized, cannot add guild on join.")
|
||||||
|
|
||||||
@bot.event
|
@bot.event
|
||||||
async def on_guild_remove(guild: discord.Guild):
|
async def on_guild_remove(guild: discord.Guild):
|
||||||
"""Removes guild from database when bot leaves."""
|
"""Removes guild from database when bot leaves."""
|
||||||
log.info(f"Left guild: {guild.name} ({guild.id})")
|
log.info(f"Left guild: {guild.name} ({guild.id})")
|
||||||
if settings_manager and settings_manager.pg_pool:
|
if bot.pg_pool:
|
||||||
try:
|
try:
|
||||||
async with settings_manager.pg_pool.acquire() as conn:
|
async with bot.pg_pool.acquire() as conn:
|
||||||
# Note: Cascading deletes should handle related settings in other tables
|
# Note: Cascading deletes should handle related settings in other tables
|
||||||
await conn.execute("DELETE FROM guilds WHERE guild_id = $1", guild.id)
|
await conn.execute("DELETE FROM guilds WHERE guild_id = $1", guild.id)
|
||||||
log.info(f"Removed guild {guild.id} from database.")
|
log.info(f"Removed guild {guild.id} from database.")
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
log.exception(f"Failed to remove guild {guild.id} from database on leave.")
|
log.exception(f"Failed to remove guild {guild.id} from database on leave.")
|
||||||
else:
|
else:
|
||||||
log.warning("Settings manager not available or pool not initialized, cannot remove guild on leave.")
|
log.warning("Bot Postgres pool not initialized, cannot remove guild on leave.")
|
||||||
|
|
||||||
|
|
||||||
# Error handling - Updated to handle custom check failures
|
# Error handling - Updated to handle custom check failures
|
||||||
@ -553,10 +574,13 @@ async def main(args): # Pass parsed args
|
|||||||
log.info("Flask server process was not running or already terminated.")
|
log.info("Flask server process was not running or already terminated.")
|
||||||
|
|
||||||
# Close database/cache pools if they were initialized
|
# Close database/cache pools if they were initialized
|
||||||
if settings_manager.pg_pool or settings_manager.redis_pool:
|
if bot.pg_pool:
|
||||||
log.info("Closing database/cache pools in main finally block...")
|
log.info("Closing Postgres pool in main finally block...")
|
||||||
await settings_manager.close_pools()
|
await bot.pg_pool.close()
|
||||||
else:
|
if bot.redis:
|
||||||
|
log.info("Closing Redis pool in main finally block...")
|
||||||
|
await bot.redis.close()
|
||||||
|
if not bot.pg_pool and not bot.redis:
|
||||||
log.info("Pools were not initialized or already closed, skipping close_pools in main.")
|
log.info("Pools were not initialized or already closed, skipping close_pools in main.")
|
||||||
|
|
||||||
# Run the main async function
|
# Run the main async function
|
||||||
|
File diff suppressed because it is too large
Load Diff
Loading…
x
Reference in New Issue
Block a user