Titan/discordbot/cleanup_db_messages.py

from config import config
from titanembeds.database import DatabaseInterface, Guilds, Messages
from titanembeds.commands import Commands
import discord
import aiohttp
import asyncio
import sys
import logging
import json
import gc
import random
from asyncio_extras import threadpool
logging.basicConfig(filename='titanbot.log',level=logging.INFO,format='%(asctime)s %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p')
logging.getLogger('TitanBot')
logging.getLogger('sqlalchemy')

###########################
#   Cleanup DB Messages   #
#                         #
# Cleans the database     #
# messages store          #
###########################

class Titan(discord.AutoShardedClient):
    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)
        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(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:')
        print(self.user.name)
        print(self.user.id)
        print('------')
        
        game = discord.Game(name="Titan is currently down for database maintenances. Bookmark https://TitanEmbeds.com/ for later access to our services!")
        await self.change_presence(status=discord.Status.do_not_disturb, activity=game)

        try:
            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...")
        self.loop.create_task(self.start_cleanup())

    async def start_cleanup(self):
        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!")

def main():
    print("Starting...")
    te = Titan()
    te.run()
    gc.collect()

if __name__ == '__main__':
    main()