diff --git a/xp/export.py b/xp/export.py new file mode 100644 index 0000000..fe54252 --- /dev/null +++ b/xp/export.py @@ -0,0 +1,88 @@ +import os +import json +import asyncio +from io import BytesIO +import discord +from discord import app_commands +from discord.ext import commands +from xp.database import get_db + +class ExportXP(commands.Cog): + def __init__(self, bot: commands.Bot): + self.bot = bot + + async def _export_data(self, lifetime: bool): + """Async function to handle database operations in a thread""" + def _db_work(): + try: + conn, cur = get_db(lifetime) + # Fetch all data at once as a list of dicts for faster processing + cur.execute("SELECT user_id, xp, level, last_message FROM xp") + rows = cur.fetchall() + conn.close() + + # Build dict directly in the thread + users = {} + for user_id, xp, level, last_msg in rows: + users[str(user_id)] = { + "xp": xp, + "level": level, + "last_message": last_msg + } + return users + except Exception as e: + print(f"Error in _db_work: {e}") + raise + + # Run everything in thread pool + users = await asyncio.to_thread(_db_work) + return {"users": users} + + @app_commands.command(name="export_xp", description="Export XP data to JSON (lifetime or annual)") + @app_commands.choices( + xp_type=[ + app_commands.Choice(name="Lifetime", value="lifetime"), + app_commands.Choice(name="Annual", value="annual") + ] + ) + async def export_xp(self, interaction: discord.Interaction, xp_type: app_commands.Choice[str]): + try: + # Defer immediately (not ephemeral so everyone can see) + await interaction.response.defer() + print(f"Export started for {xp_type.value}") + + lifetime = xp_type.value == "lifetime" + + # Get data asynchronously + print("Fetching data from database...") + data = await self._export_data(lifetime) + print(f"Data fetched: {len(data['users'])} users") + + # Create file with pretty-printed JSON + print("Encoding JSON...") + json_str = json.dumps(data, indent=2) + json_bytes = json_str.encode("utf-8") + print(f"JSON size: {len(json_bytes)} bytes") + + print("Creating Discord file...") + file = discord.File(fp=BytesIO(json_bytes), filename=f"{xp_type.value}_xp_export.json") + + # Public response + print("Sending response...") + await interaction.followup.send( + f"✅ Exported `{xp_type.value}` XP data ({len(data['users'])} users).", + file=file + ) + print("Export complete!") + + except Exception as e: + print(f"Error in export_xp: {e}") + import traceback + traceback.print_exc() + try: + await interaction.followup.send(f"❌ An error occurred during export: {str(e)}") + except: + pass + +async def setup(bot: commands.Bot): + await bot.add_cog(ExportXP(bot)) \ No newline at end of file diff --git a/xp/import.py b/xp/import.py new file mode 100644 index 0000000..0c2bf8f --- /dev/null +++ b/xp/import.py @@ -0,0 +1,123 @@ +import json +import asyncio +from io import BytesIO +import discord +from discord import app_commands +from discord.ext import commands +from xp.utils import xp_for_level +from xp.database import get_db +from moderation.loader import ModerationBase + +class ImportXP(commands.Cog): + def __init__(self, bot: commands.Bot): + self.bot = bot + + async def _import_data(self, users_data: dict, lifetime: bool): + """Async function to handle database operations in a thread""" + def _db_work(): + try: + conn, cur = get_db(lifetime) + + # Drop and recreate the XP table + cur.execute("DROP TABLE IF EXISTS xp") + cur.execute(""" + CREATE TABLE xp ( + user_id TEXT PRIMARY KEY, + xp INTEGER DEFAULT 0, + level INTEGER DEFAULT 0, + last_message INTEGER DEFAULT 0 + ) + """) + + # Batch insert for much better performance + insert_data = [] + for user_id, user_info in users_data.items(): + uid = str(user_id) + xp = int(user_info.get("xp", 0)) + level = 0 + while xp >= xp_for_level(level + 1): + level += 1 + insert_data.append((uid, xp, level, 0)) + + # Single executemany call is MUCH faster than individual inserts + cur.executemany( + "INSERT INTO xp (user_id, xp, level, last_message) VALUES (?, ?, ?, ?)", + insert_data + ) + + conn.commit() + conn.close() + return len(insert_data) + except Exception as e: + print(f"Error in _db_work: {e}") + raise + + # Run database work in thread pool + count = await asyncio.to_thread(_db_work) + return count + + @app_commands.command(name="import_xp", description="Import XP data from JSON (admin only, overwrites DB)") + @app_commands.choices( + xp_type=[ + app_commands.Choice(name="Lifetime", value="lifetime"), + app_commands.Choice(name="Annual", value="annual") + ] + ) + @ModerationBase.is_admin() + async def import_xp( + self, + interaction: discord.Interaction, + xp_type: app_commands.Choice[str], + attachment: discord.Attachment + ): + try: + # Defer immediately (not ephemeral so everyone can see) + await interaction.response.defer() + print(f"Import started for {xp_type.value}") + + lifetime = xp_type.value == "lifetime" + + if not attachment.filename.endswith(".json"): + await interaction.followup.send("❌ Please upload a valid `.json` file.") + return + + # Download and parse JSON file + print("Downloading attachment...") + file_bytes = await attachment.read() + print(f"Downloaded {len(file_bytes)} bytes") + + try: + print("Parsing JSON...") + data = json.loads(file_bytes.decode("utf-8")) + except json.JSONDecodeError as e: + await interaction.followup.send(f"❌ Invalid JSON file format: {str(e)}") + return + + users_data = data.get("users", {}) + + if not users_data: + await interaction.followup.send("❌ No user data found in the JSON file.") + return + + print(f"Found {len(users_data)} users to import") + + # Import data asynchronously + print("Importing data...") + count = await self._import_data(users_data, lifetime) + print(f"Import complete: {count} users") + + await interaction.followup.send( + f"✅ Imported `{xp_type.value}` XP data from `{attachment.filename}` — {count} users imported (existing data overwritten)." + ) + + except Exception as e: + print(f"Error in import_xp: {e}") + import traceback + traceback.print_exc() + try: + await interaction.followup.send(f"❌ An error occurred during import: {str(e)}") + except: + pass + +async def setup(bot: commands.Bot): + await bot.add_cog(ImportXP(bot)) \ No newline at end of file