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

improve message state handling

Browse source
This commit is contained in:
rambros 2026-02-23 13:43:13 +05:30
parent 6600be74c3
commit 17d9094084
3 changed files with 140 additions and 88 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
.gitignore vendored
View file

@ -20,7 +20,7 @@ wheels/
*.egg-info/ *.egg-info/
.installed.cfg .installed.cfg
*.egg *.egg
temp.txt *.txt
# Virtual Environment # Virtual Environment
venv/ venv/
@ -32,8 +32,8 @@ config.yaml
!config.example.yaml !config.example.yaml
# Logs and State # Logs and State
migration.log *.log
state.json *.json
# Temporary Test Scripts # Temporary Test Scripts
test_*.py test_*.py

116
src/core/state.py
View file

@ -5,22 +5,28 @@ from typing import Dict, Any
class MigrationState: class MigrationState:
"""Manages persistence of the migration state to allow resumability.""" """Manages persistence of the migration state to allow resumability."""
def __init__(self, state_file: str | Path = "state.json"): def __init__(self, state_file: str | Path = "state.json", messages_file: str | Path = "messages.json"):
self.state_file = Path(state_file) self.state_file = Path(state_file)
self.messages_file = Path(messages_file)
# mappings: discord_id -> fluxer_id # mappings: discord_id -> fluxer_id
self.channel_map: Dict[str, str] = {} self.channel_map: Dict[str, str] = {}
self.category_map: Dict[str, str] = {} self.category_map: Dict[str, str] = {}
self.role_map: Dict[str, str] = {} self.role_map: Dict[str, str] = {}
self.emoji_map: Dict[str, str] = {} self.emoji_map: Dict[str, str] = {}
self.sticker_map: Dict[str, str] = {} self.sticker_map: Dict[str, str] = {}
self.message_map: Dict[str, str] = {}
# tracking last message timestamp per channel to resume # message tracking
self.message_map: Dict[str, str] = {}
self.last_message_timestamps: Dict[str, str] = {} self.last_message_timestamps: Dict[str, str] = {}
self.load() self.load()
def load(self): def load(self):
migrated_state = False
migrated_messages = False
# 1. Load primary state.json
if self.state_file.exists(): if self.state_file.exists():
with open(self.state_file, "r", encoding="utf-8") as f: with open(self.state_file, "r", encoding="utf-8") as f:
data = json.load(f) data = json.load(f)
@ -29,109 +35,133 @@ class MigrationState:
self.role_map = data.get("roles", {}) self.role_map = data.get("roles", {})
self.emoji_map = data.get("emojis", {}) self.emoji_map = data.get("emojis", {})
self.sticker_map = data.get("stickers", {}) self.sticker_map = data.get("stickers", {})
self.message_map = data.get("messages", {})
self.last_message_timestamps = data.get("last_message_timestamps", {})
# Check for legacy messages in state.json
if "messages" in data or "last_message_timestamps" in data:
self.message_map = data.get("messages", {})
self.last_message_timestamps = data.get("last_message_timestamps", {})
migrated_messages = True # We found legacy data, we should write it out to messages.json later
# Legacy Migration: Move role_, emoji_, sticker_ from channel_map to dedicated maps # Legacy Migration: Move role_, emoji_, sticker_ from channel_map to dedicated maps
migrated = False
legacy_keys = list(self.channel_map.keys()) legacy_keys = list(self.channel_map.keys())
for k in legacy_keys: for k in legacy_keys:
if k.startswith("role_"): if k.startswith("role_"):
discord_id = k.replace("role_", "") discord_id = k.replace("role_", "")
self.role_map[discord_id] = self.channel_map.pop(k) self.role_map[discord_id] = self.channel_map.pop(k)
migrated = True migrated_state = True
elif k.startswith("emoji_"): elif k.startswith("emoji_"):
discord_id = k.replace("emoji_", "") discord_id = k.replace("emoji_", "")
self.emoji_map[discord_id] = self.channel_map.pop(k) self.emoji_map[discord_id] = self.channel_map.pop(k)
migrated = True migrated_state = True
elif k.startswith("sticker_"): elif k.startswith("sticker_"):
discord_id = k.replace("sticker_", "") discord_id = k.replace("sticker_", "")
self.sticker_map[discord_id] = self.channel_map.pop(k) self.sticker_map[discord_id] = self.channel_map.pop(k)
migrated = True migrated_state = True
if migrated:
self.save()
def save(self): # 2. Load separate messages.json (overrides legacy state.json)
if self.messages_file.exists():
with open(self.messages_file, "r", encoding="utf-8") as f:
msg_data = json.load(f)
self.message_map = msg_data.get("messages", {})
self.last_message_timestamps = msg_data.get("last_message_timestamps", {})
migrated_messages = False # No need to force a migrating save since it already exists
# 3. Save if we migrated any legacy data to separate maps/files
if migrated_state or migrated_messages:
self.save_state()
if migrated_messages:
self.save_messages()
def save_state(self):
"""Saves only the core server configuration (channels, roles, emojis)."""
data = { data = {
"channels": self.channel_map, "channels": self.channel_map,
"categories": self.category_map, "categories": self.category_map,
"roles": self.role_map, "roles": self.role_map,
"emojis": self.emoji_map, "emojis": self.emoji_map,
"stickers": self.sticker_map, "stickers": self.sticker_map
"last_message_timestamps": self.last_message_timestamps,
"messages": self.message_map
} }
with open(self.state_file, "w", encoding="utf-8") as f: with open(self.state_file, "w", encoding="utf-8") as f:
json.dump(data, f, indent=4) json.dump(data, f, indent=4)
def save_messages(self):
"""Saves only the message tracking data."""
data = {
"last_message_timestamps": self.last_message_timestamps,
"messages": self.message_map
}
with open(self.messages_file, "w", encoding="utf-8") as f:
json.dump(data, f, indent=4)
# --- Type Specific Getters/Setters ---
def set_channel_mapping(self, discord_id: str, fluxer_id: str): def set_channel_mapping(self, discord_id: str, fluxer_id: str):
self.channel_map[str(discord_id)] = str(fluxer_id) self.channel_map[str(discord_id)] = str(fluxer_id)
self.save() self.save_state()
def get_fluxer_channel_id(self, discord_id: str) -> str | None: def get_fluxer_channel_id(self, discord_id: str) -> str | None:
return self.channel_map.get(str(discord_id)) return self.channel_map.get(str(discord_id))
def remove_channel_mapping(self, discord_id: str): def remove_channel_mapping(self, discord_id: str):
self.channel_map.pop(str(discord_id), None) self.channel_map.pop(str(discord_id), None)
self.save() self.save_state()
def set_category_mapping(self, discord_id: str, fluxer_id: str): def set_category_mapping(self, discord_id: str, fluxer_id: str):
self.category_map[str(discord_id)] = str(fluxer_id) self.category_map[str(discord_id)] = str(fluxer_id)
self.save() self.save_state()
def get_fluxer_category_id(self, discord_id: str) -> str | None: def get_fluxer_category_id(self, discord_id: str) -> str | None:
return self.category_map.get(str(discord_id)) return self.category_map.get(str(discord_id))
def remove_category_mapping(self, discord_id: str): def remove_category_mapping(self, discord_id: str):
self.category_map.pop(str(discord_id), None) self.category_map.pop(str(discord_id), None)
self.save() self.save_state()
def set_message_mapping(self, discord_id: str, fluxer_id: str):
self.message_map[str(discord_id)] = str(fluxer_id)
self.save()
def get_fluxer_message_id(self, discord_id: str) -> str | None:
return self.message_map.get(str(discord_id))
def update_last_message_timestamp(self, channel_id: str, timestamp: str):
self.last_message_timestamps[str(channel_id)] = timestamp
self.save()
# --- Type Specific Getters/Setters ---
def set_role_mapping(self, discord_id: str, fluxer_id: str): def set_role_mapping(self, discord_id: str, fluxer_id: str):
self.role_map[str(discord_id)] = str(fluxer_id) self.role_map[str(discord_id)] = str(fluxer_id)
self.save() self.save_state()
def get_fluxer_role_id(self, discord_id: str) -> str | None: def get_fluxer_role_id(self, discord_id: str) -> str | None:
return self.role_map.get(str(discord_id)) return self.role_map.get(str(discord_id))
def remove_role_mapping(self, discord_id: str): def remove_role_mapping(self, discord_id: str):
self.role_map.pop(str(discord_id), None) self.role_map.pop(str(discord_id), None)
self.save() self.save_state()
def set_emoji_mapping(self, discord_id: str, fluxer_id: str): def set_emoji_mapping(self, discord_id: str, fluxer_id: str):
self.emoji_map[str(discord_id)] = str(fluxer_id) self.emoji_map[str(discord_id)] = str(fluxer_id)
self.save() self.save_state()
def get_fluxer_emoji_id(self, discord_id: str) -> str | None: def get_fluxer_emoji_id(self, discord_id: str) -> str | None:
return self.emoji_map.get(str(discord_id)) return self.emoji_map.get(str(discord_id))
def remove_emoji_mapping(self, discord_id: str): def remove_emoji_mapping(self, discord_id: str):
self.emoji_map.pop(str(discord_id), None) self.emoji_map.pop(str(discord_id), None)
self.save() self.save_state()
def set_sticker_mapping(self, discord_id: str, fluxer_id: str): def set_sticker_mapping(self, discord_id: str, fluxer_id: str):
self.sticker_map[str(discord_id)] = str(fluxer_id) self.sticker_map[str(discord_id)] = str(fluxer_id)
self.save() self.save_state()
def get_fluxer_sticker_id(self, discord_id: str) -> str | None: def get_fluxer_sticker_id(self, discord_id: str) -> str | None:
return self.sticker_map.get(str(discord_id)) return self.sticker_map.get(str(discord_id))
def remove_sticker_mapping(self, discord_id: str): def remove_sticker_mapping(self, discord_id: str):
self.sticker_map.pop(str(discord_id), None) self.sticker_map.pop(str(discord_id), None)
self.save() self.save_state()
# --- Message Management ---
def set_message_mapping(self, discord_id: str, fluxer_id: str):
self.message_map[str(discord_id)] = str(fluxer_id)
self.save_messages()
def get_fluxer_message_id(self, discord_id: str) -> str | None:
return self.message_map.get(str(discord_id))
def update_last_message_timestamp(self, channel_id: str, timestamp: str):
self.last_message_timestamps[str(channel_id)] = timestamp
self.save_messages()
# --- Danger Zone Clearing --- # --- Danger Zone Clearing ---
@ -139,21 +169,21 @@ class MigrationState:
"""Clears all channel and category mappings.""" """Clears all channel and category mappings."""
self.channel_map.clear() self.channel_map.clear()
self.category_map.clear() self.category_map.clear()
self.save() self.save_state()
def clear_role_mappings(self): def clear_role_mappings(self):
"""Clears all role mappings.""" """Clears all role mappings."""
self.role_map.clear() self.role_map.clear()
self.save() self.save_state()
def clear_asset_mappings(self): def clear_asset_mappings(self):
"""Clears all emoji and sticker mappings.""" """Clears all emoji and sticker mappings."""
self.emoji_map.clear() self.emoji_map.clear()
self.sticker_map.clear() self.sticker_map.clear()
self.save() self.save_state()
def clear_message_history(self): def clear_message_history(self):
"""Clears all message mappings and timestamps.""" """Clears all message mappings and timestamps."""
self.message_map.clear() self.message_map.clear()
self.last_message_timestamps.clear() self.last_message_timestamps.clear()
self.save() self.save_messages()

