added mutual guilds
This commit is contained in:
parent
c97c9d1b65
commit
db6b0a996b
22
cli.py
22
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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Reference in a new issue