class Titan(discord.Client):
def __init__(self):
super().__init__()
self.aiosession = aiohttp.ClientSession(loop=self.loop)
self.http.user_agent += ' TitanEmbeds-Bot'
self.database = DatabaseInterface(self)
self.command = Commands(self, self.database)
def _cleanup(self):
try:
self.loop.run_until_complete(self.logout())
except: # Can be ignored
pass
pending = asyncio.Task.all_tasks()
gathered = asyncio.gather(*pending)
try:
gathered.cancel()
self.loop.run_until_complete(gathered)
gathered.exception()
except: # Can be ignored
pass
def run(self):
try:
self.loop.run_until_complete(self.start(config["bot-token"]))
except discord.errors.LoginFailure:
print("Invalid bot token in config!")
finally:
try:
self._cleanup()
except Exception as e:
print("Error in cleanup:", e)
self.loop.close()
async def on_ready(self):
print('Titan [DiscordBot] [UTILITY: Cleanup database messages]')
print('Logged in as the following user:'******'------')
try:
await self.database.connect(config["database-uri"])
except Exception:
self.logger.error("Unable to connect to specified database!")
traceback.print_exc()
await self.logout()
return
print("working on this...")
async with threadpool():
with self.database.get_session() as session:
guilds = session.query(Guilds).all()
count = 0
for guild in guilds:
count += 1
print("[{}] snowflake-{} name-{}".format(
count, guild.guild_id, guild.name))
try:
channelsjson = json.loads(guild.channels)
except:
continue
active_channels = []
for channel in channelsjson:
chanid = channel["id"]
active_channels.append(chanid)
keep_these = session.query(Messages.message_id).filter(
Messages.channel_id == chanid).order_by(
Messages.timestamp.desc()).limit(50)
d = session.query(Messages).filter(
Messages.channel_id == chanid,
~Messages.message_id.in_(keep_these)).delete(
synchronize_session=False)
session.commit()
print(" --{} [{}]".format(channel["name"], d))
d = session.query(Messages).filter(
Messages.guild_id == guild.guild_id,
~Messages.channel_id.in_(active_channels)).delete(
synchronize_session=False)
session.commit()
print(" INACTIVE {}".format(d))
print("done!")
await self.logout()
class TitanCleanupDB:
def __init__(self):
super().__init__()
self.loop = asyncio.get_event_loop()
self.database = DatabaseInterface(self)
self.logger = logging.getLogger("titan_cleanupdb")
self.logger.setLevel(logging.DEBUG)
fh = logging.FileHandler("titan_cleanupdb.log")
fh.setLevel(logging.DEBUG)
session_id = str(random.randrange(100))
formatter = logging.Formatter(
"%(asctime)s - {0} - %(levelname)s - %(message)s".format(
session_id))
fh.setFormatter(formatter)
self.logger.addHandler(fh)
consoleHandler = logging.StreamHandler()
consoleHandler.setFormatter(formatter)
self.logger.addHandler(consoleHandler)
self.logger.info(
"Initialized Database Cleaning Class with session id of " +
session_id)
def _cleanup(self):
try:
self.loop.run_until_complete(self.logout())
except: # Can be ignored
pass
pending = asyncio.Task.all_tasks()
gathered = asyncio.gather(*pending)
try:
gathered.cancel()
self.loop.run_until_complete(gathered)
gathered.exception()
except: # Can be ignored
pass
def run(self):
try:
self.loop.run_until_complete(self.start_cleanup())
except Exception as e:
print("Error!", e)
finally:
try:
self._cleanup()
except Exception as e:
print("Error in cleanup:", e)
self.loop.close()
async def start_cleanup(self):
print('Titan [DiscordBot] [UTILITY: Cleanup database messages]')
print('------')
try:
self.database.connect(config["database-uri"])
except Exception:
self.logger.error("Unable to connect to specified database!")
traceback.print_exc()
return
print("working on this...")
with self.database.get_session() as session:
guilds = session.query(Guilds).all()
guilds_new = []
count = 0
for guild in guilds:
guilds_new.append(guild)
for guild in guilds_new:
count += 1
self.logger.info("[{}] snowflake-{} name-{}".format(
count, guild.guild_id, guild.name))
try:
channelsjson = json.loads(guild.channels)
except:
continue
active_channels = []
for channel in channelsjson:
chanid = channel["id"]
active_channels.append(chanid)
keep_these = session.query(Messages.message_id).filter(
Messages.channel_id == chanid).order_by(
Messages.timestamp.desc()).limit(50)
d = session.query(Messages).filter(
Messages.channel_id == chanid,
~Messages.message_id.in_(keep_these)).delete(
synchronize_session=False)
session.commit()
self.logger.info(" --{} [{}]".format(
channel["name"], d))
d = session.query(Messages).filter(
Messages.guild_id == guild.guild_id,
~Messages.channel_id.in_(active_channels)).delete(
synchronize_session=False)
session.commit()
self.logger.info(" INACTIVE {}".format(d))
self.logger.info("done!")