From 5750133b3a49c32ecd23106f6100853fcd2f9db6 Mon Sep 17 00:00:00 2001 From: rambros Date: Sat, 28 Mar 2026 02:46:19 +0530 Subject: [PATCH] use int for discord snowflake ids --- src/core/backup_database.py | 227 +++++++++++++++++++++--------------- src/core/database.py | 147 ++++++++++++++++------- src/core/exporter.py | 5 +- src/core/utils.py | 14 +++ 4 files changed, 255 insertions(+), 138 deletions(-) diff --git a/src/core/backup_database.py b/src/core/backup_database.py index 733612c..88d0df1 100644 --- a/src/core/backup_database.py +++ b/src/core/backup_database.py @@ -4,20 +4,11 @@ import json import threading from pathlib import Path from typing import Dict, Any, List, Optional, Union +from src.core.utils import parse_snowflake logger = logging.getLogger(__name__) -def parse_snowflake(value: Any) -> Optional[int]: - """Safely parses a Discord ID (Snowflake) from any input, handling 'None' strings.""" - if value is None: - return None - s = str(value).strip() - if not s or s.lower() == "none" or s == "NULL": - return None - try: - return int(s) - except ValueError: - return None + class BackupDatabase: """Manages the SQLite database for local Discord backups.""" @@ -39,47 +30,101 @@ class BackupDatabase: def _migrate_db(self): """Handles backward compatibility by renaming columns in existing databases.""" with self._lock: - # Check 'media_pool' table - self._conn.execute("") - res = self._conn.execute("SELECT count(*) FROM sqlite_master WHERE type='table' AND name='media_pool'").fetchone() - if res[0] > 0: - cols = self._conn.execute("PRAGMA table_info(media_pool)").fetchall() - col_names = [c["name"] for c in cols] - if "mime_type" in col_names and "content_type" not in col_names: - logger.info("Migrating media_pool: renaming 'mime_type' to 'content_type'") - self._conn.execute("ALTER TABLE media_pool RENAME COLUMN mime_type TO content_type") + conn = self._conn + # 1. MIME Type to Content Type Migrations + for table in ["media_pool", "server_assets"]: + res = conn.execute(f"SELECT count(*) FROM sqlite_master WHERE type='table' AND name='{table}'").fetchone() + if res and res[0] > 0: + cols = conn.execute(f"PRAGMA table_info({table})").fetchall() + col_names = [c["name"] for c in cols] + if "mime_type" in col_names and "content_type" not in col_names: + logger.info(f"Migrating {table}: renaming 'mime_type' to 'content_type'") + conn.execute(f"ALTER TABLE {table} RENAME COLUMN mime_type TO content_type") - # Check 'server_assets' table - res = self._conn.execute("SELECT count(*) FROM sqlite_master WHERE type='table' AND name='server_assets'").fetchone() - if res[0] > 0: - cols = self._conn.execute("PRAGMA table_info(server_assets)").fetchall() - col_names = [c["name"] for c in cols] - if "mime_type" in col_names and "content_type" not in col_names: - logger.info("Migrating server_assets: renaming 'mime_type' to 'content_type'") - self._conn.execute("ALTER TABLE server_assets RENAME COLUMN mime_type TO content_type") - res = self._conn.execute("SELECT count(*) FROM messages LIMIT 1").fetchone() - if res[0] > 0: - cols = self._conn.execute("PRAGMA table_info(messages)").fetchall() - id_type = next(col for col in cols if col[1] == "id")[2] - if id_type == "TEXT": - logger.info("Migrating messages: Changing id column type to integer") - self._conn.execute("ALTER TABLE messages RENAME TO messages_old") - self._conn.execute("""CREATE TABLE IF NOT EXISTS messages - ( - id INTEGER PRIMARY KEY, - channel_id TEXT, - author_id TEXT, - content TEXT, - timestamp TEXT, - type INTEGER, - message_reference TEXT, - is_pinned INTEGER, - extra_data TEXT - ) - """) - self._conn.execute("INSERT INTO messages SELECT * FROM messages_old") - self._conn.execute("DROP TABLE messages_old") - self._conn.commit() + # 2. Universal ID Migration (TEXT -> INTEGER) + # Mapping of table names to columns that must be INTEGER (Snowflakes) + id_migrations = { + "guild_profile": ["id", "owner_id"], + "roles": ["id", "permissions"], + "channels": ["id", "category_id"], + "permissions": ["channel_id", "target_id"], + "users": ["id"], + "messages": ["id", "channel_id", "author_id", "message_reference"], + "attachments": ["id", "message_id"], + "embeds": ["message_id"], + "reactions": ["message_id", "emoji_id"], + "message_stickers": ["message_id", "sticker_id"], + "threads": ["id", "parent_id"], + "forum_tags": ["id", "forum_id", "emoji_id"], + "server_assets": ["id"] + } + + for table, id_cols in id_migrations.items(): + res = conn.execute(f"SELECT count(*) FROM sqlite_master WHERE type='table' AND name='{table}'").fetchone() + if not res or res[0] == 0: + continue + + cols = conn.execute(f"PRAGMA table_info({table})").fetchall() + needs_migration = False + for col in cols: + if col[1] in id_cols and col[2] == "TEXT": + needs_migration = True + break + + if needs_migration: + logger.info(f"Migrating {table}: converting ID columns to INTEGER") + # Special Case: messages already handled id, but now generic + # We use a temporary table to handle the schema change + conn.execute(f"ALTER TABLE {table} RENAME TO {table}_old") + + # We can't easily generate the CREATE TABLE here without duplicating _init_db logic + # So we call _init_db to create the NEW table, then copy data + # But _init_db has 'IF NOT EXISTS', so we just call it once at the end? + # No, we need the table NOW for the INSERT. + # I'll just manually define the inserts or do it in _init_db. + + # Actually, a better way is to do the CREATE TABLE here for this specific table. + # I'll have to duplicate the schema from _init_db for the migration. + + # Alternatively, since we are already in _migrate_db, we can just do the + # specific CREATE TABLE for the table we are migrating. + + if table == "guild_profile": + conn.execute("CREATE TABLE guild_profile (id INTEGER PRIMARY KEY, name TEXT, description TEXT, icon_file TEXT, icon_url TEXT, banner_file TEXT, banner_url TEXT, owner_id INTEGER, last_backup TEXT, ignore_channels TEXT)") + elif table == "roles": + conn.execute("CREATE TABLE roles (id INTEGER PRIMARY KEY, name TEXT, color INTEGER, position INTEGER, permissions INTEGER, hoist INTEGER, mentionable INTEGER)") + elif table == "channels": + conn.execute("CREATE TABLE channels (id INTEGER PRIMARY KEY, name TEXT, type INTEGER, position INTEGER, category_id INTEGER, topic TEXT, nsfw INTEGER, bitrate INTEGER, slowmode_delay INTEGER)") + elif table == "permissions": + conn.execute("CREATE TABLE permissions (id INTEGER PRIMARY KEY AUTOINCREMENT, channel_id INTEGER, target_id INTEGER, target_type TEXT, allow INTEGER, deny INTEGER)") + elif table == "users": + conn.execute("CREATE TABLE users (id INTEGER PRIMARY KEY, username TEXT, display_name TEXT, avatar_file TEXT, avatar_url TEXT, roles TEXT)") + elif table == "messages": + conn.execute("CREATE TABLE messages (id INTEGER PRIMARY KEY, channel_id INTEGER, author_id INTEGER, content TEXT, timestamp TEXT, type INTEGER, message_reference INTEGER, is_pinned INTEGER, extra_data TEXT)") + elif table == "attachments": + conn.execute("CREATE TABLE attachments (id INTEGER PRIMARY KEY, message_id INTEGER, filename TEXT, size INTEGER, url TEXT, content_type TEXT, local_hash TEXT)") + elif table == "embeds": + conn.execute("CREATE TABLE embeds (id INTEGER PRIMARY KEY AUTOINCREMENT, message_id INTEGER, title TEXT, description TEXT, url TEXT, color INTEGER, timestamp TEXT, thumbnail_url TEXT, image_url TEXT, author_name TEXT, author_url TEXT, author_icon_url TEXT, footer_text TEXT, footer_icon_url TEXT, fields TEXT)") + elif table == "reactions": + conn.execute("CREATE TABLE reactions (id INTEGER PRIMARY KEY AUTOINCREMENT, message_id INTEGER, emoji_id INTEGER, emoji_name TEXT, count INTEGER)") + elif table == "message_stickers": + conn.execute("CREATE TABLE message_stickers (message_id INTEGER, sticker_id INTEGER, name TEXT, url TEXT, format_type INTEGER, local_hash TEXT, PRIMARY KEY (message_id, sticker_id))") + elif table == "threads": + conn.execute("CREATE TABLE threads (id INTEGER PRIMARY KEY, name TEXT, type INTEGER, parent_id INTEGER, message_count INTEGER, member_count INTEGER, archived INTEGER, archive_timestamp TEXT, auto_archive_duration INTEGER, locked INTEGER, applied_tags TEXT)") + elif table == "forum_tags": + conn.execute("CREATE TABLE forum_tags (id INTEGER PRIMARY KEY, forum_id INTEGER, name TEXT, moderated INTEGER, emoji_id INTEGER, emoji_name TEXT)") + elif table == "server_assets": + conn.execute("CREATE TABLE server_assets (id INTEGER PRIMARY KEY, name TEXT, type TEXT, filename TEXT, url TEXT, content_type INTEGER)") + + old_cols = [c[1] for c in conn.execute(f"PRAGMA table_info({table}_old)").fetchall()] + new_cols = [c[1] for c in conn.execute(f"PRAGMA table_info({table})").fetchall()] + common_cols = [c for c in old_cols if c in new_cols] + col_str = ", ".join(common_cols) + + conn.execute(f"INSERT INTO {table} ({col_str}) SELECT {col_str} FROM {table}_old") + conn.execute(f"DROP TABLE {table}_old") + + conn.commit() def _init_db(self): """Initializes the database schema.""" @@ -89,14 +134,14 @@ class BackupDatabase: # Guild Profile conn.execute(""" CREATE TABLE IF NOT EXISTS guild_profile ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, name TEXT, description TEXT, icon_file TEXT, icon_url TEXT, banner_file TEXT, banner_url TEXT, - owner_id TEXT, + owner_id INTEGER, last_backup TEXT, ignore_channels TEXT ) @@ -105,11 +150,11 @@ class BackupDatabase: # Roles conn.execute(""" CREATE TABLE IF NOT EXISTS roles ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, name TEXT, color INTEGER, position INTEGER, - permissions TEXT, + permissions INTEGER, hoist INTEGER, mentionable INTEGER ) @@ -118,11 +163,11 @@ class BackupDatabase: # Channels conn.execute(""" CREATE TABLE IF NOT EXISTS channels ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, name TEXT, type INTEGER, position INTEGER, - category_id TEXT, + category_id INTEGER, topic TEXT, nsfw INTEGER, bitrate INTEGER, @@ -134,8 +179,8 @@ class BackupDatabase: conn.execute(""" CREATE TABLE IF NOT EXISTS permissions ( id INTEGER PRIMARY KEY AUTOINCREMENT, - channel_id TEXT, - target_id TEXT, + channel_id INTEGER, + target_id INTEGER, target_type TEXT, allow INTEGER, deny INTEGER @@ -146,7 +191,7 @@ class BackupDatabase: # Users (Author cache) conn.execute(""" CREATE TABLE IF NOT EXISTS users ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, username TEXT, display_name TEXT, avatar_file TEXT, @@ -159,12 +204,12 @@ class BackupDatabase: conn.execute(""" CREATE TABLE IF NOT EXISTS messages ( id INTEGER PRIMARY KEY, - channel_id TEXT, - author_id TEXT, + channel_id INTEGER, + author_id INTEGER, content TEXT, timestamp TEXT, type INTEGER, - message_reference TEXT, + message_reference INTEGER, is_pinned INTEGER, extra_data TEXT ) @@ -175,8 +220,8 @@ class BackupDatabase: # Attachments conn.execute(""" CREATE TABLE IF NOT EXISTS attachments ( - id TEXT PRIMARY KEY, - message_id TEXT, + id INTEGER PRIMARY KEY, + message_id INTEGER, filename TEXT, size INTEGER, url TEXT, @@ -190,7 +235,7 @@ class BackupDatabase: conn.execute(""" CREATE TABLE IF NOT EXISTS embeds ( id INTEGER PRIMARY KEY AUTOINCREMENT, - message_id TEXT, + message_id INTEGER, title TEXT, description TEXT, url TEXT, @@ -212,8 +257,8 @@ class BackupDatabase: conn.execute(""" CREATE TABLE IF NOT EXISTS reactions ( id INTEGER PRIMARY KEY AUTOINCREMENT, - message_id TEXT, - emoji_id TEXT, + message_id INTEGER, + emoji_id INTEGER, emoji_name TEXT, count INTEGER ) @@ -223,8 +268,8 @@ class BackupDatabase: # Message Stickers conn.execute(""" CREATE TABLE IF NOT EXISTS message_stickers ( - message_id TEXT, - sticker_id TEXT, + message_id INTEGER, + sticker_id INTEGER, name TEXT, url TEXT, format_type INTEGER, @@ -237,10 +282,10 @@ class BackupDatabase: # Threads conn.execute(""" CREATE TABLE IF NOT EXISTS threads ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, name TEXT, type INTEGER, - parent_id TEXT, + parent_id INTEGER, message_count INTEGER, member_count INTEGER, archived INTEGER, @@ -255,11 +300,11 @@ class BackupDatabase: # Forum Tags (Definitions for a forum channel) conn.execute(""" CREATE TABLE IF NOT EXISTS forum_tags ( - id TEXT PRIMARY KEY, - forum_id TEXT, + id INTEGER PRIMARY KEY, + forum_id INTEGER, name TEXT, moderated INTEGER, - emoji_id TEXT, + emoji_id INTEGER, emoji_name TEXT ) """) @@ -280,7 +325,7 @@ class BackupDatabase: # Server Assets (Emojis, Stickers, etc.) conn.execute(""" CREATE TABLE IF NOT EXISTS server_assets ( - id TEXT PRIMARY KEY, + id INTEGER PRIMARY KEY, name TEXT, type TEXT, filename TEXT, @@ -299,10 +344,10 @@ class BackupDatabase: INSERT OR REPLACE INTO guild_profile (id, name, description, icon_file, icon_url, banner_file, banner_url, owner_id, last_backup, ignore_channels) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( - str(data.get("id")), data.get("name"), data.get("description"), + parse_snowflake(data.get("id")), data.get("name"), data.get("description"), data.get("icon_file"), data.get("icon_url"), data.get("banner_file"), data.get("banner_url"), - str(data.get("owner_id")), + parse_snowflake(data.get("owner_id")), data.get("last_backup"), json.dumps(data.get("ignore_channels", [])) )) self._conn.commit() @@ -323,7 +368,7 @@ class BackupDatabase: with self._lock: formatted = [ { - "id": str(r["id"]), + "id": parse_snowflake(r["id"]), "name": r["name"], "color": r["color"], "position": r["position"], @@ -370,7 +415,7 @@ class BackupDatabase: with self._lock: formatted = [ { - "id": str(a["id"]), + "id": parse_snowflake(a["id"]), "name": a.get("name"), "type": a.get("type"), "filename": a.get("filename"), @@ -430,7 +475,7 @@ class BackupDatabase: for rea in msg["reactions"]: all_reactions.append({ "message_id": msg["id"], - "emoji_id": str(rea["emoji_id"]) if rea.get("emoji_id") else None, + "emoji_id": parse_snowflake(rea["emoji_id"]) if rea.get("emoji_id") else None, "emoji_name": rea.get("emoji_name"), "count": rea.get("count", 0) }) @@ -440,7 +485,7 @@ class BackupDatabase: for st in msg["stickers"]: all_stickers.append({ "message_id": msg["id"], - "sticker_id": str(st["id"]), + "sticker_id": parse_snowflake(st["id"]), "name": st.get("name"), "url": st.get("url"), "format_type": st.get("format_type"), @@ -492,7 +537,7 @@ class BackupDatabase: def get_last_message_id(self, channel_id: str) -> Optional[str]: with self._lock: - row = self._conn.execute("SELECT id FROM messages WHERE channel_id = ? ORDER BY id DESC LIMIT 1", (str(channel_id),)).fetchone() + row = self._conn.execute("SELECT id FROM messages WHERE channel_id = ? ORDER BY id DESC LIMIT 1", (parse_snowflake(channel_id),)).fetchone() return row["id"] if row else None def get_media_by_hash(self, file_hash: str) -> Optional[Dict[str, Any]]: @@ -623,7 +668,7 @@ class BackupDatabase: """Returns forum tag definitions.""" with self._lock: if forum_id: - rows = self._conn.execute("SELECT * FROM forum_tags WHERE forum_id = ?", (str(forum_id),)).fetchall() + rows = self._conn.execute("SELECT * FROM forum_tags WHERE forum_id = ?", (parse_snowflake(forum_id),)).fetchall() else: rows = self._conn.execute("SELECT * FROM forum_tags").fetchall() return [dict(r) for r in rows] @@ -631,13 +676,13 @@ class BackupDatabase: def get_threads_by_parent(self, parent_id: str) -> List[Dict[str, Any]]: """Returns all threads belonging to a parent channel.""" with self._lock: - rows = self._conn.execute("SELECT * FROM threads WHERE parent_id = ?", (str(parent_id),)).fetchall() + rows = self._conn.execute("SELECT * FROM threads WHERE parent_id = ?", (parse_snowflake(parent_id),)).fetchall() return [dict(r) for r in rows] def get_thread(self, thread_id: str) -> Optional[Dict[str, Any]]: """Retrieves a single thread's metadata.""" with self._lock: - row = self._conn.execute("SELECT * FROM threads WHERE id = ?", (str(thread_id),)).fetchone() + row = self._conn.execute("SELECT * FROM threads WHERE id = ?", (parse_snowflake(thread_id),)).fetchone() return dict(row) if row else None def get_all_users(self) -> List[Dict[str, Any]]: @@ -647,7 +692,7 @@ class BackupDatabase: def get_user(self, user_id: str) -> Optional[Dict[str, Any]]: with self._lock: - row = self._conn.execute("SELECT * FROM users WHERE id = ?", (str(user_id),)).fetchone() + row = self._conn.execute("SELECT * FROM users WHERE id = ?", (parse_snowflake(user_id),)).fetchone() if row: data = dict(row) if data.get("roles"): @@ -673,11 +718,11 @@ class BackupDatabase: def get_messages_paged(self, channel_id: str, limit: int = 100, offset: int = 0, after_id: Optional[str] = None) -> List[Dict[str, Any]]: with self._lock: query = "SELECT * FROM messages WHERE channel_id = ?" - params = [str(channel_id)] + params = [parse_snowflake(channel_id)] if after_id: query += " AND id > ?" - params.append(str(after_id)) + params.append(parse_snowflake(after_id)) query += " ORDER BY id ASC LIMIT ? OFFSET ?" params.extend([limit, offset]) @@ -748,7 +793,7 @@ class BackupDatabase: def delete_channel_messages(self, channel_id: Union[str, int]): """Deletes all messages and related metadata for a specific channel and its threads.""" - cid = str(channel_id) + cid = parse_snowflake(channel_id) with self._lock: # 1. Identify all channel IDs involved (parent + all threads) target_ids = [cid] @@ -830,4 +875,4 @@ class BackupDatabase: self._conn.commit() self._conn.close() except Exception: - pass + pass \ No newline at end of file diff --git a/src/core/database.py b/src/core/database.py index b2230b9..aeffa49 100644 --- a/src/core/database.py +++ b/src/core/database.py @@ -6,6 +6,7 @@ from pathlib import Path from typing import Optional, Dict, Any import threading import sys +from src.core.utils import parse_snowflake logger = logging.getLogger(__name__) @@ -27,16 +28,72 @@ class MigrationDatabase: return self._local.conn def _init_db(self): - """Initialize tables if they don't exist.""" + """Initialize tables if they don't exist and handle migrations.""" conn = sqlite3.connect(self.db_path) cursor = conn.cursor() + # 1. MIME Type to Content Type Migrations (if applicable - not in this class usually) + + # 2. Universal ID Migration (TEXT -> INTEGER) + # Mapping of table names to columns that must be INTEGER (Snowflakes) + id_migrations = { + "message_mappings": ["channel_id", "source_msg_id", "target_msg_id"], + "thread_mappings": ["channel_id", "thread_id", "source_msg_id", "target_msg_id"], + "channel_tracking": ["channel_id", "last_msg_id"], + "thread_tracking": ["channel_id", "thread_id", "last_msg_id"], + "server_mappings": ["source_id", "target_id"], + "asset_mappings": ["source_id", "target_id"], + "user_alias": ["user_id"] + } + + for table, id_cols in id_migrations.items(): + cursor.execute(f"SELECT count(*) FROM sqlite_master WHERE type='table' AND name='{table}'") + res = cursor.fetchone() + if not res or res[0] == 0: + continue + + cursor.execute(f"PRAGMA table_info({table})") + cols = cursor.fetchall() + needs_migration = False + for col in cols: + if col[1] in id_cols and col[2] == "TEXT": + needs_migration = True + break + + if needs_migration: + logger.info(f"MigrationDatabase: Migrating {table}: converting ID columns to INTEGER") + cursor.execute(f"ALTER TABLE {table} RENAME TO {table}_old") + + if table == "message_mappings": + cursor.execute("CREATE TABLE message_mappings (channel_id INTEGER, source_msg_id INTEGER, target_msg_id INTEGER, timestamp TEXT, PRIMARY KEY (channel_id, source_msg_id))") + elif table == "thread_mappings": + cursor.execute("CREATE TABLE thread_mappings (channel_id INTEGER, thread_id INTEGER, source_msg_id INTEGER, target_msg_id INTEGER, timestamp TEXT, PRIMARY KEY (channel_id, thread_id, source_msg_id))") + elif table == "channel_tracking": + cursor.execute("CREATE TABLE channel_tracking (channel_id INTEGER PRIMARY KEY, last_msg_id INTEGER, last_msg_ts TEXT, msg_count INTEGER DEFAULT 0, file_count INTEGER DEFAULT 0)") + elif table == "thread_tracking": + cursor.execute("CREATE TABLE thread_tracking (channel_id INTEGER, thread_id INTEGER, last_msg_id INTEGER, last_msg_ts TEXT, msg_count INTEGER DEFAULT 0, file_count INTEGER DEFAULT 0, completed INTEGER DEFAULT 0, PRIMARY KEY (channel_id, thread_id))") + elif table == "server_mappings": + cursor.execute("CREATE TABLE server_mappings (category TEXT, source_id INTEGER, target_id INTEGER, PRIMARY KEY (category, source_id))") + elif table == "asset_mappings": + cursor.execute("CREATE TABLE asset_mappings (category TEXT, source_id INTEGER, target_id INTEGER, PRIMARY KEY (category, source_id))") + elif table == "user_alias": + cursor.execute("CREATE TABLE user_alias (user_id INTEGER PRIMARY KEY, alias TEXT UNIQUE)") + + old_cols = [c[1] for c in cursor.execute(f"PRAGMA table_info({table}_old)").fetchall()] + new_cols = [c[1] for c in cursor.execute(f"PRAGMA table_info({table})").fetchall()] + common_cols = [c for c in old_cols if c in new_cols] + col_str = ", ".join(common_cols) + + cursor.execute(f"INSERT OR IGNORE INTO {table} ({col_str}) SELECT {col_str} FROM {table}_old") + cursor.execute(f"DROP TABLE {table}_old") + + # Initial Creation / Ensure Schema Correctness # Table for message mappings: SourceID -> TargetID cursor.execute(""" CREATE TABLE IF NOT EXISTS message_mappings ( - channel_id TEXT, - source_msg_id TEXT, - target_msg_id TEXT, + channel_id INTEGER, + source_msg_id INTEGER, + target_msg_id INTEGER, timestamp TEXT, PRIMARY KEY (channel_id, source_msg_id) ) @@ -45,10 +102,10 @@ class MigrationDatabase: # Table for thread mappings cursor.execute(""" CREATE TABLE IF NOT EXISTS thread_mappings ( - channel_id TEXT, - thread_id TEXT, - source_msg_id TEXT, - target_msg_id TEXT, + channel_id INTEGER, + thread_id INTEGER, + source_msg_id INTEGER, + target_msg_id INTEGER, timestamp TEXT, PRIMARY KEY (channel_id, thread_id, source_msg_id) ) @@ -57,8 +114,8 @@ class MigrationDatabase: # Table for per-channel stats and tracking cursor.execute(""" CREATE TABLE IF NOT EXISTS channel_tracking ( - channel_id TEXT PRIMARY KEY, - last_msg_id TEXT, + channel_id INTEGER PRIMARY KEY, + last_msg_id INTEGER, last_msg_ts TEXT, msg_count INTEGER DEFAULT 0, file_count INTEGER DEFAULT 0 @@ -68,9 +125,9 @@ class MigrationDatabase: # Table for per-thread stats cursor.execute(""" CREATE TABLE IF NOT EXISTS thread_tracking ( - channel_id TEXT, - thread_id TEXT, - last_msg_id TEXT, + channel_id INTEGER, + thread_id INTEGER, + last_msg_id INTEGER, last_msg_ts TEXT, msg_count INTEGER DEFAULT 0, file_count INTEGER DEFAULT 0, @@ -89,8 +146,8 @@ class MigrationDatabase: cursor.execute(""" CREATE TABLE IF NOT EXISTS server_mappings ( category TEXT, - source_id TEXT, - target_id TEXT, + source_id INTEGER, + target_id INTEGER, PRIMARY KEY (category, source_id) ) """) @@ -99,8 +156,8 @@ class MigrationDatabase: cursor.execute(""" CREATE TABLE IF NOT EXISTS asset_mappings ( category TEXT, - source_id TEXT, - target_id TEXT, + source_id INTEGER, + target_id INTEGER, PRIMARY KEY (category, source_id) ) """) @@ -136,7 +193,7 @@ class MigrationDatabase: # Table for auto-generated user aliases (user_id -> alias) cursor.execute(""" CREATE TABLE IF NOT EXISTS user_alias ( - user_id TEXT PRIMARY KEY, + user_id INTEGER PRIMARY KEY, alias TEXT UNIQUE ) """) @@ -152,7 +209,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "INSERT OR REPLACE INTO message_mappings (channel_id, source_msg_id, target_msg_id, timestamp) VALUES (?, ?, ?, ?)", - (channel_id, source_id, target_id, timestamp) + (parse_snowflake(channel_id), parse_snowflake(source_id), parse_snowflake(target_id), timestamp) ) conn.commit() @@ -160,7 +217,7 @@ class MigrationDatabase: conn = self._get_conn() row = conn.execute( "SELECT target_msg_id FROM message_mappings WHERE channel_id = ? AND source_msg_id = ?", - (channel_id, source_id) + (parse_snowflake(channel_id), parse_snowflake(source_id)) ).fetchone() return row["target_msg_id"] if row else None @@ -205,7 +262,7 @@ class MigrationDatabase: conn = self._get_conn() # Check for existing alias - row = conn.execute("SELECT alias FROM user_alias WHERE user_id = ?", (str(user_id),)).fetchone() + row = conn.execute("SELECT alias FROM user_alias WHERE user_id = ?", (parse_snowflake(user_id),)).fetchone() if row: return row["alias"] @@ -215,7 +272,7 @@ class MigrationDatabase: new_alias = self._generate_alias() conn.execute( "INSERT INTO user_alias (user_id, alias) VALUES (?, ?)", - (str(user_id), new_alias) + (parse_snowflake(user_id), new_alias) ) conn.commit() return new_alias @@ -239,7 +296,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "INSERT OR REPLACE INTO server_mappings (category, source_id, target_id) VALUES (?, ?, ?)", - (category, str(source_id), str(target_id)) + (category, parse_snowflake(source_id), parse_snowflake(target_id)) ) conn.commit() @@ -247,7 +304,7 @@ class MigrationDatabase: conn = self._get_conn() row = conn.execute( "SELECT target_id FROM server_mappings WHERE category = ? AND source_id = ?", - (category, str(source_id)) + (category, parse_snowflake(source_id)) ).fetchone() return row["target_id"] if row else None @@ -263,7 +320,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "DELETE FROM server_mappings WHERE category = ? AND source_id = ?", - (category, str(source_id)) + (category, parse_snowflake(source_id)) ) conn.commit() @@ -281,7 +338,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "INSERT OR REPLACE INTO asset_mappings (category, source_id, target_id) VALUES (?, ?, ?)", - (category, str(source_id), str(target_id)) + (category, parse_snowflake(source_id), parse_snowflake(target_id)) ) conn.commit() @@ -289,7 +346,7 @@ class MigrationDatabase: conn = self._get_conn() row = conn.execute( "SELECT target_id FROM asset_mappings WHERE category = ? AND source_id = ?", - (category, str(source_id)) + (category, parse_snowflake(source_id)) ).fetchone() return row["target_id"] if row else None @@ -305,7 +362,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "DELETE FROM asset_mappings WHERE category = ? AND source_id = ?", - (category, str(source_id)) + (category, parse_snowflake(source_id)) ) conn.commit() @@ -332,23 +389,23 @@ class MigrationDatabase: def update_channel_tracking(self, channel_id: str, last_msg_id: str = None, last_msg_ts: str = None, msg_inc: int = 0, file_inc: int = 0): conn = self._get_conn() # Initialize if missing - conn.execute("INSERT OR IGNORE INTO channel_tracking (channel_id) VALUES (?)", (channel_id,)) + conn.execute("INSERT OR IGNORE INTO channel_tracking (channel_id) VALUES (?)", (parse_snowflake(channel_id),)) if last_msg_id: - conn.execute("UPDATE channel_tracking SET last_msg_id = ? WHERE channel_id = ?", (last_msg_id, channel_id)) + conn.execute("UPDATE channel_tracking SET last_msg_id = ? WHERE channel_id = ?", (parse_snowflake(last_msg_id), parse_snowflake(channel_id))) if last_msg_ts: - conn.execute("UPDATE channel_tracking SET last_msg_ts = ? WHERE channel_id = ?", (last_msg_ts, channel_id)) + conn.execute("UPDATE channel_tracking SET last_msg_ts = ? WHERE channel_id = ?", (last_msg_ts, parse_snowflake(channel_id))) if msg_inc != 0 or file_inc != 0: conn.execute( "UPDATE channel_tracking SET msg_count = msg_count + ?, file_count = file_count + ? WHERE channel_id = ?", - (msg_inc, file_inc, channel_id) + (msg_inc, file_inc, parse_snowflake(channel_id)) ) conn.commit() def get_channel_tracking(self, channel_id: str) -> Dict[str, Any]: conn = self._get_conn() - row = conn.execute("SELECT * FROM channel_tracking WHERE channel_id = ?", (channel_id,)).fetchone() + row = conn.execute("SELECT * FROM channel_tracking WHERE channel_id = ?", (parse_snowflake(channel_id),)).fetchone() if row: return dict(row) return {"last_msg_id": None, "last_msg_ts": None, "msg_count": 0, "file_count": 0} @@ -358,7 +415,7 @@ class MigrationDatabase: conn = self._get_conn() conn.execute( "INSERT OR REPLACE INTO thread_mappings (channel_id, thread_id, source_msg_id, target_msg_id, timestamp) VALUES (?, ?, ?, ?, ?)", - (channel_id, thread_id, source_id, target_id, timestamp) + (parse_snowflake(channel_id), parse_snowflake(thread_id), parse_snowflake(source_id), parse_snowflake(target_id), timestamp) ) conn.commit() @@ -366,31 +423,31 @@ class MigrationDatabase: conn = self._get_conn() row = conn.execute( "SELECT target_msg_id FROM thread_mappings WHERE channel_id = ? AND thread_id = ? AND source_msg_id = ?", - (channel_id, thread_id, source_id) + (parse_snowflake(channel_id), parse_snowflake(thread_id), parse_snowflake(source_id)) ).fetchone() return row["target_msg_id"] if row else None def update_thread_tracking(self, channel_id: str, thread_id: str, last_msg_id: str = None, last_msg_ts: str = None, msg_inc: int = 0, file_inc: int = 0, completed: int = None): conn = self._get_conn() - conn.execute("INSERT OR IGNORE INTO thread_tracking (channel_id, thread_id) VALUES (?, ?)", (channel_id, thread_id)) + conn.execute("INSERT OR IGNORE INTO thread_tracking (channel_id, thread_id) VALUES (?, ?)", (parse_snowflake(channel_id), parse_snowflake(thread_id))) if last_msg_id: - conn.execute("UPDATE thread_tracking SET last_msg_id = ? WHERE channel_id = ? AND thread_id = ?", (last_msg_id, channel_id, thread_id)) + conn.execute("UPDATE thread_tracking SET last_msg_id = ? WHERE channel_id = ? AND thread_id = ?", (parse_snowflake(last_msg_id), parse_snowflake(channel_id), parse_snowflake(thread_id))) if last_msg_ts: - conn.execute("UPDATE thread_tracking SET last_msg_ts = ? WHERE channel_id = ? AND thread_id = ?", (last_msg_ts, channel_id, thread_id)) + conn.execute("UPDATE thread_tracking SET last_msg_ts = ? WHERE channel_id = ? AND thread_id = ?", (last_msg_ts, parse_snowflake(channel_id), parse_snowflake(thread_id))) if completed is not None: - conn.execute("UPDATE thread_tracking SET completed = ? WHERE channel_id = ? AND thread_id = ?", (completed, channel_id, thread_id)) + conn.execute("UPDATE thread_tracking SET completed = ? WHERE channel_id = ? AND thread_id = ?", (completed, parse_snowflake(channel_id), parse_snowflake(thread_id))) if msg_inc != 0 or file_inc != 0: conn.execute( "UPDATE thread_tracking SET msg_count = msg_count + ?, file_count = file_count + ? WHERE channel_id = ? AND thread_id = ?", - (msg_inc, file_inc, channel_id, thread_id) + (msg_inc, file_inc, parse_snowflake(channel_id), parse_snowflake(thread_id)) ) conn.commit() def get_thread_tracking(self, channel_id: str, thread_id: str) -> Dict[str, Any]: conn = self._get_conn() - row = conn.execute("SELECT * FROM thread_tracking WHERE channel_id = ? AND thread_id = ?", (channel_id, thread_id)).fetchone() + row = conn.execute("SELECT * FROM thread_tracking WHERE channel_id = ? AND thread_id = ?", (parse_snowflake(channel_id), parse_snowflake(thread_id))).fetchone() if row: return dict(row) return {"last_msg_id": None, "last_msg_ts": None, "msg_count": 0, "file_count": 0} @@ -398,10 +455,10 @@ class MigrationDatabase: def clear_channel_data(self, channel_id: str): """Purge all mappings and tracking data for a specific channel and its threads.""" conn = self._get_conn() - conn.execute("DELETE FROM message_mappings WHERE channel_id = ?", (channel_id,)) - conn.execute("DELETE FROM thread_mappings WHERE channel_id = ?", (channel_id,)) - conn.execute("DELETE FROM channel_tracking WHERE channel_id = ?", (channel_id,)) - conn.execute("DELETE FROM thread_tracking WHERE channel_id = ?", (channel_id,)) + conn.execute("DELETE FROM message_mappings WHERE channel_id = ?", (parse_snowflake(channel_id),)) + conn.execute("DELETE FROM thread_mappings WHERE channel_id = ?", (parse_snowflake(channel_id),)) + conn.execute("DELETE FROM channel_tracking WHERE channel_id = ?", (parse_snowflake(channel_id),)) + conn.execute("DELETE FROM thread_tracking WHERE channel_id = ?", (parse_snowflake(channel_id),)) conn.commit() logger.info(f"Cleared all tracking and mapping data for channel: {channel_id}") diff --git a/src/core/exporter.py b/src/core/exporter.py index e60aaf7..7ad8926 100644 --- a/src/core/exporter.py +++ b/src/core/exporter.py @@ -109,7 +109,7 @@ class DiscordExporter: "name": r.name, "color": r.color.value, "position": r.position, - "permissions": str(r.permissions.value), + "permissions": r.permissions.value, "hoist": 1 if r.hoist else 0, "mentionable": 1 if r.mentionable else 0 }) @@ -563,9 +563,10 @@ class DiscordExporter: }) # 5. Message data + from src.core.utils import parse_snowflake message_reference = None if msg.reference and msg.reference.message_id: - message_reference = str(msg.reference.message_id) + message_reference = parse_snowflake(msg.reference.message_id) # 5.5 Forwarded snapshots content = msg.content or "" diff --git a/src/core/utils.py b/src/core/utils.py index 1375ef5..8402395 100644 --- a/src/core/utils.py +++ b/src/core/utils.py @@ -1,5 +1,19 @@ +from typing import Any, Optional import re import logging + +def parse_snowflake(value: Any) -> Optional[int]: + """Safely parses a Discord ID (Snowflake) from any input, handling 'None' strings.""" + if value is None: + return None + s = str(value).strip() + if not s or s.lower() == "none" or s == "NULL": + return None + try: + return int(s) + except ValueError: + return None + from src.core.state import MigrationState logger = logging.getLogger(__name__)