Use redis to store discord api related objects, so far messages only

Keeping the table for now in case the redis method does not work as hoped
This commit is contained in:
Jeremy Zhang 2018-07-16 03:50:31 +00:00
parent 9faef4f761
commit bc129289fc
10 changed files with 346 additions and 139 deletions

View File

@ -1,5 +1,6 @@
from config import config
from titanembeds.database import DatabaseInterface
from titanembeds.redisqueue import RedisQueue
from titanembeds.commands import Commands
from titanembeds.socketio import SocketIOInterface
from titanembeds.poststats import DiscordBotsOrg, BotsDiscordPw
@ -21,6 +22,7 @@ class Titan(discord.AutoShardedClient):
self.aiosession = aiohttp.ClientSession(loop=self.loop)
self.http.user_agent += ' TitanEmbeds-Bot'
self.database = DatabaseInterface(self)
self.redisqueue = RedisQueue(self, config["redis-uri"])
self.command = Commands(self, self.database)
self.socketio = SocketIOInterface(self, config["redis-uri"])
@ -57,6 +59,8 @@ class Titan(discord.AutoShardedClient):
async def start(self):
await self.database.connect(config["database-uri"])
await self.redisqueue.connect()
self.loop.create_task(self.redisqueue.subscribe())
await super().start(config["bot-token"])
async def on_ready(self):
@ -78,6 +82,7 @@ class Titan(discord.AutoShardedClient):
async def on_message(self, message):
await self.socketio.on_message(message)
await self.database.push_message(message)
await self.redisqueue.push_message(message)
msg_arr = message.content.split() # split the message
if len(message.content.split()) > 1 and message.guild: #making sure there is actually stuff in the message and have arguments and check if it is sent in server (not PM)
@ -92,11 +97,13 @@ class Titan(discord.AutoShardedClient):
async def on_message_edit(self, message_before, message_after):
await self.database.update_message(message_after)
await self.redisqueue.update_message(message_after)
await self.socketio.on_message_update(message_after)
async def on_message_delete(self, message):
self.delete_list.append(message.id)
await self.database.delete_message(message)
await self.redisqueue.delete_message(message)
await self.socketio.on_message_delete(message)
async def on_guild_join(self, guild):

View File

@ -0,0 +1,110 @@
from titanembeds.utils import get_formatted_message
from urllib.parse import urlparse
import asyncio_redis
import json
import discord
import asyncio
import traceback
import sys
import re
class RedisQueue:
def __init__(self, bot, redis_uri):
self.bot = bot
self.redis_uri = redis_uri
async def connect(self):
url_parsed = urlparse(self.redis_uri)
url_path = 0
if url_parsed.path and len(url_parsed.path) > 2:
url_path = int(url_parsed.path[1:])
self.sub_connection = await asyncio_redis.Connection.create(
host = url_parsed.hostname or "localhost",
port = url_parsed.port or 6379,
password = url_parsed.password,
db = url_path
)
self.connection = await asyncio_redis.Pool.create(
host = url_parsed.hostname or "localhost",
port = url_parsed.port or 6379,
password = url_parsed.password,
db = url_path,
poolsize = 10
)
async def subscribe(self):
await self.bot.wait_until_ready()
subscriber = await self.sub_connection.start_subscribe()
await subscriber.subscribe(["discord-api-req"])
while True:
reply = await subscriber.next_published()
request = json.loads(reply.value)
resource = request["resource"]
self.dispatch(resource, request["key"], request["params"])
def dispatch(self, event, key, params):
method = "on_" + event
if hasattr(self, method):
self.bot.loop.create_task(self._run_event(method, key, params))
async def _run_event(self, event, key, params):
try:
await getattr(self, event)(key, params)
except asyncio.CancelledError:
pass
except Exception:
try:
await self.on_error(event)
except asyncio.CancelledError:
pass
async def on_error(self, event_method):
print('Ignoring exception in {}'.format(event_method), file=sys.stderr)
traceback.print_exc()
async def set_scan_json(self, key, dict_key, dict_value_pattern):
unformatted_item = None
formatted_item = None
exists = await self.connection.exists(key)
if exists:
members = await self.connection.smembers(key)
for member in members:
the_member = await member
parsed = json.loads(the_member)
if re.match(str(dict_value_pattern), str(parsed[dict_key])):
unformatted_item = the_member
formatted_item = parsed
break
return (unformatted_item, formatted_item)
async def on_get_channel_messages(self, key, params):
channel = self.bot.get_channel(int(params["channel_id"]))
if not channel or not isinstance(channel, discord.channel.TextChannel):
return
await self.connection.delete([key])
messages = []
async for message in channel.history(limit=50):
formatted = get_formatted_message(message)
messages.append(json.dumps(formatted))
await self.connection.sadd(key, messages)
async def push_message(self, message):
if message.guild:
key = "Queue/channels/{}/messages".format(message.channel.id)
exists = await self.connection.exists(key)
if exists:
message = get_formatted_message(message)
await self.connection.sadd(key, [json.dumps(message)])
async def delete_message(self, message):
if message.guild:
key = "Queue/channels/{}/messages".format(message.channel.id)
exists = await self.connection.exists(key)
if exists:
unformatted_item, formatted_item = await self.set_scan_json(key, "id", message.id)
if formatted_item:
await self.connection.srem(key, [unformatted_item])
async def update_message(self, message):
await self.delete_message(message)
await self.push_message(message)

