From 25d80c82050b4af8b522a7b54099f9b0a5bb85cf Mon Sep 17 00:00:00 2001 From: Jeremy Zhang Date: Sat, 23 Jun 2018 06:11:19 +0000 Subject: [PATCH] Build cleanup db into the main titan code, executing it on bot startup/ready --- discordbot/titanembeds/bot.py | 3 ++ discordbot/titanembeds/cleanup_db.py | 51 ++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+) create mode 100644 discordbot/titanembeds/cleanup_db.py diff --git a/discordbot/titanembeds/bot.py b/discordbot/titanembeds/bot.py index 25818ee..fc52eb9 100644 --- a/discordbot/titanembeds/bot.py +++ b/discordbot/titanembeds/bot.py @@ -3,6 +3,7 @@ from titanembeds.database import DatabaseInterface from titanembeds.commands import Commands from titanembeds.socketio import SocketIOInterface from titanembeds.poststats import DiscordBotsOrg, BotsDiscordPw +from titanembeds.cleanup_db import CleanupDatabase from collections import deque import discord import aiohttp @@ -23,6 +24,7 @@ class Titan(discord.AutoShardedClient): self.database = DatabaseInterface(self) self.command = Commands(self, self.database) self.socketio = SocketIOInterface(self, config["redis-uri"]) + self.cleanup_db = CleanupDatabase(self, self.database) self.delete_list = deque(maxlen=100) # List of msg ids to prevent duplicate delete @@ -78,6 +80,7 @@ class Titan(discord.AutoShardedClient): self.discordBotsOrg = DiscordBotsOrg(self.user.id, config.get("discord-bots-org-token", None)) self.botsDiscordPw = BotsDiscordPw(self.user.id, config.get("bots-discord-pw-token", None)) await self.postStats() + await self.cleanup_db.start_cleanup() async def on_message(self, message): await self.socketio.on_message(message) diff --git a/discordbot/titanembeds/cleanup_db.py b/discordbot/titanembeds/cleanup_db.py new file mode 100644 index 0000000..9710f4f --- /dev/null +++ b/discordbot/titanembeds/cleanup_db.py @@ -0,0 +1,51 @@ +from titanembeds.database.guilds import Guilds +from titanembeds.database.messages import Messages +import logging +import json +import random + +class CleanupDatabase: + def __init__(self, bot, db): + self.bot = bot + self.db = db + 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) + self.logger.info("Initialized Database Cleaning Class with session id of " + session_id) + + async def start_cleanup(self): + self.logger.info("Started cleaning up task") + self.logger.info("Cleaning up unused guilds") + await self.db.remove_unused_guilds(list(self.bot.guilds)) + self.logger.info("Started cleaning up messages task, hopefully cleaned up unused guilds") + self.bot.loop.run_in_executor(None, self._start_cleanup) + + def _start_cleanup(self): + with self.db.get_session() as session: + guilds = session.query(Guilds).all() + count = 0 + for guild in guilds: + 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!") \ No newline at end of file