fix: async thread on discord but

This commit is contained in:
Muhamad Aji Wibisono 2025-06-02 19:22:04 +07:00
parent 1d67a87b82
commit a0f9efdb67
2 changed files with 46 additions and 9 deletions

View file

@ -10,11 +10,12 @@ import logging
import discord
from discord.ext import commands
import datetime
import asyncio
logger = logging.getLogger(__name__)
class DiscordConnector:
class DiscordConnector(commands.Bot):
"""Class for retrieving guild, channel, and message history from Discord."""
def __init__(self, token: str = None):
@ -31,23 +32,28 @@ class DiscordConnector:
intents.members = True # Required to fetch member information
super().__init__(command_prefix="!", intents=intents) # command_prefix is required but not strictly used here
self.token = token
self._ready = False
self._bot_task = None # Holds the async bot task
# Event to confirm bot is ready
@self.event
async def on_ready():
logger.info(f"Logged in as {self.user} (ID: {self.user.id})")
self._ready = True
async def start_bot(self):
"""Starts the bot to connect to Discord."""
logger.info("Starting Discord bot...")
if not self.token:
raise ValueError("Discord bot token not set. Call set_token(token) first.")
await self.start(self.token)
logger.info("Discord bot started successfully.")
async def close_bot(self):
"""Closes the bot's connection to Discord."""
logger.info("Closing Discord bot connection...")
await self.close()
logger.info("Discord bot connection closed.")
def set_token(self, token: str) -> None:
@ -57,14 +63,21 @@ class DiscordConnector:
Args:
token (str): The Discord bot token.
"""
logger.info("Setting Discord bot token.")
self.token = token
logger.info("Token set successfully. You can now start the bot with start_bot().")
async def _wait_until_ready(self):
"""Helper to wait until the bot is connected and ready."""
if not self._ready:
logger.info("Bot not yet ready, waiting...")
await self.wait_until_ready()
logger.info("Bot is now ready.")
logger.info("Waiting for the bot to be ready...")
# Give the event loop a chance to switch to the bot's startup task.
# This allows self.start() to begin initializing the client.
# Terrible solution, but necessary to avoid blocking the event loop.
await asyncio.sleep(1) # Yield control to the event loop
await self.wait_until_ready()
logger.info("Bot is ready.")
async def get_guilds(self) -> list[dict]:
"""
@ -78,6 +91,7 @@ class DiscordConnector:
ValueError: If the token is not set.
"""
await self._wait_until_ready()
logger.info("Fetching guilds...")
guilds_data = []
for guild in self.guilds:
@ -89,6 +103,8 @@ class DiscordConnector:
"member_count": member_count,
}
)
logger.info(f"Fetched {len(guilds_data)} guilds.")
return guilds_data
async def get_text_channels(self, guild_id: str) -> list[dict]:
@ -106,6 +122,7 @@ class DiscordConnector:
discord.NotFound: If the guild is not found.
"""
await self._wait_until_ready()
logger.info(f"Fetching text channels for guild ID: {guild_id}")
guild = self.get_guild(int(guild_id))
if not guild:
@ -118,6 +135,8 @@ class DiscordConnector:
channels_data.append(
{"id": str(channel.id), "name": channel.name, "type": "text"}
)
logger.info(f"Fetched {len(channels_data)} text channels from guild {guild_id}.")
return channels_data
async def get_channel_history(
@ -146,6 +165,7 @@ class DiscordConnector:
discord.Forbidden: If the bot does not have permissions to read history in the channel.
"""
await self._wait_until_ready()
logger.info(f"Fetching message history for channel ID: {channel_id}")
channel = self.get_channel(int(channel_id))
if not channel:
@ -190,7 +210,8 @@ class DiscordConnector:
except discord.HTTPException as e:
logger.error(f"Failed to fetch messages from channel {channel_id}: {e}")
return []
logger.info(f"Fetched {len(messages_data)} messages from channel {channel_id}.")
return messages_data
async def get_user_info(self, guild_id: str, user_id: str) -> dict | None:
@ -211,6 +232,7 @@ class DiscordConnector:
permissions to view members.
"""
await self._wait_until_ready()
logger.info(f"Fetching user info for user ID: {user_id} in guild ID: {guild_id}")
guild = self.get_guild(int(guild_id))
if not guild:
@ -221,12 +243,15 @@ class DiscordConnector:
member = await guild.fetch_member(int(user_id))
if member:
roles = [role.name for role in member.roles if role.name != "@everyone"]
logger.info(f"User {user_id} found in guild {guild_id}.")
return {
"id": str(member.id),
"name": member.name,
"joined_at": member.joined_at.isoformat() if member.joined_at else None,
"roles": roles,
}
logger.warning(f"User {user_id} not found in guild {guild_id}.")
return None
except discord.NotFound:
logger.warning(f"User {user_id} not found in guild {guild_id}.")

View file

@ -15,6 +15,7 @@ from app.connectors.discord_connector import DiscordConnector
from discord import DiscordException
from slack_sdk.errors import SlackApiError
import logging
import asyncio
from app.utils.document_converters import generate_content_hash
@ -952,6 +953,9 @@ async def index_discord_messages(
if not discord_token:
return 0, "Discord token not found in connector config"
logger.info(f"Starting Discord indexing for connector {connector_id}")
logger.info(f"discord token is: {discord_token[:4]}... (truncated for security)")
# Initialize Discord client
discord_client = DiscordConnector(token=discord_token)
@ -975,13 +979,21 @@ async def index_discord_messages(
skipped_guilds = []
try:
await discord_client.start_bot()
logger.info("Starting Discord bot to fetch guilds")
discord_client._bot_task = asyncio.create_task(discord_client.start_bot())
await discord_client._wait_until_ready()
logger.info("Fetching Discord guilds")
guilds = await discord_client.get_guilds()
logger.info(f"Found {len(guilds)} guilds")
except Exception as e:
logger.error(f"Failed to get Discord guilds: {str(e)}", exc_info=True)
await discord_client.close_bot()
return 0, f"Failed to get Discord guilds: {str(e)}"
if not guilds:
logger.info("No Discord guilds found to index")
await discord_client.close_bot()
return 0, "No Discord guilds found"