From 7762862623c3c0f9528345a56fec8481847b71ef Mon Sep 17 00:00:00 2001
From: Jeremy Zhang <jeremyzhang5000@gmail.com>
Date: Wed, 3 May 2017 22:16:49 -0700
Subject: [PATCH] Discordbot database setup inital/partial

---
 discordbot/requirements.txt                 |  3 +-
 discordbot/run.py                           | 13 ++++-
 discordbot/titanembeds/__init__.py          |  1 +
 discordbot/titanembeds/bot.py               | 64 +++++++++++++++++----
 discordbot/titanembeds/database/__init__.py | 36 +++++++++---
 5 files changed, 92 insertions(+), 25 deletions(-)

diff --git a/discordbot/requirements.txt b/discordbot/requirements.txt
index ee91f00..400bc22 100644
--- a/discordbot/requirements.txt
+++ b/discordbot/requirements.txt
@@ -1,2 +1,3 @@
 discord.py
-sqlalchemy-aio
\ No newline at end of file
+sqlalchemy
+asyncio_extras
diff --git a/discordbot/run.py b/discordbot/run.py
index 1ec70b8..c6b058c 100644
--- a/discordbot/run.py
+++ b/discordbot/run.py
@@ -1,4 +1,11 @@
-from titanembeds.bot import client
-from config import config
+from titanembeds import Titan
+import gc
 
-client.run(config["bot-token"])
+def main():
+    print("Starting...")
+    te = Titan()
+    te.run()
+    gc.collect()
+
+if __name__ == '__main__':
+    main()
diff --git a/discordbot/titanembeds/__init__.py b/discordbot/titanembeds/__init__.py
index e69de29..9960b6e 100644
--- a/discordbot/titanembeds/__init__.py
+++ b/discordbot/titanembeds/__init__.py
@@ -0,0 +1 @@
+from titanembeds.bot import Titan
diff --git a/discordbot/titanembeds/bot.py b/discordbot/titanembeds/bot.py
index 55457f9..b538bed 100644
--- a/discordbot/titanembeds/bot.py
+++ b/discordbot/titanembeds/bot.py
@@ -1,17 +1,57 @@
 from config import config
+from titanembeds.database import DatabaseInterface
 import discord
+import aiohttp
+import asyncio
 
-client = discord.Client()
+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)
 
-@client.event
-async def on_ready():
-    print('Titan -- DiscordBot')
-    print('Logged in as the following user:')
-    print(client.user.name)
-    print(client.user.id)
-    print('------')
-    await test()
+    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
 
-async def test():
-    from titanembeds.database import db, Guilds, session
-    session.query(Guilds).all()
\ No newline at end of file
+    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]')
+        print('Logged in as the following user:')
+        print(self.user.name)
+        print(self.user.id)
+        print('------')
+
+        await self.change_presence(
+            game=discord.Game(name="Get your own @ https://TitanEmbeds.tk/"), status=discord.Status.online
+        )
+
+        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
diff --git a/discordbot/titanembeds/database/__init__.py b/discordbot/titanembeds/database/__init__.py
index 0c8a8db..56431c6 100644
--- a/discordbot/titanembeds/database/__init__.py
+++ b/discordbot/titanembeds/database/__init__.py
@@ -1,17 +1,35 @@
-from config import config
-from sqlalchemy_aio import ASYNCIO_STRATEGY
+from contextlib import contextmanager
+from asyncio_extras import threadpool
 import sqlalchemy as db
+from sqlalchemy.engine import Engine, create_engine
+from sqlalchemy.orm import sessionmaker, Session
 from sqlalchemy.ext.declarative import declarative_base
 
 Base = declarative_base()
 
-from .guilds import Guilds
+from titanembeds.database.guilds import Guilds
 
-engine = db.create_engine(config["database-uri"])
+class DatabaseInterface(object):
+    # Courtesy of https://github.com/SunDwarf/Jokusoramame
+    def __init__(self, bot):
+        self.bot = bot
 
-Base.metadata.create_all(engine)
+        self.engine = None  # type: Engine
+        self._sessionmaker = None  # type: sessionmaker
 
-from sqlalchemy.orm import sessionmaker
-DBSession = sessionmaker()
-DBSession.bind = engine
-session = DBSession()
\ No newline at end of file
+    async def connect(self, dburi):
+        async with threadpool():
+            self.engine = create_engine(dburi)
+            self._sessionmaker = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+    @contextmanager
+    def get_session(self) -> Session:
+        session = self._sessionmaker()  # type: Session
+        try:
+            yield session
+            session.commit()
+        except:
+            session.rollback()
+            raise
+        finally:
+            session.close()