diff --git a/src/core/exporter.py b/src/core/exporter.py index c62cc28..a006190 100644 --- a/src/core/exporter.py +++ b/src/core/exporter.py @@ -18,6 +18,7 @@ class DiscordExporter: self.server_name = "" self.server_id = "" 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.is_running = True self.db: Optional[BackupDatabase] = None @@ -62,6 +63,20 @@ class DiscordExporter: hash_sha256.update(chunk) 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): """Saves server metadata to the SQLite database.""" metadata = await self.reader.get_server_metadata() @@ -476,6 +491,10 @@ class DiscordExporter: # Queue for deferred download 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 # 0: Regular User, 1: Bot, 2: Webhook, 3: System u_type = 0 @@ -519,12 +538,19 @@ class DiscordExporter: # 1. Author handling 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) # 1.5 Mentions handling (ensure all mentioned users are saved) if 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) if u_ment: new_users.append(u_ment) diff --git a/src/ui/shuttle_ops.py b/src/ui/shuttle_ops.py index d137832..021ef21 100644 --- a/src/ui/shuttle_ops.py +++ b/src/ui/shuttle_ops.py @@ -329,8 +329,13 @@ class OperationPane(Container): 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) - for bid in ("#op_backup_msgs", "#op_backup_sync", "#op_autotest"): - for btn in self.query(bid): btn.disabled = not enabled + for btn in self.query("#op_backup_msgs"): + 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"): btn.display = self.has_backup @@ -2379,6 +2384,10 @@ class OperationPane(Container): await self.exporter.export_roles() 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 total_chans = len(selected_channels) modal.write(f"\n[bold cyan]Backing up {total_chans} channels...[/bold cyan]")