View File

@ -1,7 +1,5 @@
import socketio
from titanembeds.utils import get_message_author, get_message_mentions, get_roles_list, get_attachments_list, get_embeds_list
import time
from email import utils as emailutils
from titanembeds.utils import get_message_author, get_message_mentions, get_roles_list, get_attachments_list, get_embeds_list, get_formatted_message, get_formatted_user, get_formatted_emojis, get_formatted_guild, get_formatted_channel, get_formatted_role
import discord
class SocketIOInterface:
@ -9,178 +7,67 @@ class SocketIOInterface:
self.io = socketio.AsyncRedisManager(redis_uri, write_only=True, channel='flask-socketio')
self.bot = bot
def format_datetime(self, datetimeobj):
return emailutils.formatdate(time.mktime(datetimeobj.timetuple())) # https://stackoverflow.com/questions/3453177/convert-python-datetime-to-rfc-2822
def get_formatted_message(self, message):
edit_ts = message.edited_at
if not edit_ts:
edit_ts = None
else:
edit_ts = self.format_datetime(edit_ts)
msg = {
"id": str(message.id),
"channel_id": str(message.channel.id),
"content": message.content,
"author": get_message_author(message),
"timestamp": self.format_datetime(message.created_at),
"edited_timestamp": edit_ts,
}
if hasattr(message, "mentions"):
msg["mentions"] = get_message_mentions(message.mentions)
if hasattr(message, "attachments"):
msg["attachments"] = get_attachments_list(message.attachments)
if hasattr(message, "embeds"):
msg["embeds"] = get_embeds_list(message.embeds)
if hasattr(message, "author"):
nickname = None
if hasattr(message.author, 'nick') and message.author.nick:
nickname = message.author.nick
msg["author"]["nickname"] = nickname
if hasattr(message, "mentions"):
for mention in msg["mentions"]:
mention["nickname"] = None
member = message.guild.get_member(mention["id"])
if member:
mention["nickname"] = member.nick
return msg
async def on_message(self, message):
if message.guild:
msg = self.get_formatted_message(message)
msg = get_formatted_message(message)
await self.io.emit('MESSAGE_CREATE', data=msg, room=str("CHANNEL_"+str(message.channel.id)), namespace='/gateway')
async def on_message_delete(self, message):
if message.guild:
msg = self.get_formatted_message(message)
msg = get_formatted_message(message)
await self.io.emit('MESSAGE_DELETE', data=msg, room=str("CHANNEL_"+str(message.channel.id)), namespace='/gateway')
async def on_message_update(self, message):
if message.guild:
msg = self.get_formatted_message(message)
msg = get_formatted_message(message)
await self.io.emit('MESSAGE_UPDATE', data=msg, room=str("CHANNEL_"+str(message.channel.id)), namespace='/gateway')
def get_formatted_user(self, user):
userobj = {
"avatar": user.avatar,
"avatar_url": user.avatar_url,
"color": str(user.color)[1:],
"discriminator": user.discriminator,
"game": None,
"hoist-role": None,
"id": str(user.id),
"status": str(user.status),
"username": user.name,
"nick": None,
}
if userobj["color"] == "000000":
userobj["color"] = None
# if userobj["avatar_url"][len(userobj["avatar_url"])-15:] != ".jpg":
# userobj["avatar_url"] = userobj["avatar_url"][:len(userobj["avatar_url"])-14] + ".jpg"
if user.nick:
userobj["nick"] = user.nick
if hasattr(user, "activity") and user.activity:
userobj["activity"] = {
"name": user.activity.name
}
roles = sorted(user.roles, key=lambda k: k.position, reverse=True)
for role in roles:
if role.hoist:
userobj["hoist-role"] = {
"id": str(role.id),
"name": role.name,
"position": role.position,
}
break
return userobj
async def on_guild_member_add(self, member):
user = self.get_formatted_user(member)
user = get_formatted_user(member)
await self.io.emit('GUILD_MEMBER_ADD', data=user, room=str("GUILD_"+str(member.guild.id)), namespace='/gateway')
async def on_guild_member_remove(self, member):
user = self.get_formatted_user(member)
user = get_formatted_user(member)
await self.io.emit('GUILD_MEMBER_REMOVE', data=user, room=str("GUILD_"+str(member.guild.id)), namespace='/gateway')
async def on_guild_member_update(self, member):
user = self.get_formatted_user(member)
user = get_formatted_user(member)
await self.io.emit('GUILD_MEMBER_UPDATE', data=user, room=str("GUILD_"+str(member.guild.id)), namespace='/gateway')
def get_formatted_emojis(self, emojis):
emotes = []
for emo in emojis:
emotes.append({
"id": str(emo.id),
"managed": emo.managed,
"name": emo.name,
"require_colons": emo.require_colons,
"roles": get_roles_list(emo.roles),
"url": emo.url,
})
return emotes
async def on_guild_emojis_update(self, emojis):
emotes = self.get_formatted_emojis(emojis)
emotes = get_formatted_emojis(emojis)
await self.io.emit('GUILD_EMOJIS_UPDATE', data=emotes, room=str("GUILD_"+str(emojis[0].guild.id)), namespace='/gateway')
def get_formatted_guild(self, guild):
guil = {
"id": str(guild.id),
"name": guild.name,
"icon": guild.icon,
"icon_url": guild.icon_url,
}
return guil
async def on_guild_update(self, guild):
guildobj = self.get_formatted_guild(guild)
guildobj = get_formatted_guild(guild)
await self.io.emit('GUILD_UPDATE', data=guildobj, room=str("GUILD_"+str(guild.id)), namespace='/gateway')
def get_formatted_channel(self, channel):
chan = {
"id": str(channel.id),
"guild_id": str(channel.guild.id),
}
return chan
async def on_channel_delete(self, channel):
if str(channel.type) != "text":
return
chan = self.get_formatted_channel(channel)
chan = get_formatted_channel(channel)
await self.io.emit('CHANNEL_DELETE', data=chan, room=str("GUILD_"+str(channel.guild.id)), namespace='/gateway')
async def on_channel_create(self, channel):
if str(channel.type) != "text":
return
chan = self.get_formatted_channel(channel)
chan = get_formatted_channel(channel)
await self.io.emit('CHANNEL_CREATE', data=chan, room=str("GUILD_"+str(channel.guild.id)), namespace='/gateway')
async def on_channel_update(self, channel):
if not isinstance(channel, discord.channel.TextChannel) and not isinstance(channel, discord.channel.CategoryChannel):
return
chan = self.get_formatted_channel(channel)
chan = get_formatted_channel(channel)
await self.io.emit('CHANNEL_UPDATE', data=chan, room=str("GUILD_"+str(channel.guild.id)), namespace='/gateway')
def get_formatted_role(self, role):
rol = {
"id": str(role.id),
"guild_id": str(role.guild.id),
"name": role.name,
"color": role.color.value,
"hoist": role.hoist,
"position": role.position,
"permissions": role.permissions.value,
}
return rol
async def on_guild_role_create(self, role):
rol = self.get_formatted_role(role)
rol = get_formatted_role(role)
await self.io.emit('GUILD_ROLE_CREATE', data=rol, room=str("GUILD_"+str(role.guild.id)), namespace='/gateway')
async def on_guild_role_update(self, role):
rol = self.get_formatted_role(role)
rol = get_formatted_role(role)
await self.io.emit('GUILD_ROLE_UPDATE', data=rol, room=str("GUILD_"+str(role.guild.id)), namespace='/gateway')
async def on_guild_role_delete(self, role):
rol = self.get_formatted_role(role)
rol = get_formatted_role(role)
await self.io.emit('GUILD_ROLE_DELETE', data=rol, room=str("GUILD_"+str(role.guild.id)), namespace='/gateway')

