k.eaven/Pterodisbot
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 5 Wiki Activity

feat: Multi-guild support implementation

Browse Source 
- Added multi-guild support with comma-separated AllowedGuildIDs config
- Per-guild embed tracking with isolated state
- Per-server action locking to prevent race conditions across guilds
- Cross-guild embed updates after power actions
- Debug logging for embed update triggers
- Fixed resource usage field formatting (brackets now properly closed)
- Updated config.ini.sample, generate_config.py for backward compatibility
- Updated README.md with multi-guild documentation
- Updated tests for multi-guild support
This commit is contained in:
k.eaven committed 2026-10-30 10:37:13 +00:00
1 parent e68b248033
commit bec06bd909
10 files changed
+802 -408
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

embeds.py
+155 -56
View File
@@ -3,10 +3,12 @@ Embed Management Module
Handles creation, tracking, and persistence of Discord server status embeds.
Manages embed locations storage and provides utilities for embed lifecycle.
Supports multi-guild deployments where each guild tracks servers independently.
"""
import json
import logging
from collections import defaultdict
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple
@@ -24,14 +26,15 @@ logger = logging.getLogger("pterodisbot")
class EmbedManager:
"""
Manages server status embeds lifecycle.
Manages server status embeds lifecycle across multiple guilds.
Handles:
- Embed creation with real-time server data
- Persistent storage of embed locations (channel_id, message_id)
- Embed refresh and recreation
- Persistent storage of embed locations per guild (guild_id -> server_id -> location)
- Embed refresh and recreation per guild
- Cleanup of deleted/missing embeds
- Rate limit protection during bulk operations
- Migration from legacy single-guild storage format
Attributes:
config: Application configuration
@@ -58,15 +61,16 @@ class EmbedManager:
self.api = api
self.metrics_manager = metrics_manager
self.storage_path = config.embed_storage_path
self.embed_locations: Dict[str, Dict[str, int]] = {}
self.previous_states: Dict[str, Tuple[str, float, Optional[float]]] = {}
# Structure: {guild_id: {server_id: {"channel_id": str, "message_id": str}}}
self.embed_locations: Dict[int, Dict[str, Dict[str, str]]] = defaultdict(dict)
# Previous states tracking moved to bot.py for multi-guild support
# Ensure storage directory exists
self.storage_path.parent.mkdir(parents=True, exist_ok=True)
logger.info("Initialized EmbedManager")
logger.info("Initialized EmbedManager with multi-guild support")
async def load_locations(self) -> None:
"""Load saved embed locations from JSON storage file."""
"""Load saved embed locations from JSON storage file with migration support."""
logger.debug("Loading embed locations from storage")
if not self.storage_path.exists():
logger.info("No existing embed locations file found")
@@ -74,8 +78,47 @@ class EmbedManager:
try:
with open(self.storage_path, "r") as f:
self.embed_locations = json.load(f)
logger.info(f"Loaded {len(self.embed_locations)} embed locations from storage")
data = json.load(f)
# Detect format:
# New format: {guild_id: {server_id: location}} where values are dicts with channel_id/message_id
# Old format: {server_id: location} where location has channel_id/message_id directly
is_new_format = False
if data:
first_key = next(iter(data.keys()))
first_value = data[first_key]
# New format: value is a dict of server locations
# Old format: value is a location dict with channel_id/message_id
if isinstance(first_value, dict) and "channel_id" not in first_value:
# Value is {server_id: {channel_id, message_id}} - new format
is_new_format = True
elif isinstance(first_value, dict) and "channel_id" in first_value:
# Value is {channel_id, message_id} - old format
is_new_format = False
# If first_key is numeric and >= 17 digits, also treat as new format
elif first_key.isdigit() and len(first_key) >= 17:
is_new_format = True
if is_new_format:
# New format: {guild_id: {server_id: location}}
self.embed_locations = defaultdict(dict)
for guild_id_str, servers in data.items():
guild_id = int(guild_id_str)
self.embed_locations[guild_id] = servers
total_embeds = sum(len(s) for s in self.embed_locations.values())
logger.info(f"Loaded embed locations for {len(self.embed_locations)} guilds ({total_embeds} total embeds)")
else:
# Legacy format: {server_id: location} - migrate to first allowed guild
legacy_guild_id = self.config.allowed_guild_ids[0] if self.config.allowed_guild_ids else 0
if legacy_guild_id:
self.embed_locations = defaultdict(dict)
self.embed_locations[legacy_guild_id] = data
logger.info(f"Migrated legacy embed locations to guild {legacy_guild_id}")
# Save in new format immediately
await self.save_locations()
else:
logger.warning("No allowed guilds configured, cannot migrate legacy embed locations")
except Exception as e:
logger.error(f"Failed to load embed locations: {str(e)}")
@@ -83,40 +126,72 @@ class EmbedManager:
"""Save current embed locations to JSON storage file."""
logger.debug("Saving embed locations to storage")
try:
# Convert defaultdict to regular dict for JSON serialization
# Guild IDs as strings for JSON compatibility
data = {str(gid): servers for gid, servers in self.embed_locations.items()}
with open(self.storage_path, "w") as f:
json.dump(self.embed_locations, f, indent=2)
json.dump(data, f, indent=2)
logger.debug("Successfully saved embed locations")
except Exception as e:
logger.error(f"Failed to save embed locations: {str(e)}")
async def track_embed(self, server_id: str, message: discord.Message) -> None:
def get_guild_locations(self, guild_id: int) -> Dict[str, Dict[str, str]]:
"""
Track a newly created embed in storage.
Get all tracked embeds for a specific guild.
Args:
guild_id: Discord guild ID
Returns:
Dictionary mapping server_id to location info
"""
return self.embed_locations.get(guild_id, {})
def get_all_tracked_servers(self) -> List[Tuple[int, str]]:
"""
Get all (guild_id, server_id) pairs currently being tracked.
Returns:
List of (guild_id, server_id) tuples
"""
result = []
for guild_id, servers in self.embed_locations.items():
for server_id in servers:
result.append((guild_id, server_id))
return result
async def track_embed(self, guild_id: int, server_id: str, message: discord.Message) -> None:
"""
Track a newly created embed in storage for a specific guild.
Args:
guild_id: Discord guild ID
server_id: The server's panel identifier
message: Discord message containing the embed
"""
logger.debug(f"Tracking new embed for server {server_id} in channel {message.channel.id}")
self.embed_locations[server_id] = {
logger.debug(f"Tracking new embed for server {server_id} in guild {guild_id}, channel {message.channel.id}")
self.embed_locations[guild_id][server_id] = {
"channel_id": str(message.channel.id),
"message_id": str(message.id),
}
await self.save_locations()
async def untrack_embed(self, server_id: str) -> None:
async def untrack_embed(self, guild_id: int, server_id: str) -> None:
"""
Remove a server's embed from tracking.
Remove a server's embed from tracking for a specific guild.
Args:
guild_id: Discord guild ID
server_id: The server's panel identifier
"""
self.embed_locations.pop(server_id, None)
self.previous_states.pop(server_id, None)
self.embed_locations[guild_id].pop(server_id, None)
# Clean up empty guild entries
if not self.embed_locations[guild_id]:
del self.embed_locations[guild_id]
await self.save_locations()
async def create_server_embed(
self, server_data: ServerData, resources: ServerResources
self, server_data: ServerData, resources: ServerResources, guild_id: int
) -> Tuple[discord.Embed, discord.ui.View]:
"""
Create a status embed and interactive view for a server.
@@ -124,6 +199,7 @@ class EmbedManager:
Args:
server_data: Server information from panel
resources: Current resource usage data
guild_id: Discord guild ID this embed belongs to
Returns:
Tuple of (embed, view) ready for display
@@ -133,7 +209,7 @@ class EmbedManager:
name = attributes.name
description = attributes.description or "No description available"
logger.debug(f"Building status embed for server {name} ({identifier})")
logger.debug(f"Building status embed for server {name} ({identifier}) in guild {guild_id}")
# Parse resource data
current_state = resources.attributes.current_state.title()
@@ -165,7 +241,7 @@ class EmbedManager:
embed.set_footer(text="Last updated")
# Create interactive view with control buttons
# Create interactive view with control buttons (pass guild_id for validation)
from views import ServerStatusView
view = ServerStatusView(
@@ -173,9 +249,10 @@ class EmbedManager:
server_name=name,
api=self.api,
server_data=server_data.model_dump(),
guild_id=guild_id,
)
logger.debug(f"Successfully built status components for {name}")
logger.debug(f"Successfully built status components for {name} in guild {guild_id}")
return embed, view
def _add_resource_fields(
@@ -198,10 +275,10 @@ class EmbedManager:
disk_limit = limits.disk
# Format limit values - display ∞ for unlimited (0 limit)
def format_limit(value: int, unit: str = "") -> str:
def format_limit(value: int) -> str:
if value == 0:
return f"{'∞':>8} {unit}"
return f"{value:>8} {unit}"
return f"{'∞':>8}"
return f"{value:>8}"
# Format uptime
uptime_ms = resource_attrs.uptime
@@ -227,9 +304,9 @@ class EmbedManager:
# Resource usage text box
usage_text = (
f"```properties\n"
f"CPU : [{cpu_usage:>8} / {format_limit(cpu_limit, '%')}\n"
f"Memory : [{memory_usage:>8} / {format_limit(memory_limit, 'MiB')}\n"
f"Disk : [{disk_usage:>8} / {format_limit(disk_limit, 'MiB')}\n"
f"CPU : [{cpu_usage:>8} / {format_limit(cpu_limit)}] %\n"
f"Memory : [{memory_usage:>8} / {format_limit(memory_limit)}] MiB\n"
f"Disk : [{disk_usage:>8} / {format_limit(disk_limit)}] MiB\n"
f"```"
)
@@ -261,16 +338,19 @@ class EmbedManager:
return embed
async def refresh_all_embeds(self) -> Tuple[int, int]:
async def refresh_all_embeds(self, guild_id: int) -> Tuple[int, int]:
"""
Perform a complete refresh of all server status embeds.
Perform a complete refresh of all server status embeds for a specific guild.
Creates new embeds and deletes old ones to prevent duplication.
Args:
guild_id: Discord guild ID to refresh embeds for
Returns:
Tuple of (deleted_count, created_count)
"""
logger.info("Starting full refresh of all server embeds")
logger.info(f"Starting full refresh of all server embeds for guild {guild_id}")
from bot import PterodactylBot # Avoid circular import
bot = PterodactylBot.get_instance()
@@ -278,6 +358,11 @@ class EmbedManager:
created_count = 0
skipped_count = 0
guild_locations = self.get_guild_locations(guild_id)
if not guild_locations:
logger.info(f"No embeds tracked for guild {guild_id}, nothing to refresh")
return 0, 0
try:
# Get current server list
servers = await self.api.get_servers()
@@ -291,14 +376,19 @@ class EmbedManager:
new_embeds = {}
for server_id, server_data in server_cache.items():
if server_id not in self.embed_locations:
if server_id not in guild_locations:
skipped_count += 1
continue
channel_id = self.embed_locations[server_id]["channel_id"]
channel_id = guild_locations[server_id]["channel_id"]
channel = bot.get_channel(int(channel_id))
if not channel:
logger.warning(f"Channel {channel_id} not found for server {server_id}")
logger.warning(f"Channel {channel_id} not found for server {server_id} in guild {guild_id}")
continue
# Verify channel belongs to the correct guild
if channel.guild.id != guild_id:
logger.warning(f"Channel {channel_id} not in guild {guild_id}, skipping")
continue
try:
@@ -306,7 +396,7 @@ class EmbedManager:
server_obj = ServerData.model_validate(server_data)
resources_obj = ServerResources.model_validate(resources)
embed, view = await self.create_server_embed(server_obj, resources_obj)
embed, view = await self.create_server_embed(server_obj, resources_obj, guild_id)
message = await channel.send(embed=embed, view=view)
new_embeds[server_id] = {
@@ -314,48 +404,53 @@ class EmbedManager:
"message_id": str(message.id),
}
created_count += 1
logger.info(f"Created new embed for server {server_data['attributes']['name']}")
logger.info(f"Created new embed for server {server_data['attributes']['name']} in guild {guild_id}")
await asyncio.sleep(1) # Rate limit protection
except Exception as e:
logger.error(f"Failed to create new embed for server {server_id}: {str(e)}")
logger.error(f"Failed to create new embed for server {server_id} in guild {guild_id}: {str(e)}")
if not new_embeds:
logger.warning("No new embeds created during refresh")
logger.warning(f"No new embeds created during refresh for guild {guild_id}")
return 0, 0
# Delete old embeds
for server_id, location in list(self.embed_locations.items()):
for server_id, location in list(guild_locations.items()):
try:
channel = bot.get_channel(int(location["channel_id"]))
if channel:
# Verify channel belongs to the correct guild
if channel.guild.id != guild_id:
logger.warning(f"Channel {location['channel_id']} not in guild {guild_id}, skipping delete")
continue
try:
message = await channel.fetch_message(int(location["message_id"]))
await message.delete()
deleted_count += 1
logger.debug(f"Deleted old embed for server {server_id}")
logger.debug(f"Deleted old embed for server {server_id} in guild {guild_id}")
await asyncio.sleep(0.5)
except discord.NotFound:
logger.debug(f"Old embed for server {server_id} already deleted")
logger.debug(f"Old embed for server {server_id} in guild {guild_id} already deleted")
except Exception as e:
logger.error(f"Failed to delete old embed for server {server_id}: {str(e)}")
logger.error(f"Failed to delete old embed for server {server_id} in guild {guild_id}: {str(e)}")
except Exception as e:
logger.error(f"Error processing old embed for server {server_id}: {str(e)}")
logger.error(f"Error processing old embed for server {server_id} in guild {guild_id}: {str(e)}")
logger.info(f"Deleted {deleted_count} old embeds, created {created_count} new embeds")
logger.info(f"Guild {guild_id}: Deleted {deleted_count} old embeds, created {created_count} new embeds")
# Update storage
self.embed_locations = new_embeds
# Update storage for this guild only
self.embed_locations[guild_id] = new_embeds
await self.save_locations()
return deleted_count, created_count
except Exception as e:
logger.error(f"Critical error during embed refresh: {str(e)}")
logger.error(f"Critical error during embed refresh for guild {guild_id}: {str(e)}")
raise
async def update_single_embed(
self,
guild_id: int,
server_id: str,
location: Dict[str, str],
server_data: ServerData,
@@ -364,9 +459,10 @@ class EmbedManager:
cpu_usage: float,
) -> bool:
"""
Update a single server's status embed.
Update a single server's status embed for a specific guild.
Args:
guild_id: Discord guild ID
server_id: Server identifier
location: Embed location dict with channel_id and message_id
server_data: Current server data
@@ -381,11 +477,16 @@ class EmbedManager:
bot = PterodactylBot.get_instance()
try:
embed, view = await self.create_server_embed(server_data, resources)
embed, view = await self.create_server_embed(server_data, resources, guild_id)
channel = bot.get_channel(int(location["channel_id"]))
if not channel:
logger.warning(f"Channel {location['channel_id']} not found for server {server_id}")
logger.warning(f"Channel {location['channel_id']} not found for server {server_id} in guild {guild_id}")
return False
# Verify channel belongs to the correct guild
if channel.guild.id != guild_id:
logger.warning(f"Channel {location['channel_id']} not in guild {guild_id}")
return False
message = await channel.fetch_message(int(location["message_id"]))
@@ -406,17 +507,15 @@ class EmbedManager:
# Empty files list removes existing attachments
await message.edit(embed=embed, view=view, attachments=files)
logger.debug(f"Updated status for {server_data.attributes.name}")
logger.debug(f"Updated status for {server_data.attributes.name} in guild {guild_id}")
return True
except discord.NotFound:
logger.warning(f"Embed for server {server_id} not found, removing from tracking")
self.embed_locations.pop(server_id, None)
self.previous_states.pop(server_id, None)
await self.save_locations()
logger.warning(f"Embed for server {server_id} in guild {guild_id} not found, removing from tracking")
await self.untrack_embed(guild_id, server_id)
return False
except Exception as e:
logger.error(f"Failed to update status for server {server_id}: {str(e)}")
logger.error(f"Failed to update status for server {server_id} in guild {guild_id}: {str(e)}")
return False