Added entire modular starboard system

pull/9/head
sgoudham 4 years ago
parent b0ce830ea5
commit dcbfddcb95

@ -78,6 +78,8 @@ class Bot(commands.Bot):
# Instance variables for cache
self.enso_cache = {}
self.modmail_cache = {}
self.starboard_cache = {}
self.starboard_messages_cache = {}
self.member_cache = MyCoolCache(100)
async def create_connection():
@ -125,6 +127,31 @@ class Bot(commands.Bot):
except asyncpg.PostgresError as e:
print("PostGres Error: Modmail Records Could Not Be Loaded Into Cache On Startup", e)
# Query to get all records of starboards within guilds
try:
results = await conn.fetch("SELECT * FROM starboard")
# Store the information for starboard within cache
for row in results:
self.starboard_cache[row["guild_id"]] = {"channel_id": row["channel_id"]}
# Catch errors
except asyncpg.PostgresError as e:
print("PostGres Error: Starboard Records Could Not Be Loaded Into Cache On Startup", e)
# Query to get all records of starboard messages within guilds
try:
results = await conn.fetch("SELECT * FROM starboard_messages")
for row in results:
self.starboard_messages_cache[(row["root_message_id"], row["guild_id"])] = {
"star_message_id": row["star_message_id"],
"stars": row["stars"]}
# Catch errors
except asyncpg.PostgresError as e:
print("PostGres Error: Starboard Records Could Not Be Loaded Into Cache On Startup", e)
# Release connection back to pool
await pool.release(conn)
@ -183,8 +210,135 @@ class Bot(commands.Bot):
del self.enso_cache[guild_id]
async def check_cache(self, member_id, guild_id):
"""Checks if member is in the member cache"""
# Return key-value pair if member is already in the cache
if (member_id, guild_id) in self.member_cache.cache:
return self.member_cache.cache[member_id, guild_id]
else:
# Setup pool connection
pool = self.db
async with pool.acquire() as conn:
# Get the author's/members row from the Members Table
try:
select_query = """SELECT * FROM members WHERE member_id = $1 and guild_id = $2"""
result = await conn.fetchrow(select_query, member_id, guild_id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: Member {member_id} From Guild {guild_id}"
"Record Could Not Be Retrieved When Checking Cache", e)
# Store it in cache
else:
dict_items = {"married": result["married"],
"married_date": result["married_date"],
"muted_roles": result["muted_roles"],
"roles": result["roles"]}
self.member_cache.store_cache((member_id, guild_id), dict_items)
return self.member_cache.cache[member_id, guild_id]
# Release connection back to pool
finally:
await pool.release(conn)
# --------------------------------------------!End Cache Section!---------------------------------------------------
# --------------------------------------------!Starboard Section!---------------------------------------------------
def cache_store_starboard(self, guild_id, channel_id):
"""Storing starboard within cache"""
self.starboard_cache[guild_id] = {"channel_id": channel_id}
def get_starboard(self, guild_id):
"""Returning the starboard of the guild"""
starboard = self.starboard_cache.get(guild_id)
return starboard.get("channel_id") if starboard else None
def update_starboard(self, guild_id, channel_id):
"""Update the starboard channel"""
self.starboard_cache[guild_id]["channel_id"] = channel_id
def delete_starboard(self, guild_id):
"""Deleting the starboard of the guild"""
del self.starboard_cache[guild_id]
def delete_starboard_messages(self, guild_id):
# Array to store keys to be removed
keys_to_remove = []
# For every starboard message in cache
for (root_msg_id, guild_id) in self.starboard_messages_cache:
# if the guild_id passed in is equal to the guild_id within the cache
if guild_id == guild_id:
# Store key within array
keys_to_remove.append((root_msg_id, guild_id))
# Iterate through the array and then pop the keys from cache
for key in keys_to_remove:
self.starboard_messages_cache.pop(key)
def cache_store_starboard_message(self, root_message_id, guild_id, star_message_id):
"""Store the starboard messages within cache"""
self.starboard_messages_cache[root_message_id, guild_id] = {"star_message_id": star_message_id,
"stars": 1}
def update_starboard_message(self, root_message_id, guild_id, reactions):
"""Update the stored starboard message"""
self.starboard_messages_cache[root_message_id, guild_id]["stars"] = reactions
async def check_starboard_messages_cache(self, root_message_id, guild_id):
"""Check if the message is already in the cache"""
# Return value if message is already in the cache
if (root_message_id, guild_id) in self.starboard_messages_cache:
return self.starboard_messages_cache[root_message_id, guild_id]["star_message_id"], \
self.starboard_messages_cache[root_message_id, guild_id]["stars"]
else:
# Setup pool connection
pool = self.db
async with pool.acquire() as conn:
# Get the starboard row from the starboard_messages table
try:
select_query = """SELECT * FROM starboard_messages WHERE root_message_id = $1 AND guild_id = $2"""
result = await conn.fetchrow(select_query, root_message_id, guild_id)
# Catch errors
except asyncpg.PostgresError as e:
print(e)
# Store it in cache
else:
if result:
self.starboard_messages_cache[root_message_id, guild_id] = {"channel_id": result["channel_id"],
"star_message_id":
result["star_message_id"],
"stars": result["stars"]}
return self.starboard_messages_cache[root_message_id, guild_id]["star_message_id"], \
self.starboard_messages_cache[root_message_id, guild_id]["stars"]
else:
return None, 0
# Release connection back to pool
finally:
await pool.release(conn)
# --------------------------------------------!EndStarbard Section!-------------------------------------------------
# --------------------------------------------!Modmail Section!-----------------------------------------------------
def cache_store_modmail(self, guild_id, modmail_channel, message, modmail_logging_channel):
@ -507,6 +661,24 @@ class Bot(commands.Bot):
# Remove any/all members stored in cache from that guild
self.member_cache.remove_many(guild.id)
# Delete any starboard information upon leaving the guild
try:
delete = """DELETE FROM starboard WHERE guild_id = $1"""
rowcount = await conn.execute(delete, guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print(
f"PostGres Error: On Guild Remove Event Starboard/Starboard Messages Were Not Deleted For {guild.id}",
e)
# Delete all information about the starboard and any messages stored
else:
print(rowcount, f"Starboard deleted successfully from Guild {guild}")
if self.get_starboard(guild.id):
self.delete_starboard(guild.id)
self.delete_starboard_messages(guild.id)
# Release connection back to pool
await pool.release(conn)
@ -668,6 +840,8 @@ class Bot(commands.Bot):
# Get the modlogs channel (channel or none)
modlogs = self.get_modlog_for_guild(channel.guild.id)
# Get the starboard (record or none)
starboard = self.get_starboard(channel.guild.id)
# Get the modmail record - (normal and logging channels)
modmail_record = self.get_modmail(channel.guild.id)
@ -699,67 +873,56 @@ class Bot(commands.Bot):
finally:
await pool.release(conn)
# If modmail channels are deleted, delete the entire system
if channel.id == modmail_channel or channel.id == modmail_logging_channel:
# Set up pool connection
# Delete all of the starboard information when the channel is deleted from the guild
if channel.id == starboard:
# Setup pool connection
async with pool.acquire() as conn:
# Remove the moderatormail record from the database
# Delete any starboard information upon leaving the guild
try:
delete = """DELETE FROM moderatormail WHERE guild_id = $1"""
await conn.execute(delete, channel.guild.id)
delete = """DELETE FROM starboard WHERE guild_id = $1"""
rowcount = await conn.execute(delete, channel.guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: ModeratorMail Record Could Not Be Deleted for Guild {channel.guild.id}", e)
print(
f"PostGres Error: On Guild Remove Event Starboard/Starboard Messages Were Not Deleted For {channel.guild.id}",
e)
# Delete from cache
# Delete all information about the starboard and any messages stored
else:
self.delete_modmail(channel.guild.id)
print(rowcount, f"Starboard deleted successfully from Guild {channel.guild}")
if self.get_starboard(channel.guild.id):
self.delete_starboard(channel.guild.id)
self.delete_starboard_messages(channel.guild.id)
# Release connection back to pool
finally:
await pool.release(conn)
# --------------------------------------------!End Events Section!----------------------------------------------
async def check_cache(self, member_id, guild_id):
"""Checks if member is in the member cache"""
# Return key-value pair if member is already in the cache
if (member_id, guild_id) in self.member_cache.cache:
return self.member_cache.cache[member_id, guild_id]
else:
await pool.release(conn)
# Setup pool connection
pool = self.db
# If modmail channels are deleted, delete the entire system
if channel.id == modmail_channel or channel.id == modmail_logging_channel:
# Set up pool connection
async with pool.acquire() as conn:
# Get the author's/members row from the Members Table
# Remove the moderatormail record from the database
try:
select_query = """SELECT * FROM members WHERE member_id = $1 and guild_id = $2"""
result = await conn.fetchrow(select_query, member_id, guild_id)
delete = """DELETE FROM moderatormail WHERE guild_id = $1"""
await conn.execute(delete, channel.guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: Member {member_id} From Guild {guild_id}"
"Record Could Not Be Retrieved When Checking Cache", e)
print(f"PostGres Error: ModeratorMail Record Could Not Be Deleted for Guild {channel.guild.id}", e)
# Store it in cache
# Delete from cache
else:
dict_items = {"married": result["married"],
"married_date": result["married_date"],
"muted_roles": result["muted_roles"],
"roles": result["roles"]}
self.member_cache.store_cache((member_id, guild_id), dict_items)
return self.member_cache.cache[(member_id, guild_id)]
self.delete_modmail(channel.guild.id)
# Release connection back to pool
finally:
await pool.release(conn)
# --------------------------------------------!End Events Section!--------------------------------------------------
def execute(self):
"""Load the cogs and then run the bot"""

@ -0,0 +1,309 @@
# Ensō~Chan - A Multi Purpose Discord Bot That Has Everything Your Server Needs!
# Copyright (C) 2020 Hamothy
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import datetime
import asyncpg
from discord import TextChannel, Embed, NotFound
from discord.ext.commands import Cog, group, bot_has_permissions, has_permissions
async def send_starboard_and_update_db(self, payload, action):
"""Send the starboard embed and update database/cache"""
if (starboard := self.bot.get_starboard(payload.guild_id)) and payload.emoji.name == "":
message = await self.bot.get_channel(payload.channel_id).fetch_message(payload.message_id)
if not message.author.bot and payload.member.id != message.author.id:
channel = self.bot.get_channel(starboard)
msg_id, stars = await self.bot.check_starboard_messages_cache(message.id, payload.guild_id)
new_stars = stars + 1 if action == "added" else stars - 1
embed = Embed(title=f"Starred Message | {new_stars} :star:",
description=f"{message.content or 'View Attachment'}",
colour=message.author.colour,
timestamp=datetime.datetime.utcnow())
embed.set_author(name=message.author, icon_url=message.author.avatar_url)
embed.set_footer(text=f"ID: {message.id}")
embed.add_field(name="Original Message",
value=f"**Channel:** {message.channel.mention}\n[Jump To Message]({message.jump_url})",
inline=False)
if len(message.attachments):
embed.set_image(url=message.attachments[0].url)
if new_stars <= 0:
try:
star_message = await channel.fetch_message(msg_id)
await star_message.delete()
except NotFound:
pass
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Insert the starboard message in the database
try:
update_query = """DELETE FROM starboard_messages WHERE root_message_id = $1 AND guild_id = $2"""
await conn.execute(update_query, message.id, payload.guild_id)
# Catch errors
except asyncpg.PostgresError as e:
print(
f"PostGres Error: Starboard_Message Record Could Not Be Deleted For Guild {payload.guild_id}",
e)
# Update cache
else:
self.bot.delete_starboard_messages(payload.guild_id)
# Release connection back to pool
finally:
await pool.release(conn)
return
if not stars:
star_message = await channel.send(embed=embed)
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Insert the starboard message in the database
try:
update_query = """INSERT INTO starboard_messages (root_message_id, guild_id, star_message_id)
VALUES ($1, $2, $3)"""
await conn.execute(update_query, message.id, payload.guild_id, star_message.id)
# Catch errors
except asyncpg.PostgresError as e:
print(
f"PostGres Error: Starboard_Message Record Could Not Be Inserted For Guild {payload.guild_id}",
e)
# Update cache
else:
self.bot.cache_store_starboard_message(message.id, payload.guild_id, star_message.id)
# Release connection back to pool
finally:
await pool.release(conn)
else:
try:
star_message = await channel.fetch_message(msg_id)
await star_message.edit(embed=embed)
except NotFound:
pass
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Update the stars that the message has in the database
try:
update_query = """UPDATE starboard_messages
SET stars = $1
WHERE root_message_id = $2 AND guild_id = $3"""
await conn.execute(update_query, new_stars, message.id, payload.guild_id)
# Catch errors
except asyncpg.PostgresError as e:
print(
f"PostGres Error: Starboard_Message Record Could Not Be Updated For Guild {payload.guild_id}",
e)
# Update cache
else:
self.bot.update_starboard_message(message.id, payload.guild_id, new_stars)
# Release connection back to pool
finally:
await pool.release(conn)
else:
if action == "added":
await message.remove_reaction(payload.emoji, payload.member)
async def get_starboard_from_db(self, ctx, action):
"""Get the starboard record from DB"""
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Get the row of the guild from the starboard table
try:
select_query = """SELECT * FROM starboard WHERE guild_id = $1"""
result = await conn.fetchrow(select_query, ctx.guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print("PostGres Error: Starboard Record Could Not Be Retrieved For Starboard Setup", e)
# Throw error if the guild already exists
else:
if action == "setup" and result:
text = "**Starboard** Already Setup!" \
f"\nDo **{ctx.prefix}help starboard** to find out more!"
await self.bot.generate_embed(ctx, desc=text)
return None
elif (action == "update" or action == "delete") and not result:
text = "**Starboard** Not Setup!" \
f"\nDo **{ctx.prefix}help modlogs** to find out more!"
await self.bot.generate_embed(ctx, desc=text)
return None
# Release the connection back to the pool
finally:
await pool.release(conn)
return not None
# Set up the cog
class Starboard(Cog):
"""Starboard feature!"""
def __init__(self, bot):
self.bot = bot
@group(name="starboard", case_insensitive=True, usage="`<setup|update|delete>`")
@bot_has_permissions(embed_links=True)
@has_permissions(manage_guild=True)
async def starboard(self, ctx):
"""
Starboard! Let the community star messages!
"""
@starboard.command(name="setup")
@bot_has_permissions(embed_links=True)
async def sb_setup(self, ctx, starboard_channel: TextChannel):
"""
Setup Starboard
First Argument: Input Channel(Mention or ID) where starred messages will be sent
"""
if await get_starboard_from_db(self, ctx, "setup"):
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Insert the information about the starboard into database
try:
insert_query = """INSERT INTO starboard (guild_id, channel_id) VALUES ($1, $2)"""
await conn.execute(insert_query, ctx.guild.id, starboard_channel.id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: Starboard Record Could Not Be Inserted For Guild {ctx.guild.id}", e)
# Send confirmation message
else:
text = "**Starboard** is successfully set up!" \
f"\nRefer to **{ctx.prefix}help modmail** for more information"
await self.bot.generate_embed(ctx, desc=text)
# Store into cache
self.bot.cache_store_starboard(ctx.guild.id, starboard_channel.id)
# Release connection back into pool
finally:
await pool.release(conn)
@starboard.command(name="update")
@bot_has_permissions(embed_links=True)
async def sb_update(self, ctx, starboard_channel: TextChannel):
"""
Update the channel that the starred messages are sent to
You can Mention or use the Channel ID
"""
if await get_starboard_from_db(self, ctx, "update"):
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Update the starboard channel in the database
try:
update_query = """UPDATE starboard SET channel_id = $1 WHERE guild_id = $2"""
await conn.execute(update_query, starboard_channel.id, ctx.guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: Starboard Record Could Not Be Updated For Guild {ctx.guild.id}", e)
# Send confirmation that the channel has been updated
else:
text = "**Channel Updated**" \
f"\nNew Starred Messages will be sent to {starboard_channel.mention}"
await self.bot.generate_embed(ctx, desc=text)
# Update cache
self.bot.update_starboard(ctx.guild.id, starboard_channel.id)
# Release connection back to pool
finally:
await pool.release(conn)
@starboard.command(name="delete")
@bot_has_permissions(embed_links=True)
async def sb_delete(self, ctx):
"""Delete the starboard from the guild"""
if await get_starboard_from_db(self, ctx, "delete"):
# Setup up pool connection
pool = self.bot.db
async with pool.acquire() as conn:
# Remove the starboard record from the database
try:
delete_query = """DELETE FROM starboard WHERE guild_id = $1"""
await conn.execute(delete_query, ctx.guild.id)
# Catch errors
except asyncpg.PostgresError as e:
print(f"PostGres Error: Starboard Record Could Not Be Deleted for Guild {ctx.guild.id}", e)
# Sending confirmation message that the starboard has been deleted
else:
text = "**Starboard** successfully deleted!" \
f"\nDo **{ctx.prefix}help starboard** to find out more!"
await self.bot.generate_embed(ctx, desc=text)
# Delete from cache
self.bot.delete_starboard(ctx.guild.id)
# Release connection back to pool
finally:
await pool.release(conn)
@Cog.listener()
async def on_raw_reaction_remove(self, payload):
"""Removing reaction when a star is removed from the message"""
await send_starboard_and_update_db(self, payload, action="removed")
@Cog.listener()
async def on_raw_reaction_add(self, payload):
"""Listening for star reactions for any guilds that have starboard enabled"""
await send_starboard_and_update_db(self, payload, action="added")
def setup(bot):
bot.add_cog(Starboard(bot))
Loading…
Cancel
Save