View File

@ -1,4 +1,76 @@
import discord
import time
from email import utils as emailutils
def format_datetime(datetimeobj):
return emailutils.formatdate(time.mktime(datetimeobj.timetuple())) # https://stackoverflow.com/questions/3453177/convert-python-datetime-to-rfc-2822
def get_formatted_message(message):
edit_ts = message.edited_at
if not edit_ts:
edit_ts = None
else:
edit_ts = format_datetime(edit_ts)
msg = {
"id": str(message.id),
"channel_id": str(message.channel.id),
"content": message.content,
"author": get_message_author(message),
"timestamp": format_datetime(message.created_at),
"edited_timestamp": edit_ts,
}
if hasattr(message, "mentions"):
msg["mentions"] = get_message_mentions(message.mentions)
if hasattr(message, "attachments"):
msg["attachments"] = get_attachments_list(message.attachments)
if hasattr(message, "embeds"):
msg["embeds"] = get_embeds_list(message.embeds)
if hasattr(message, "author"):
nickname = None
if hasattr(message.author, 'nick') and message.author.nick:
nickname = message.author.nick
msg["author"]["nickname"] = nickname
if hasattr(message, "mentions"):
for mention in msg["mentions"]:
mention["nickname"] = None
member = message.guild.get_member(mention["id"])
if member:
mention["nickname"] = member.nick
return msg
def get_formatted_user(user):
userobj = {
"avatar": user.avatar,
"avatar_url": user.avatar_url,
"color": str(user.color)[1:],
"discriminator": user.discriminator,
"game": None,
"hoist-role": None,
"id": str(user.id),
"status": str(user.status),
"username": user.name,
"nick": None,
}
if userobj["color"] == "000000":
userobj["color"] = None
# if userobj["avatar_url"][len(userobj["avatar_url"])-15:] != ".jpg":
# userobj["avatar_url"] = userobj["avatar_url"][:len(userobj["avatar_url"])-14] + ".jpg"
if user.nick:
userobj["nick"] = user.nick
if hasattr(user, "activity") and user.activity:
userobj["activity"] = {
"name": user.activity.name
}
roles = sorted(user.roles, key=lambda k: k.position, reverse=True)
for role in roles:
if role.hoist:
userobj["hoist-role"] = {
"id": str(role.id),
"name": role.name,
"position": role.position,
}
break
return userobj
def get_message_author(message):
if not hasattr(message, "author"):
@ -12,6 +84,47 @@ def get_message_author(message):
"avatar": author.avatar
}
return obj
def get_formatted_emojis(emojis):
emotes = []
for emo in emojis:
emotes.append({
"id": str(emo.id),
"managed": emo.managed,
"name": emo.name,
"require_colons": emo.require_colons,
"roles": get_roles_list(emo.roles),
"url": emo.url,
})
return emotes
def get_formatted_guild(guild):
guil = {
"id": str(guild.id),
"name": guild.name,
"icon": guild.icon,
"icon_url": guild.icon_url,
}
return guil
def get_formatted_channel(channel):
chan = {
"id": str(channel.id),
"guild_id": str(channel.guild.id),
}
return chan
def get_formatted_role(role):
rol = {
"id": str(role.id),
"guild_id": str(role.guild.id),
"name": role.name,
"color": role.color.value,
"hoist": role.hoist,
"position": role.position,
"permissions": role.permissions.value,
}
return rol
def get_message_mentions(mentions):
ments = []

