backup user roles also

This commit is contained in:
rambros 2026-04-01 13:05:49 +05:30
parent e5ff2ae37f
commit d63ed0c440
2 changed files with 38 additions and 3 deletions

View file

@ -18,6 +18,7 @@ class DiscordExporter:
self.server_name = "" self.server_name = ""
self.server_id = "" self.server_id = ""
self.user_cache = {} self.user_cache = {}
self.member_cache: Dict[int, Any] = {} # Pre-fetched member objects (id -> Member)
self.base_dir = Path(base_dir) if base_dir else Path(".") self.base_dir = Path(base_dir) if base_dir else Path(".")
self.is_running = True self.is_running = True
self.db: Optional[BackupDatabase] = None self.db: Optional[BackupDatabase] = None
@ -62,6 +63,20 @@ class DiscordExporter:
hash_sha256.update(chunk) hash_sha256.update(chunk)
return hash_sha256.hexdigest() return hash_sha256.hexdigest()
async def prefetch_members(self):
"""Pre-fetches all guild members into a local cache for role resolution.
msg.author is a discord.User (no roles). This cache allows us to
resolve roles without an API call per message during message export.
"""
try:
members = await self.reader.get_members()
self.member_cache = {m.id: m for m in members}
logger.info(f"Pre-fetched {len(self.member_cache)} members for role resolution.")
except Exception as e:
logger.warning(f"Could not pre-fetch members (roles will be empty): {e}")
self.member_cache = {}
async def export_metadata(self): async def export_metadata(self):
"""Saves server metadata to the SQLite database.""" """Saves server metadata to the SQLite database."""
metadata = await self.reader.get_server_metadata() metadata = await self.reader.get_server_metadata()
@ -476,6 +491,10 @@ class DiscordExporter:
# Queue for deferred download # Queue for deferred download
self._pending_avatars.append((user_id, avatar, av_target)) self._pending_avatars.append((user_id, avatar, av_target))
roles = []
if hasattr(user, "roles"):
roles = [str(r.id) for r in user.roles if not r.is_default()]
# Determine user type # Determine user type
# 0: Regular User, 1: Bot, 2: Webhook, 3: System # 0: Regular User, 1: Bot, 2: Webhook, 3: System
u_type = 0 u_type = 0
@ -519,12 +538,19 @@ class DiscordExporter:
# 1. Author handling # 1. Author handling
is_webhook = bool(getattr(msg, "webhook_id", None)) is_webhook = bool(getattr(msg, "webhook_id", None))
u_data = await self._format_user(msg.author, is_webhook=is_webhook) author = msg.author
# msg.author is discord.User (no roles). Resolve to Member for role data.
if not is_webhook:
member = self.member_cache.get(msg.author.id)
if member:
author = member
u_data = await self._format_user(author, is_webhook=is_webhook)
if u_data: new_users.append(u_data) if u_data: new_users.append(u_data)
# 1.5 Mentions handling (ensure all mentioned users are saved) # 1.5 Mentions handling (ensure all mentioned users are saved)
if msg.mentions: if msg.mentions:
for mention in msg.mentions: for mention in msg.mentions:
# Mentions can be Member objects already, so roles work naturally
u_ment = await self._format_user(mention, is_webhook=False) u_ment = await self._format_user(mention, is_webhook=False)
if u_ment: new_users.append(u_ment) if u_ment: new_users.append(u_ment)

View file

@ -329,8 +329,13 @@ class OperationPane(Container):
for pne in self.query("#op_target_pane"): pne.display = False for pne in self.query("#op_target_pane"): pne.display = False
enabled = (v.get("discord_token") and v.get("discord_server") and not d_missing) enabled = (v.get("discord_token") and v.get("discord_server") and not d_missing)
for bid in ("#op_backup_msgs", "#op_backup_sync", "#op_autotest"): for btn in self.query("#op_backup_msgs"):
for btn in self.query(bid): btn.disabled = not enabled btn.disabled = not enabled
for btn in self.query("#op_backup_sync"):
btn.display = self.has_backup
btn.disabled = not (enabled and self.has_backup)
for btn in self.query("#op_autotest"):
btn.disabled = not enabled
for btn in self.query("#op_backup_stats"): for btn in self.query("#op_backup_stats"):
btn.display = self.has_backup btn.display = self.has_backup
@ -2379,6 +2384,10 @@ class OperationPane(Container):
await self.exporter.export_roles() await self.exporter.export_roles()
await self.exporter.export_assets() await self.exporter.export_assets()
# Pre-fetch all members once for role resolution during message export
modal.set_status("Pre-fetching server members...")
await self.exporter.prefetch_members()
# 2. Channel Messages # 2. Channel Messages
total_chans = len(selected_channels) total_chans = len(selected_channels)
modal.write(f"\n[bold cyan]Backing up {total_chans} channels...[/bold cyan]") modal.write(f"\n[bold cyan]Backing up {total_chans} channels...[/bold cyan]")