Messages database

discordbot/titanembeds/bot.py

@@ -55,3 +55,10 @@ class Titan(discord.Client):
             traceback.print_exc()
             await self.logout()
             return
+
+    async def on_message(self, message):
+        await self.database.push_message(message)
+        # TODO: Will add command handler + ban/kick command
+
+    async def on_message_edit(self, message_before, message_after):
+        await self.database.update_message(message_after)

discordbot/titanembeds/database/__init__.py

@@ -5,9 +5,12 @@ from sqlalchemy.engine import Engine, create_engine
 from sqlalchemy.orm import sessionmaker, Session
 from sqlalchemy.ext.declarative import declarative_base
 
+import json
+
 Base = declarative_base()
 
 from titanembeds.database.guilds import Guilds
+from titanembeds.database.messages import Messages
 
 class DatabaseInterface(object):
     # Courtesy of https://github.com/SunDwarf/Jokusoramame
@@ -33,3 +36,40 @@ class DatabaseInterface(object):
             raise
         finally:
             session.close()
+
+    async def push_message(self, message):
+        if message.server:
+            async with threadpool():
+                with self.get_session() as session:
+                    edit_ts = message.edited_timestamp
+                    if not edit_ts:
+                        edit_ts = None
+                    else:
+                        edit_ts = str(edit_ts)
+
+                    msg = Messages(
+                        message.server.id,
+                        message.channel.id,
+                        message.id,
+                        message.content,
+                        str(message.timestamp),
+                        edit_ts,
+                        json.dumps(message.mentions),
+                        json.dumps(message.attachments)
+                    )
+                    session.add(msg)
+                    session.commit()
+
+    async def update_message(self, message):
+        if message.server:
+            async with threadpool():
+                with self.get_session() as session:
+                    msg = session.query(Messages) \
+                        .filter(Messages.guild_id == message.server.id) \
+                        .filter(Messages.channel_id == message.channel.id) \
+                        .filter(Messages.message_id == message.id).first()
+                    msg.content = message.content
+                    msg.edited_timestamp = message.edited_timestamp
+                    msg.mentions = json.dumps(message.mentions)
+                    msg.attachments = json.dumps(message.attachments)
+                    session.commit()

discordbot/titanembeds/database/messages.py

@@ -0,0 +1,26 @@
+from titanembeds.database import db, Base
+
+class Messages(Base):
+    __tablename__ = "messages"
+    id = db.Column(db.Integer, primary_key=True)    # Auto incremented id
+    guild_id = db.Column(db.String(255))            # Discord guild id
+    channel_id = db.Column(db.String(255))          # Channel id
+    message_id = db.Column(db.String(255))          # Message snowflake
+    content = db.Column(db.Text())                  # Message contents
+    timestamp = db.Column(db.TIMESTAMP)             # Timestamp of when content is created
+    edited_timestamp = db.Column(db.TIMESTAMP)      # Timestamp of when content is edited
+    mentions = db.Column(db.Text())                 # Mentions serialized
+    attachments = db.Column(db.Text())              # serialized attachments
+
+    def __init__(self, guild_id, channel_id, message_id, content, timestamp, edited_timestamp, mentions, attachments):
+        self.guild_id = guild_id
+        self.channel_id = channel_id
+        self.message_id = message_id
+        self.content = content
+        self.timestamp = timestamp
+        self.edited_timestamp = edited_timestamp
+        self.mentions = mentions
+        self.attachments = attachments
+
+    def __repr__(self):
+        return '<Messages {0} {1} {2} {3} {4}>'.format(self.id, self.guild_id, self.guild_id, self.channel_id, self.message_id)