View File

@ -14,4 +14,5 @@ Flask-Babel
patreon
flask-redis
gino
sqlalchemy==1.2.8
sqlalchemy==1.2.8
asyncio_redis

View File

@ -1,6 +1,6 @@
from titanembeds.database import db, Guilds, UnauthenticatedUsers, UnauthenticatedBans, AuthenticatedUsers, GuildMembers, Messages, get_channel_messages, list_all_guild_members, get_guild_member, get_administrators_list, get_badges, DiscordBotsOrgTransactions
from titanembeds.database import db, Guilds, UnauthenticatedUsers, UnauthenticatedBans, AuthenticatedUsers, GuildMembers, Messages, list_all_guild_members, get_guild_member, get_administrators_list, get_badges, DiscordBotsOrgTransactions
from titanembeds.decorators import valid_session_required, discord_users_only, abort_if_guild_disabled
from titanembeds.utils import check_guild_existance, guild_accepts_visitors, guild_query_unauth_users_bool, get_client_ipaddr, discord_api, rate_limiter, channel_ratelimit_key, guild_ratelimit_key, user_unauthenticated, checkUserRevoke, checkUserBanned, update_user_status, check_user_in_guild, get_guild_channels, guild_webhooks_enabled, guild_unauthcaptcha_enabled, get_member_roles, get_online_embed_user_keys, redis_store
from titanembeds.utils import check_guild_existance, guild_accepts_visitors, guild_query_unauth_users_bool, get_client_ipaddr, discord_api, rate_limiter, channel_ratelimit_key, guild_ratelimit_key, user_unauthenticated, checkUserRevoke, checkUserBanned, update_user_status, check_user_in_guild, get_guild_channels, guild_webhooks_enabled, guild_unauthcaptcha_enabled, get_member_roles, get_online_embed_user_keys, redis_store, redisqueue
from titanembeds.oauth import user_has_permission, generate_avatar_url, check_user_can_administrate_guild
import titanembeds.constants as constants
from flask import Blueprint, abort, jsonify, session, request, url_for
@ -214,7 +214,7 @@ def get_all_users(guild_id):
def fetch():
guild_id = request.args.get("guild_id")
channel_id = request.args.get('channel_id')
after_snowflake = request.args.get('after', None, type=int)
after_snowflake = request.args.get('after', 0, type=int)
if user_unauthenticated():
key = session['user_keys'][guild_id]
else:
@ -233,7 +233,7 @@ def fetch():
if not chan.get("read") or chan["channel"]["type"] != "text":
status_code = 401
else:
messages = get_channel_messages(guild_id, channel_id, after_snowflake)
messages = redisqueue.get_channel_messages(guild_id, channel_id, after_snowflake)
status_code = 200
response = jsonify(messages=messages, status=status)
response.status_code = status_code
@ -245,7 +245,7 @@ def fetch():
def fetch_visitor():
guild_id = request.args.get("guild_id")
channel_id = request.args.get('channel_id')
after_snowflake = request.args.get('after', None, type=int)
after_snowflake = request.args.get('after', 0, type=int)
if not guild_accepts_visitors(guild_id):
abort(403)
messages = {}
@ -255,7 +255,7 @@ def fetch_visitor():
if not chan.get("read") or chan["channel"]["type"] != "text":
status_code = 401
else:
messages = get_channel_messages(guild_id, channel_id, after_snowflake)
messages = redisqueue.get_channel_messages(guild_id, channel_id, after_snowflake)
status_code = 200
response = jsonify(messages=messages)
response.status_code = status_code

