diff --git a/discordbot/titanembeds/bot.py b/discordbot/titanembeds/bot.py index b538bed..1c33940 100644 --- a/discordbot/titanembeds/bot.py +++ b/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) diff --git a/discordbot/titanembeds/database/__init__.py b/discordbot/titanembeds/database/__init__.py index 56431c6..3b03a80 100644 --- a/discordbot/titanembeds/database/__init__.py +++ b/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() diff --git a/discordbot/titanembeds/database/messages.py b/discordbot/titanembeds/database/messages.py new file mode 100644 index 0000000..d4bad0c --- /dev/null +++ b/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 ''.format(self.id, self.guild_id, self.guild_id, self.channel_id, self.message_id)