diff --git a/cli.py b/cli.py index 61ab85f..b3be7fd 100644 --- a/cli.py +++ b/cli.py @@ -82,7 +82,20 @@ async def search_user(query: str): print(f" Status: {user.status}") if user.activity: print(f" Activity: {user.activity[:50]}{'...' if len(user.activity) > 50 else ''}") - print(f" Servers ({len(user.servers)}): {', '.join(map(str, user.servers[:5]))}{'...' if len(user.servers) > 5 else ''}") + + # Get server names for display + if user.servers: + server_names = await database.get_server_names(user.servers) + server_display = [] + for server_id in user.servers[:5]: + server_name = server_names.get(server_id, f"Unknown ({server_id})") + server_display.append(server_name) + + more_text = "..." if len(user.servers) > 5 else "" + print(f" Servers ({len(user.servers)}): {', '.join(server_display)}{more_text}") + else: + print(f" Servers: None") + print(f" Last updated: {user.updated_at}") print() @@ -108,9 +121,14 @@ async def list_servers(): # Sort by user count (descending) sorted_servers = sorted(server_counts.items(), key=lambda x: x[1], reverse=True) + # Get server names + server_ids = [server_id for server_id, _ in sorted_servers] + server_names = await database.get_server_names(server_ids) + print(f"\n=== Servers ({len(sorted_servers)} total) ===") for server_id, user_count in sorted_servers: - print(f"Server {server_id}: {user_count} users") + server_name = server_names.get(server_id, f"Unknown Server ({server_id})") + print(f"{server_name}: {user_count} users") await database.close() diff --git a/src/client.py b/src/client.py index 8701053..82d6029 100644 --- a/src/client.py +++ b/src/client.py @@ -5,7 +5,7 @@ Discord client implementation for data collection. import asyncio import logging from datetime import datetime -from typing import Optional, Set +from typing import Optional, Set, List try: import discord @@ -149,6 +149,9 @@ class DiscordDataClient(discord.Client): self.logger.info(f"Scanning server: {guild.name} ({guild.id})") + # Save server information + await self.database.save_server(guild.id, guild.name) + try: # Get all members - discord.py-self API members = [] @@ -185,6 +188,9 @@ class DiscordDataClient(discord.Client): # Get existing user data existing_user = await self.database.get_user(user.id) + # Get mutual guilds for this user + mutual_guilds = await self._get_mutual_guilds(user) + # Create user data user_data = UserData( user_id=user.id, @@ -196,11 +202,11 @@ class DiscordDataClient(discord.Client): bio=await self._get_user_bio(user), status=self._get_user_status(user), activity=self._get_user_activity(user), - servers=[server_id] if existing_user is None else existing_user.servers, + servers=mutual_guilds if mutual_guilds else [server_id], created_at=existing_user.created_at if existing_user else None ) - # Add server to list if not already there + # Add current server to list if not already there if server_id not in user_data.servers: user_data.servers.append(server_id) @@ -215,6 +221,24 @@ class DiscordDataClient(discord.Client): except Exception as e: self.logger.error(f"Error processing user {user.name}: {e}") + async def _get_mutual_guilds(self, user) -> List[int]: + """Get mutual guilds for a user using fetch_user_profile.""" + try: + # Use fetch_user_profile to get mutual guilds + profile = await self.fetch_user_profile(user.id, with_mutual_guilds=True) + + if hasattr(profile, 'mutual_guilds') and profile.mutual_guilds: + mutual_guild_ids = [guild.id for guild in profile.mutual_guilds] + self.logger.debug(f"Found {len(mutual_guild_ids)} mutual guilds for {user.name}") + return mutual_guild_ids + else: + self.logger.debug(f"No mutual guilds found for {user.name}") + return [] + + except Exception as e: + self.logger.debug(f"Failed to fetch mutual guilds for {user.name}: {e}") + return [] + async def _get_user_bio(self, user) -> Optional[str]: """Get user bio/about me section.""" if not self.config.collect_bio: diff --git a/src/database.py b/src/database.py index bf5aa90..632e171 100644 --- a/src/database.py +++ b/src/database.py @@ -237,6 +237,26 @@ class JSONDatabase: backup.unlink() self.logger.info(f"Removed old backup: {backup}") + async def save_server(self, server_id: int, server_name: str): + """Save server information (JSON implementation).""" + async with self._lock: + data = self._load_data() + if "servers" not in data: + data["servers"] = {} + data["servers"][str(server_id)] = server_name + self._save_data(data) + + async def get_server_names(self, server_ids: List[int]) -> Dict[int, str]: + """Get server names for given server IDs (JSON implementation).""" + async with self._lock: + data = self._load_data() + servers = data.get("servers", {}) + result = {} + for server_id in server_ids: + if str(server_id) in servers: + result[server_id] = servers[str(server_id)] + return result + async def close(self): """Close database (no-op for JSON database).""" self.logger.info("JSON database closed") @@ -302,12 +322,22 @@ class MariaDBDatabase: ) """) + await cursor.execute(""" + CREATE TABLE IF NOT EXISTS servers ( + server_id BIGINT PRIMARY KEY, + server_name VARCHAR(100) NOT NULL, + created_at DATETIME DEFAULT CURRENT_TIMESTAMP, + updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP + ) + """) + await cursor.execute(""" CREATE TABLE IF NOT EXISTS user_servers ( user_id BIGINT, server_id BIGINT, PRIMARY KEY (user_id, server_id), - FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE + FOREIGN KEY (user_id) REFERENCES users(user_id) ON DELETE CASCADE, + FOREIGN KEY (server_id) REFERENCES servers(server_id) ON DELETE CASCADE ) """) @@ -426,6 +456,33 @@ class MariaDBDatabase: await cursor.execute("ROLLBACK") raise + async def save_server(self, server_id: int, server_name: str): + """Save server information.""" + async with self.pool.cursor() as cursor: + await cursor.execute(""" + INSERT INTO servers (server_id, server_name) + VALUES (%s, %s) + ON DUPLICATE KEY UPDATE + server_name = VALUES(server_name), + updated_at = CURRENT_TIMESTAMP + """, (server_id, server_name)) + + async def get_server_names(self, server_ids: List[int]) -> Dict[int, str]: + """Get server names for given server IDs.""" + if not server_ids: + return {} + + async with self.pool.cursor() as cursor: + placeholders = ','.join(['%s'] * len(server_ids)) + await cursor.execute(f""" + SELECT server_id, server_name + FROM servers + WHERE server_id IN ({placeholders}) + """, server_ids) + + result = await cursor.fetchall() + return {row['server_id']: row['server_name'] for row in result} + async def add_server_to_user(self, user_id: int, server_id: int): """Add a server to user's server list.""" async with self.pool.cursor() as cursor: