ad3laid3/disco-reaper
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix progress display for forums & threads

Browse source
This commit is contained in:
rambros 2026-03-02 01:59:57 +05:30
parent b18ee0c5d3
commit 4e9c61e256
1 changed files with 64 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

89
src/disco_reaper/exporter.py
View file

@ -265,8 +265,8 @@ class DiscordExporter:
return data return data
async def export_channel_messages(self, channel_id: int, progress_callback=None, force=False): async def export_channel_messages(self, channel_id: int, progress_callback=None, force=False, accumulated_count=0):
"""Fetches and saves message history for a channel, handling incremental sync.""" """Fetches and saves message history for a channel, handling incremental sync. Returns the total messages processed."""
channel = await self.reader.get_channel(channel_id) channel = await self.reader.get_channel(channel_id)
if not channel: if not channel:
logger.error(f"Channel not found: {channel_id}") logger.error(f"Channel not found: {channel_id}")
@ -356,14 +356,19 @@ class DiscordExporter:
msg_data = await self._format_message(msg, asset_dir, base_filename, avatar_dir, avatar_rel_base) msg_data = await self._format_message(msg, asset_dir, base_filename, avatar_dir, avatar_rel_base)
messages.append(msg_data) messages.append(msg_data)
new_count += 1 new_count += 1
accumulated_count += 1
if progress_callback: if progress_callback:
await progress_callback(channel_name, new_count) await progress_callback(channel_name, accumulated_count)
except discord.Forbidden: except discord.Forbidden:
logger.error(f"403 Forbidden: Missing Access to read messages in {channel_name} ({channel_id})") logger.error(f"403 Forbidden: Missing Access to read messages in {channel_name} ({channel_id})")
if not messages: return 0 if not messages: return accumulated_count
except Exception as e: except Exception as e:
logger.error(f"Error fetching messages for {channel_name}: {e}") logger.error(f"Error fetching messages for {channel_name}: {e}")
if not messages: return 0 if not messages: return accumulated_count
# If it's a forum or a channel with no new messages, we still want the UI to register that we've started it.
if new_count == 0 and progress_callback:
await progress_callback(channel_name, accumulated_count)
# 2. Handle Threads and collect counts accurately # 2. Handle Threads and collect counts accurately
all_threads = [] all_threads = []
@ -405,6 +410,8 @@ class DiscordExporter:
"threadCount": thread_count, "threadCount": thread_count,
"lastMessageID": str(messages[-1]["messageID"]) if messages else None, "lastMessageID": str(messages[-1]["messageID"]) if messages else None,
"threadMessagesCount": thread_msg_count, "threadMessagesCount": thread_msg_count,
"totalAttachmentSizeBytes": sum(m.get("totalFileSizeBytes", 0) for m in messages),
"numberOfAttachments": sum(m.get("numberOfFiles", 0) for m in messages),
"lastBackup": discord.utils.utcnow().isoformat(), "lastBackup": discord.utils.utcnow().isoformat(),
"messages": messages "messages": messages
} }
@ -429,9 +436,9 @@ class DiscordExporter:
# If it's a forum, also export its threads into the sub-directory # If it's a forum, also export its threads into the sub-directory
if is_forum: if is_forum:
await self.export_threads(channel_id, progress_callback=progress_callback, force=force) accumulated_count = await self.export_threads(channel_id, progress_callback=progress_callback, force=force, accumulated_count=accumulated_count)
return count return accumulated_count
async def _format_message(self, msg, asset_dir, asset_prefix, avatar_dir, avatar_rel_base): async def _format_message(self, msg, asset_dir, asset_prefix, avatar_dir, avatar_rel_base):
"""Formats a single message to match the reference format.""" """Formats a single message to match the reference format."""
@ -575,6 +582,8 @@ class DiscordExporter:
"content": msg_content, "content": msg_content,
"userID": user_id, "userID": user_id,
"attachments": attachments, "attachments": attachments,
"numberOfFiles": len(attachments),
"totalFileSizeBytes": sum(a["fileSizeBytes"] for a in attachments),
"embeds": [e.to_dict() for e in msg.embeds], "embeds": [e.to_dict() for e in msg.embeds],
"stickers": stickers, "stickers": stickers,
"reactions": reactions "reactions": reactions
@ -600,8 +609,8 @@ class DiscordExporter:
return data return data
async def export_threads(self, channel_id: int, progress_callback=None, force=False): async def export_threads(self, channel_id: int, progress_callback=None, force=False, accumulated_count=0):
"""Exports active and archived threads for a channel.""" """Exports active and archived threads for a channel. Returns accumulated message count."""
channel = await self.reader.get_channel(channel_id) channel = await self.reader.get_channel(channel_id)
if not hasattr(channel, "threads") and not hasattr(channel, "public_archived_threads"): if not hasattr(channel, "threads") and not hasattr(channel, "public_archived_threads"):
return 0 return 0
@ -636,7 +645,11 @@ class DiscordExporter:
logger.info(f"Found {len(all_threads)} threads in {channel.name}. Starting backup...") logger.info(f"Found {len(all_threads)} threads in {channel.name}. Starting backup...")
for thread in all_threads: for thread in all_threads:
# Whenever a forum thread backup starts, populate the forum root json with the starter message. # First backup the full thread — this creates {thread_id}.json with totalAttachmentSizeBytes
accumulated_count = await self.export_channel_messages(thread.id, progress_callback=progress_callback, force=force, accumulated_count=accumulated_count)
thread_count += 1
# Then populate the forum root JSON with the starter message
if is_forum: if is_forum:
logger.info(f"Adding starter message for thread: {thread.name} ({thread.id})") logger.info(f"Adding starter message for thread: {thread.name} ({thread.id})")
try: try:
@ -664,6 +677,23 @@ class DiscordExporter:
# Store applied tag IDs (as strings) — names are resolvable via the forum's available_tags # Store applied tag IDs (as strings) — names are resolvable via the forum's available_tags
msg_data["tags"] = [str(tid) for tid in getattr(thread, "_applied_tags", [])] msg_data["tags"] = [str(tid) for tid in getattr(thread, "_applied_tags", [])]
# Enrich totalFileSizeBytes with the child thread's totalAttachmentSizeBytes
# (the thread JSON has already been written above)
thread_json = backup_root / str(channel_id) / f"{thread.id}.json"
if thread_json.exists():
try:
with open(thread_json, "r", encoding="utf-8") as f:
thread_data = json.load(f)
child_size = thread_data.get("totalAttachmentSizeBytes", 0)
msg_data["totalFileSizeBytes"] = msg_data.get("totalFileSizeBytes", 0) + child_size
child_count = thread_data.get("numberOfAttachments", 0)
msg_data["numberOfFiles"] = msg_data.get("numberOfFiles", 0) + child_count
logger.debug(f"Enriched files for {thread.name}: +{child_size} bytes, +{child_count} files from child thread")
except Exception as e:
logger.error(f"Failed to read thread JSON for size enrichment: {e}")
if forum_json_file.exists(): if forum_json_file.exists():
with open(forum_json_file, "r", encoding="utf-8") as f: with open(forum_json_file, "r", encoding="utf-8") as f:
try: try:
@ -675,18 +705,29 @@ class DiscordExporter:
if "messages" not in forum_data: if "messages" not in forum_data:
forum_data["messages"] = [] forum_data["messages"] = []
# Avoid duplicates # Avoid duplicates — update if already exists (e.g. sync run)
if not any(m["messageID"] == msg_data["messageID"] for m in forum_data["messages"]): existing = next((m for m in forum_data["messages"] if m["messageID"] == msg_data["messageID"]), None)
forum_data["messages"].append(msg_data) if existing:
forum_data["messageCount"] = len(forum_data["messages"]) existing.update(msg_data)
# Keep chronological order logger.debug(f"Updated starter message for {thread.name} in forum JSON")
forum_data["messages"].sort(key=lambda x: x["timestamp"])
with open(forum_json_file, "w", encoding="utf-8") as f:
json.dump(forum_data, f, indent=4, ensure_ascii=False)
logger.info(f"Appended starter message for {thread.name} to {forum_json_file.name}")
else: else:
logger.debug(f"Starter message for {thread.name} already in JSON") forum_data["messages"].append(msg_data)
forum_data["messageCount"] = len(forum_data["messages"])
# Recalculate forum totalAttachmentSizeBytes from enriched starter messages
forum_data["totalAttachmentSizeBytes"] = sum(
m.get("totalFileSizeBytes", 0) for m in forum_data["messages"]
)
# Recalculate forum numberOfAttachments from enriched starter messages
forum_data["numberOfAttachments"] = sum(
m.get("numberOfFiles", 0) for m in forum_data["messages"]
)
# Keep chronological order
forum_data["messages"].sort(key=lambda x: x["timestamp"])
with open(forum_json_file, "w", encoding="utf-8") as f:
json.dump(forum_data, f, indent=4, ensure_ascii=False)
logger.info(f"Appended starter message for {thread.name} to {forum_json_file.name}")
else: else:
logger.warning(f"Forum JSON file does not exist: {forum_json_file}") logger.warning(f"Forum JSON file does not exist: {forum_json_file}")
@ -694,8 +735,4 @@ class DiscordExporter:
logger.warning(f"No starter message found for thread: {thread.name}") logger.warning(f"No starter message found for thread: {thread.name}")
except Exception as e: except Exception as e:
logger.error(f"Error adding starter message for {thread.name}: {e}") logger.error(f"Error adding starter message for {thread.name}: {e}")
return accumulated_count
await self.export_channel_messages(thread.id, progress_callback=progress_callback, force=force)
thread_count += 1
return thread_count