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}")
|
print(f" Status: {user.status}")
|
||||||
if user.activity:
|
if user.activity:
|
||||||
print(f" Activity: {user.activity[:50]}{'...' if len(user.activity) > 50 else ''}")
|
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(f" Last updated: {user.updated_at}")
|
||||||
print()
|
print()
|
||||||
|
|
||||||
|
@ -108,9 +121,14 @@ async def list_servers():
|
||||||
# Sort by user count (descending)
|
# Sort by user count (descending)
|
||||||
sorted_servers = sorted(server_counts.items(), key=lambda x: x[1], reverse=True)
|
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) ===")
|
print(f"\n=== Servers ({len(sorted_servers)} total) ===")
|
||||||
for server_id, user_count in sorted_servers:
|
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()
|
await database.close()
|
||||||
|
|
||||||
|
|
|
@ -5,7 +5,7 @@ Discord client implementation for data collection.
|
||||||
import asyncio
|
import asyncio
|
||||||
import logging
|
import logging
|
||||||
from datetime import datetime
|
from datetime import datetime
|
||||||
from typing import Optional, Set
|
from typing import Optional, Set, List
|
||||||
|
|
||||||
try:
|
try:
|
||||||
import discord
|
import discord
|
||||||
|
@ -149,6 +149,9 @@ class DiscordDataClient(discord.Client):
|
||||||
|
|
||||||
self.logger.info(f"Scanning server: {guild.name} ({guild.id})")
|
self.logger.info(f"Scanning server: {guild.name} ({guild.id})")
|
||||||
|
|
||||||
|
# Save server information
|
||||||
|
await self.database.save_server(guild.id, guild.name)
|
||||||
|
|
||||||
try:
|
try:
|
||||||
# Get all members - discord.py-self API
|
# Get all members - discord.py-self API
|
||||||
members = []
|
members = []
|
||||||
|
@ -185,6 +188,9 @@ class DiscordDataClient(discord.Client):
|
||||||
# Get existing user data
|
# Get existing user data
|
||||||
existing_user = await self.database.get_user(user.id)
|
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
|
# Create user data
|
||||||
user_data = UserData(
|
user_data = UserData(
|
||||||
user_id=user.id,
|
user_id=user.id,
|
||||||
|
@ -196,11 +202,11 @@ class DiscordDataClient(discord.Client):
|
||||||
bio=await self._get_user_bio(user),
|
bio=await self._get_user_bio(user),
|
||||||
status=self._get_user_status(user),
|
status=self._get_user_status(user),
|
||||||
activity=self._get_user_activity(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
|
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:
|
if server_id not in user_data.servers:
|
||||||
user_data.servers.append(server_id)
|
user_data.servers.append(server_id)
|
||||||
|
|
||||||
|
@ -215,6 +221,24 @@ class DiscordDataClient(discord.Client):
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
self.logger.error(f"Error processing user {user.name}: {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]:
|
async def _get_user_bio(self, user) -> Optional[str]:
|
||||||
"""Get user bio/about me section."""
|
"""Get user bio/about me section."""
|
||||||
if not self.config.collect_bio:
|
if not self.config.collect_bio:
|
||||||
|
|
|
@ -237,6 +237,26 @@ class JSONDatabase:
|
||||||
backup.unlink()
|
backup.unlink()
|
||||||
self.logger.info(f"Removed old backup: {backup}")
|
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):
|
async def close(self):
|
||||||
"""Close database (no-op for JSON database)."""
|
"""Close database (no-op for JSON database)."""
|
||||||
self.logger.info("JSON database closed")
|
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("""
|
await cursor.execute("""
|
||||||
CREATE TABLE IF NOT EXISTS user_servers (
|
CREATE TABLE IF NOT EXISTS user_servers (
|
||||||
user_id BIGINT,
|
user_id BIGINT,
|
||||||
server_id BIGINT,
|
server_id BIGINT,
|
||||||
PRIMARY KEY (user_id, server_id),
|
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")
|
await cursor.execute("ROLLBACK")
|
||||||
raise
|
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):
|
async def add_server_to_user(self, user_id: int, server_id: int):
|
||||||
"""Add a server to user's server list."""
|
"""Add a server to user's server list."""
|
||||||
async with self.pool.cursor() as cursor:
|
async with self.pool.cursor() as cursor:
|
||||||
|
|
Loading…
Reference in a new issue