View File

@ -27,7 +27,7 @@ class DiscordREST:
def _get_bucket(self, key):
value = redis_store.get(self.global_redis_prefix + key)
if value:
value = value.decode("utf-8")
value = value
return value
def _set_bucket(self, key, value):

View File

@ -47,7 +47,7 @@ def user_has_permission(permission, index):
def get_user_guilds():
cache = redis_store.get("OAUTH/USERGUILDS/"+str(make_user_cache_key()))
if cache:
return cache.decode("utf-8")
return cache
req = discordrest_from_user("/users/@me/guilds")
if req.status_code != 200:
if hasattr(request, "sid"):

View File

@ -0,0 +1,88 @@
from titanembeds.utils import redis_store
from titanembeds.database import get_guild_member
import json
import time
class RedisQueue:
def __init__(self):
pass # Nothing really to initialize
def get(self, key, resource, params, *, data_type="str"):
key = "Queue" + key
data = self._get(key, data_type)
payload = {
"key": key,
"resource": resource,
"params": params
}
loop_count = 0
while not data and loop_count < 10:
if loop_count % 5 == 0:
redis_store.publish("discord-api-req", json.dumps(payload))
time.sleep(0.5)
data = self._get(key, data_type)
loop_count += 1
redis_store.expire(key, 60 * 5)
if data == None:
return None
if data_type == "set":
data = list(data)
data_parsed = []
for d in data:
data_parsed.append(json.loads(d))
return data_parsed
return json.loads(data)
def _get(self, key, data_type):
if data_type == "set":
return redis_store.smembers(key)
else:
return redis_store.get(key)
def get_channel_messages(self, guild_id, channel_id, after_snowflake=0):
key = "/channels/{}/messages".format(channel_id)
q = self.get(key, "get_channel_messages", {"channel_id": channel_id}, data_type="set")
msgs = []
snowflakes = []
guild_members = {}
for x in q:
if x["id"] in snowflakes or int(x["id"]) <= int(after_snowflake):
continue
snowflakes.append(x["id"])
message = {
"attachments": x["attachments"],
"timestamp": x["timestamp"],
"id": x["id"],
"edited_timestamp": x["edited_timestamp"],
"author": x["author"],
"content": x["content"],
"channel_id": str(x["channel_id"]),
"mentions": x["mentions"],
"embeds": x["embeds"],
}
if message["author"]["id"] not in guild_members:
member = get_guild_member(guild_id, message["author"]["id"])
guild_members[message["author"]["id"]] = member
else:
member = guild_members[message["author"]["id"]]
message["author"]["nickname"] = None
if member:
message["author"]["nickname"] = member.nickname
message["author"]["avatar"] = member.avatar
message["author"]["discriminator"] = member.discriminator
message["author"]["username"] = member.username
for mention in message["mentions"]:
if mention["id"] not in guild_members:
author = get_guild_member(guild_id, mention["id"])
guild_members[mention["id"]] = author
else:
author = guild_members[mention["id"]]
mention["nickname"] = None
if author:
mention["nickname"] = author.nickname
mention["avatar"] = author.avatar
mention["username"] = author.username
mention["discriminator"] = author.discriminator
msgs.append(message)
sorted_msgs = sorted(msgs, key=lambda k: k['id'], reverse=True)
return sorted_msgs

View File

@ -13,11 +13,13 @@ import hashlib
import time
import json
redis_store = FlaskRedis()
redis_store = FlaskRedis(charset="utf-8", decode_responses=True)
from titanembeds.discordrest import DiscordREST
from titanembeds.redisqueue import RedisQueue
discord_api = DiscordREST(config['bot-token'])
redisqueue = RedisQueue()
def get_client_ipaddr():
if request.headers.getlist("X-Forwarded-For"):
@ -181,7 +183,6 @@ def get_online_embed_user_keys(guild_id="*", user_type=None):
usrs[utype] = []
keys = redis_store.keys("MemberPresence/{}/{}/*".format(guild_id, utype))
for key in keys:
key = str(key, "utf-8")
client_key = key.split("/")[-1]
usrs[utype].append(client_key)
return usrs