106
src/ui/app.py
View file

@ -203,7 +203,7 @@ class MigrationCLI:
console.print("(Q) Exit") console.print("(Q) Exit")
choice = Prompt.ask("\nSelect an option", choices=["1", "2", "3", "4", "5", "6", "7", "Q", "q"], default="Q").upper() choice = Prompt.ask("\nSelect an option [1/2/3/4/5/6/7/Q]", choices=["1", "2", "3", "4", "5", "6", "7", "Q", "q"], default="Q", show_choices=False).upper()
if choice == "1": if choice == "1":
await self.clone_server_template() await self.clone_server_template()
@ -381,7 +381,7 @@ class MigrationCLI:
console.print("[bold red](F) Force re-clone, creates duplicate channels![/bold red]") console.print("[bold red](F) Force re-clone, creates duplicate channels![/bold red]")
console.print("[bold yellow](B) Back[/bold yellow]") console.print("[bold yellow](B) Back[/bold yellow]")
choice = Prompt.ask("Select an option", choices=["Y", "F", "B"], default="Y").upper() choice = Prompt.ask("Select an option [Y/F/B]", choices=["Y", "y", "F", "f", "B", "b"], default="Y", show_choices=False).upper()
if choice == "B": if choice == "B":
await self.engine.close_connections() await self.engine.close_connections()
@ -430,7 +430,7 @@ class MigrationCLI:
console.print("(2) Sync Category Permissions & Channel Permissions") console.print("(2) Sync Category Permissions & Channel Permissions")
console.print("(B) Back") console.print("(B) Back")
choice = Prompt.ask("Select an option", choices=["1", "2", "B", "b"], default="B").upper() choice = Prompt.ask("Select an option [1/2/B]", choices=["1", "2", "B", "b"], default="B", show_choices=False).upper()
if choice == "B": if choice == "B":
return return
@ -469,7 +469,7 @@ class MigrationCLI:
console.print("[bold red](F) Force Overwrite[/bold red]") console.print("[bold red](F) Force Overwrite[/bold red]")
console.print("[bold yellow](B) Back[/bold yellow]") console.print("[bold yellow](B) Back[/bold yellow]")
sub_choice = Prompt.ask("Select an option", choices=["Y", "F", "B"], default="Y").upper() sub_choice = Prompt.ask("Select an option [Y/F/B]", choices=["Y", "y", "F", "f", "B", "b"], default="Y", show_choices=False).upper()
if sub_choice == "B": if sub_choice == "B":
return return
@ -524,7 +524,7 @@ class MigrationCLI:
console.print("[bold green](Y) Proceed with Permission synchronization[/bold green]") console.print("[bold green](Y) Proceed with Permission synchronization[/bold green]")
console.print("[bold yellow](B) Back[/bold yellow]") console.print("[bold yellow](B) Back[/bold yellow]")
sub_choice = Prompt.ask("Select an option", choices=["Y", "B"], default="Y").upper() sub_choice = Prompt.ask("Select an option [Y/B]", choices=["Y", "y", "B", "b"], default="Y", show_choices=False).upper()
if sub_choice == "B": if sub_choice == "B":
return return
@ -598,7 +598,7 @@ class MigrationCLI:
console.print("(3) Sync Emojis and Stickers") console.print("(3) Sync Emojis and Stickers")
console.print("(B) Back") console.print("(B) Back")
choice = Prompt.ask("Select an option", choices=["1", "2", "3", "B", "b"], default="B").upper() choice = Prompt.ask("Select an option [1/2/3/B]", choices=["1", "2", "3", "B", "b"], default="B", show_choices=False).upper()
if choice == "B": if choice == "B":
return return
@ -625,7 +625,7 @@ class MigrationCLI:
console.print("[bold red](F) Force Overwrite[/bold red]") console.print("[bold red](F) Force Overwrite[/bold red]")
console.print("[bold yellow](B) Back[/bold yellow]") console.print("[bold yellow](B) Back[/bold yellow]")
choice = Prompt.ask("Select an option", choices=["Y", "F", "B"], default="Y").upper() choice = Prompt.ask("Select an option [Y/F/B]", choices=["Y", "y", "F", "f", "B", "b"], default="Y", show_choices=False).upper()
if choice == "B": if choice == "B":
return return
@ -692,7 +692,7 @@ class MigrationCLI:
console.print("(4) Sync Everything") console.print("(4) Sync Everything")
console.print("(B) Back") console.print("(B) Back")
choice = Prompt.ask("Select an option", choices=["1", "2", "3", "4", "B", "b"], default="B").upper() choice = Prompt.ask("Select an option [1/2/3/4/B]", choices=["1", "2", "3", "4", "B", "b"], default="B", show_choices=False).upper()
if choice == "B": if choice == "B":
return return
@ -721,7 +721,7 @@ class MigrationCLI:
await self.engine.close_connections() await self.engine.close_connections()
async def migrate_message_history(self): async def migrate_message_history(self):
console.print("\n[bold]Message History Migration (Manual Selection)[/bold]") console.print("\n[bold]Message History Migration[/bold]")
if not self.tokens_valid: if not self.tokens_valid:
console.print("[bold red]Error: You must have a valid configuration (Option 6) to migrate messages.[/bold red]") console.print("[bold red]Error: You must have a valid configuration (Option 6) to migrate messages.[/bold red]")
@ -837,34 +837,51 @@ class MigrationCLI:
if first_msg: if first_msg:
# Link format: https://discord.com/channels/GUILD_ID/CHANNEL_ID/MESSAGE_ID # Link format: https://discord.com/channels/GUILD_ID/CHANNEL_ID/MESSAGE_ID
first_msg_link = f"https://discord.com/channels/{self.config.discord_server_id}/{source_channel.id}/{first_msg.id}" first_msg_link = f"https://discord.com/channels/{self.config.discord_server_id}/{source_channel.id}/{first_msg.id}"
console.print(f"\n[bold green]Channel Found![/bold green]") server_name = self.validation_results.get("discord_server_name", "server")
console.print(f"Oldest Message Link: {first_msg_link}")
console.print(f"Author: [blue]{first_msg.author.name}[/blue]")
console.print(f"Content Preview: {first_msg.content[:100]}...")
console.print("\n(Y) [green]Yes, start from oldest message[/green]") while True:
console.print("(M) Provide Specific message link or ID") console.print(f"\n[bold green]Channel Found![/bold green]")
console.print("(B) Back") console.print(f"Oldest Message Link: {first_msg_link}")
console.print(f"Author: [blue]{first_msg.author.name}[/blue]")
start_mode = Prompt.ask("Start migration", choices=["Y", "M", "B"], default="Y").upper() console.print(f"Content Preview: {first_msg.content[:100]}...")
if start_mode == "B": console.print("\n(Y) [green]Yes, start from oldest message[/green]")
return console.print("(M) Provide Specific message link or ID")
elif start_mode == "M": console.print("(B) Back")
custom_link = Prompt.ask("Enter Discord Message Link")
try: start_mode = Prompt.ask("Start migration [Y/M/B]", choices=["Y", "y", "M", "m", "B", "b"], default="Y", show_choices=False).upper()
# Extract message ID from end of link
custom_id = int(custom_link.split("/")[-1]) if start_mode == "B":
# Check if message exists to give feedback return
msg = await self.engine.discord_reader.get_message(source_channel.id, custom_id) elif start_mode == "Y":
if msg: break
# To include this message, we start 'after' the ID before it elif start_mode == "M":
after_id = custom_id - 1 prompt_msg = "Enter Discord Message Link"
console.print(f"[green]Confirmed: Starting from {msg.author.name}'s message.[/green]") valid_message_found = False
else: while True:
console.print("[red]Warning: Message ID not found in this channel. Starting from beginning.[/red]") custom_link = Prompt.ask(prompt_msg)
except Exception as e: if str(custom_link).upper() in ["B", "BACK"]:
console.print(f"[red]Error parsing link: {e}. Starting from beginning.[/red]") break
try:
# Extract message ID from end of link
custom_id = int(str(custom_link).strip().split("/")[-1])
# Check if message exists to give feedback
msg = await self.engine.discord_reader.get_message(source_channel.id, custom_id)
if msg:
# To include this message, we start 'after' the ID before it
after_id = custom_id - 1
console.print(f"\n[green]Confirmed: Starting from {msg.author.name}'s message.[/green]")
valid_message_found = True
break
else:
console.print("\n[red]Error parsing link: 404 Not Found (error code: 10008): Unknown Message.[/red]")
prompt_msg = f"[yellow]Provide a valid Message link or Message ID from \"{server_name}\" (or 'B' to go back)[/yellow]"
except Exception as e:
console.print(f"\n[red]Error parsing link: {e}[/red]")
prompt_msg = f"[yellow]Provide a valid Message link or Message ID from \"{server_name}\" (or 'B' to go back)[/yellow]"
if valid_message_found:
break
else: else:
console.print("[yellow]Source channel appears to be empty. Nothing to migrate.[/yellow]") console.print("[yellow]Source channel appears to be empty. Nothing to migrate.[/yellow]")
return return
@ -968,9 +985,10 @@ class MigrationCLI:
console.print("(B) Back") console.print("(B) Back")
choice = Prompt.ask( choice = Prompt.ask(
"Select a Danger Zone option", "Select a Danger Zone option [1/2/3/4/B]",
choices=["1", "2", "3", "4", "B", "b"], choices=["1", "2", "3", "4", "B", "b"],
default="B" default="B",
show_choices=False
).upper() ).upper()
if choice == "B": if choice == "B":
@ -979,7 +997,8 @@ class MigrationCLI:
# ---- (1) Delete all Channels & Categories ---- # ---- (1) Delete all Channels & Categories ----
if choice == "1": if choice == "1":
console.print("") console.print("")
console.print("[bold red]This will DELETE every channel and category in the Fluxer community.[/bold red]") community_name = self.validation_results.get("fluxer_community_name", "Unknown")
console.print(f"[bold red]This will DELETE every channel and category in the Fluxer community: \"{community_name}\"[/bold red]")
if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"): if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"):
return return
if not Confirm.ask("[bold red]Last chance \u2013 this cannot be undone. Continue?[/bold red]"): if not Confirm.ask("[bold red]Last chance \u2013 this cannot be undone. Continue?[/bold red]"):
@ -1010,7 +1029,8 @@ class MigrationCLI:
# ---- (2) Reset Channel & Category Permissions ---- # ---- (2) Reset Channel & Category Permissions ----
elif choice == "2": elif choice == "2":
console.print("") console.print("")
console.print("[bold red]This will RESET all permission overwrites on every channel and category.[/bold red]") community_name = self.validation_results.get("fluxer_community_name", "Unknown")
console.print(f"[bold red]This will RESET all permission overwrites on every channel and category in the Fluxer community: \"{community_name}\"[/bold red]")
if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"): if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"):
return return
if not Confirm.ask("[bold red]Last chance \u2013 all custom permission overrides will be wiped. Continue?[/bold red]"): if not Confirm.ask("[bold red]Last chance \u2013 all custom permission overrides will be wiped. Continue?[/bold red]"):
@ -1041,7 +1061,8 @@ class MigrationCLI:
# ---- (3) Delete all Roles ---- # ---- (3) Delete all Roles ----
elif choice == "3": elif choice == "3":
console.print("") console.print("")
console.print("[bold red]This will DELETE all roles in the Fluxer community.[/bold red]") community_name = self.validation_results.get("fluxer_community_name", "Unknown")
console.print(f"[bold red]This will DELETE all roles in the Fluxer community: \"{community_name}\"[/bold red]")
console.print("[dim]Managed roles (including the bot's own role) and @everyone are automatically protected.[/dim]") console.print("[dim]Managed roles (including the bot's own role) and @everyone are automatically protected.[/dim]")
if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"): if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"):
return return
@ -1073,7 +1094,8 @@ class MigrationCLI:
# ---- (4) Delete all Emojis & Stickers ---- # ---- (4) Delete all Emojis & Stickers ----
elif choice == "4": elif choice == "4":
console.print("") console.print("")
console.print("[bold red]This will DELETE all custom emojis and stickers in the Fluxer community.[/bold red]") community_name = self.validation_results.get("fluxer_community_name", "Unknown")
console.print(f"[bold red]This will DELETE all custom emojis and stickers in the Fluxer community: \"{community_name}\"[/bold red]")
if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"): if not Confirm.ask("[bold red]Are you absolutely sure?[/bold red]"):
return return
if not Confirm.ask("[bold red]Last chance \u2013 this cannot be undone. Continue?[/bold red]"): if not Confirm.ask("[bold red]Last chance \u2013 this cannot be undone. Continue?[/bold red]"):