Xargana/discordb
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

fixed stuff

Browse source
This commit is contained in:
Xargana 2025-07-14 12:59:57 +03:00
parent db6b0a996b
commit 22f59ea2ad
2 changed files with 47 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

25
src/client.py
View file

@ -155,16 +155,26 @@ class DiscordDataClient(discord.Client):
try: try:
# Get all members - discord.py-self API # Get all members - discord.py-self API
members = [] members = []
member_iterator = await guild.fetch_members() member_result = await guild.fetch_members()
async for member in member_iterator:
members.append(member) # Handle different return types from fetch_members
if hasattr(member_result, '__aiter__'):
# It's an async iterator
async for member in member_result:
members.append(member)
elif isinstance(member_result, list):
# It's already a list
members = member_result
else:
# Try to iterate over it
for member in member_result:
members.append(member)
for member in members: for member in members:
if not member.bot: if not member.bot:
await self._process_user(member, guild.id) # Rate limiting before processing
# Rate limiting
await self.rate_limiter.wait() await self.rate_limiter.wait()
await self._process_user(member, guild.id)
self.logger.info(f"Processed {len(members)} members from {guild.name}") self.logger.info(f"Processed {len(members)} members from {guild.name}")
@ -224,6 +234,9 @@ class DiscordDataClient(discord.Client):
async def _get_mutual_guilds(self, user) -> List[int]: async def _get_mutual_guilds(self, user) -> List[int]:
"""Get mutual guilds for a user using fetch_user_profile.""" """Get mutual guilds for a user using fetch_user_profile."""
try: try:
# Rate limiting before API call
await self.rate_limiter.wait()
# Use fetch_user_profile to get mutual guilds # Use fetch_user_profile to get mutual guilds
profile = await self.fetch_user_profile(user.id, with_mutual_guilds=True) profile = await self.fetch_user_profile(user.id, with_mutual_guilds=True)

53
src/database.py
View file

@ -7,7 +7,7 @@ import json
import shutil import shutil
from datetime import datetime from datetime import datetime
from typing import Dict, List, Optional, Any from typing import Dict, List, Optional, Any
from dataclasses import dataclass, asdict from dataclasses import dataclass, asdict, field
from pathlib import Path from pathlib import Path
import logging import logging
@ -35,13 +35,11 @@ class UserData:
bio: Optional[str] = None bio: Optional[str] = None
status: Optional[str] = None status: Optional[str] = None
activity: Optional[str] = None activity: Optional[str] = None
servers: List[int] = None servers: List[int] = field(default_factory=list)
created_at: datetime = None created_at: datetime = None
updated_at: datetime = None updated_at: datetime = None
def __post_init__(self): def __post_init__(self):
if self.servers is None:
self.servers = []
current_time = datetime.utcnow() current_time = datetime.utcnow()
if self.created_at is None: if self.created_at is None:
self.created_at = current_time self.created_at = current_time
@ -458,38 +456,41 @@ class MariaDBDatabase:
async def save_server(self, server_id: int, server_name: str): async def save_server(self, server_id: int, server_name: str):
"""Save server information.""" """Save server information."""
async with self.pool.cursor() as cursor: async with self._lock:
await cursor.execute(""" async with self.pool.cursor() as cursor:
INSERT INTO servers (server_id, server_name) await cursor.execute("""
VALUES (%s, %s) INSERT INTO servers (server_id, server_name)
ON DUPLICATE KEY UPDATE VALUES (%s, %s)
server_name = VALUES(server_name), ON DUPLICATE KEY UPDATE
updated_at = CURRENT_TIMESTAMP server_name = VALUES(server_name),
""", (server_id, server_name)) updated_at = CURRENT_TIMESTAMP
""", (server_id, server_name))
async def get_server_names(self, server_ids: List[int]) -> Dict[int, str]: async def get_server_names(self, server_ids: List[int]) -> Dict[int, str]:
"""Get server names for given server IDs.""" """Get server names for given server IDs."""
if not server_ids: if not server_ids:
return {} return {}
async with self.pool.cursor() as cursor: async with self._lock:
placeholders = ','.join(['%s'] * len(server_ids)) async with self.pool.cursor() as cursor:
await cursor.execute(f""" placeholders = ','.join(['%s'] * len(server_ids))
SELECT server_id, server_name await cursor.execute(f"""
FROM servers SELECT server_id, server_name
WHERE server_id IN ({placeholders}) FROM servers
""", server_ids) WHERE server_id IN ({placeholders})
""", server_ids)
result = await cursor.fetchall() result = await cursor.fetchall()
return {row['server_id']: row['server_name'] for row in result} 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._lock:
await cursor.execute(""" async with self.pool.cursor() as cursor:
INSERT IGNORE INTO user_servers (user_id, server_id) await cursor.execute("""
VALUES (%s, %s) INSERT IGNORE INTO user_servers (user_id, server_id)
""", (user_id, server_id)) VALUES (%s, %s)
""", (user_id, server_id))
async def get_all_users(self) -> List[UserData]: async def get_all_users(self) -> List[UserData]:
"""Get all users from the database.""" """Get all users from the database."""