From 328dd2a1bc14295381d20a3a08e43850e8ba0dbf Mon Sep 17 00:00:00 2001 From: Zevaryx Date: Thu, 5 Aug 2021 09:43:38 -0600 Subject: [PATCH] Migrate to mongoengine instead of custom ORM --- jarvis/__init__.py | 29 +- jarvis/cogs/admin/ban.py | 19 +- jarvis/cogs/admin/kick.py | 4 +- jarvis/cogs/admin/lock.py | 10 +- jarvis/cogs/admin/lockdown.py | 48 +-- jarvis/cogs/admin/mute.py | 15 +- jarvis/cogs/admin/purge.py | 17 +- jarvis/cogs/admin/roleping.py | 45 ++- jarvis/cogs/admin/warning.py | 26 +- jarvis/cogs/autoreact.py | 31 +- jarvis/cogs/ctc2.py | 20 +- jarvis/cogs/error.py | 4 +- jarvis/cogs/jokes.py | 56 ++-- jarvis/cogs/modlog/member.py | 85 +++-- jarvis/cogs/modlog/message.py | 10 +- jarvis/cogs/owner.py | 18 +- jarvis/cogs/remindme.py | 20 +- jarvis/cogs/rolegiver.py | 26 +- jarvis/cogs/settings.py | 16 +- jarvis/cogs/starboard.py | 36 ++- jarvis/cogs/verify.py | 12 +- jarvis/config.py | 11 +- jarvis/db/__init__.py | 14 - jarvis/db/models.py | 197 +++++++++++ jarvis/db/types.py | 594 ---------------------------------- jarvis/events/guild.py | 5 +- jarvis/events/member.py | 12 +- jarvis/events/message.py | 31 +- jarvis/tasks/unban.py | 33 +- jarvis/tasks/unlock.py | 20 +- jarvis/tasks/unmute.py | 22 +- jarvis/tasks/unwarn.py | 15 +- requirements.txt | 1 + 33 files changed, 542 insertions(+), 960 deletions(-) delete mode 100644 jarvis/db/__init__.py create mode 100644 jarvis/db/models.py delete mode 100644 jarvis/db/types.py diff --git a/jarvis/__init__.py b/jarvis/__init__.py index a10ba7c..65abcee 100644 --- a/jarvis/__init__.py +++ b/jarvis/__init__.py @@ -5,11 +5,11 @@ from discord import Intents from discord.ext import commands from discord.utils import find from discord_slash import SlashCommand +from mongoengine import connect from psutil import Process from jarvis import logo, tasks, utils from jarvis.config import get_config -from jarvis.db import DBManager from jarvis.events import guild, member, message if asyncio.get_event_loop().is_closed(): @@ -25,11 +25,9 @@ jarvis = commands.Bot( ) slash = SlashCommand(jarvis, sync_commands=True, sync_on_cog_reload=True) jarvis_self = Process() -__version__ = "1.9.6" +__version__ = "1.10.0" jconfig = get_config() -db = DBManager(jconfig.mongo["connect"]).mongo -jarvis_db = db[jconfig.mongo["database"]] @jarvis.event @@ -59,28 +57,39 @@ def run(ctx=None): global restart_ctx if ctx: restart_ctx = ctx + connect( + db="ctc2", + alias="ctc2", + authentication_source="admin", + **jconfig.mongo["connect"], + ) + connect( + db=jconfig.mongo["database"], + alias="main", + authentication_source="admin", + **jconfig.mongo["connect"], + ) + jconfig.get_db_config() for extension in utils.get_extensions(): jarvis.load_extension(extension) - config = get_config() print( " https://discord.com/api/oauth2/authorize?client_id=" + "{}&permissions=8&scope=bot%20applications.commands".format( - config.client_id + jconfig.client_id ) ) - jarvis.max_messages = config.max_messages + jarvis.max_messages = jconfig.max_messages tasks.init() # Add event listeners - if config.events: + if jconfig.events: listeners = [ guild.GuildEventHandler(jarvis), member.MemberEventHandler(jarvis), message.MessageEventHandler(jarvis), ] - - jarvis.run(config.token, bot=True, reconnect=True) + jarvis.run(jconfig.token, bot=True, reconnect=True) for cog in jarvis.cogs: session = getattr(cog, "_session", None) if session: diff --git a/jarvis/cogs/admin/ban.py b/jarvis/cogs/admin/ban.py index bf6b7d6..72f8f86 100644 --- a/jarvis/cogs/admin/ban.py +++ b/jarvis/cogs/admin/ban.py @@ -9,7 +9,7 @@ from discord_slash import SlashContext, cog_ext from discord_slash.model import ButtonStyle from discord_slash.utils.manage_commands import create_choice, create_option -from jarvis.db.types import Ban, Unban +from jarvis.db.models import Ban, Unban from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -40,7 +40,7 @@ class BanCog(CacheCog): type=type, duration=duration, active=active, - ).insert() + ).save() embed = build_embed( title="User Banned", @@ -68,7 +68,7 @@ class BanCog(CacheCog): guild=ctx.guild.id, admin=ctx.author.id, reason=reason, - ).insert() + ).save() embed = build_embed( title="User Unbanned", @@ -280,9 +280,9 @@ class BanCog(CacheCog): # We take advantage of the previous checks to save CPU cycles if not discord_ban_info: if isinstance(user, int): - database_ban_info = Ban.get( + database_ban_info = Ban.objects( guild=ctx.guild.id, user=user, active=True - ) + ).first() else: search = { "guild": ctx.guild.id, @@ -291,7 +291,7 @@ class BanCog(CacheCog): } if discrim: search["discrim"] = discrim - database_ban_info = Ban.get(**search) + database_ban_info = Ban.objects(**search).first() if not discord_ban_info and not database_ban_info: await ctx.send(f"Unable to find user {orig_user}", hidden=True) @@ -308,7 +308,7 @@ class BanCog(CacheCog): ) else: database_ban_info.active = False - database_ban_info.update() + database_ban_info.save() _ = Unban( user=database_ban_info.user, username=database_ban_info.username, @@ -316,7 +316,7 @@ class BanCog(CacheCog): guild=ctx.guild.id, admin=ctx.author.id, reason=reason, - ).insert() + ).save() await ctx.send( "Unable to find user in Discord, " + "but removed entry from database." @@ -371,8 +371,7 @@ class BanCog(CacheCog): search["active"] = True if type > 0: search["type"] = types[type] - bans = Ban.get_many(**search) - bans.sort(key=lambda x: x.created_at, reverse=True) + bans = Ban.objects(**search).order_by("-created_at") db_bans = [] fields = [] for ban in bans: diff --git a/jarvis/cogs/admin/kick.py b/jarvis/cogs/admin/kick.py index 32184d2..2f3f06b 100644 --- a/jarvis/cogs/admin/kick.py +++ b/jarvis/cogs/admin/kick.py @@ -2,7 +2,7 @@ from discord import User from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Kick +from jarvis.db.models import Kick from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -85,4 +85,4 @@ async def _kick(self, ctx: SlashContext, user: User, reason=None): reason=reason, admin=ctx.author.id, guild=ctx.guild.id, - ).insert() + ).save() diff --git a/jarvis/cogs/admin/lock.py b/jarvis/cogs/admin/lock.py index c8c557f..2762bcd 100644 --- a/jarvis/cogs/admin/lock.py +++ b/jarvis/cogs/admin/lock.py @@ -5,7 +5,7 @@ from discord.ext import commands from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Lock +from jarvis.db.models import Lock from jarvis.utils.cachecog import CacheCog from jarvis.utils.permissions import admin_or_permissions @@ -97,7 +97,7 @@ class LockCog(CacheCog): admin=ctx.author.id, reason=reason, duration=duration, - ).insert() + ).save() await ctx.send(f"{channel.mention} locked for {duration} minute(s)") @cog_ext.cog_slash( @@ -120,7 +120,9 @@ class LockCog(CacheCog): ): if not channel: channel = ctx.channel - lock = Lock.get(guild=ctx.guild.id, channel=channel.id, active=True) + lock = Lock.objects( + guild=ctx.guild.id, channel=channel.id, active=True + ).first() if not lock: await ctx.send(f"{channel.mention} not locked.", hidden=True) return @@ -130,5 +132,5 @@ class LockCog(CacheCog): except Exception: continue # Just continue on error lock.active = False - lock.update() + lock.save() await ctx.send(f"{channel.mention} unlocked") diff --git a/jarvis/cogs/admin/lockdown.py b/jarvis/cogs/admin/lockdown.py index 20a7f43..39f4883 100644 --- a/jarvis/cogs/admin/lockdown.py +++ b/jarvis/cogs/admin/lockdown.py @@ -1,13 +1,10 @@ from datetime import datetime -import pymongo from discord.ext import commands from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.config import get_config -from jarvis.db import DBManager -from jarvis.db.types import Lock +from jarvis.db.models import Lock from jarvis.utils.cachecog import CacheCog from jarvis.utils.permissions import admin_or_permissions @@ -15,10 +12,6 @@ from jarvis.utils.permissions import admin_or_permissions class LockdownCog(CacheCog): def __init__(self, bot: commands.Bot): super().__init__(bot) - config = get_config() - self.db = DBManager(config.mongo["connect"]).mongo[ - config.mongo["database"] - ] @cog_ext.cog_subcommand( base="lockdown", @@ -63,20 +56,18 @@ class LockdownCog(CacheCog): except Exception: continue # Just continue on error updates.append( - pymongo.InsertOne( - { - "channel": channel.id, - "guild": ctx.guild.id, - "admin": ctx.author.id, - "reason": reason, - "duration": duration, - "active": True, - "created_at": datetime.utcnow(), - } + Lock( + channel=channel.id, + guild=ctx.guild.id, + admin=ctx.author.id, + reason=reason, + duration=duration, + active=True, + created_at=datetime.utcnow(), ) ) if updates: - self.db.locks.bulk_write(updates) + Lock.objects().insert(updates) await ctx.send(f"Server locked for {duration} minute(s)") @cog_ext.cog_subcommand( @@ -91,8 +82,8 @@ class LockdownCog(CacheCog): ): channels = ctx.guild.channels roles = ctx.guild.roles - updates = [] - locks = Lock.get_many(guild=ctx.guild.id, active=True) + update = False + locks = Lock.objects(guild=ctx.guild.id, active=True) if not locks: await ctx.send("No lockdown detected.", hidden=True) return @@ -103,16 +94,9 @@ class LockdownCog(CacheCog): await self._unlock_channel(channel, role, ctx.author) except Exception: continue # Just continue on error - updates.append( - pymongo.UpdateOne( - { - "channel": channel.id, - "guild": ctx.guild.id, - "admin": ctx.author.id, - }, - {"$set": {"active": False}}, - ) + update = True + if update: + Lock.objects(guild=ctx.guild.id, active=True).update( + set__active=False ) - if updates: - self.db.locks.bulk_write(updates) await ctx.send("Server unlocked") diff --git a/jarvis/cogs/admin/mute.py b/jarvis/cogs/admin/mute.py index 9a4361c..7519c4f 100644 --- a/jarvis/cogs/admin/mute.py +++ b/jarvis/cogs/admin/mute.py @@ -4,7 +4,7 @@ from discord.utils import get from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Mute, Setting +from jarvis.db.models import Mute, Setting from jarvis.utils import build_embed from jarvis.utils.field import Field from jarvis.utils.permissions import admin_or_permissions @@ -73,7 +73,7 @@ class MuteCog(commands.Cog): guild=ctx.guild.id, duration=duration, active=True if duration >= 0 else False, - ).insert() + ).save() embed = build_embed( title="User Muted", @@ -102,7 +102,9 @@ class MuteCog(commands.Cog): ) @admin_or_permissions(mute_members=True) async def _unmute(self, ctx: SlashContext, user: Member): - mute_setting = Setting.get(guild=ctx.guild.id, setting="mute") + mute_setting = Setting.objects( + guild=ctx.guild.id, setting="mute" + ).first() if not mute_setting: await ctx.send( "Please configure a mute role with " @@ -118,10 +120,9 @@ class MuteCog(commands.Cog): await ctx.send("User is not muted.", hidden=True) return - mutes = Mute.get_many(guild=ctx.guild.id, user=user.id) - for mute in mutes: - mute.active = False - mute.update() + _ = Mute.objects(guild=ctx.guild.id, user=user.id).update( + set__active=False + ) embed = build_embed( title="User Unmwaruted", description=f"{user.mention} has been unmuted", diff --git a/jarvis/cogs/admin/purge.py b/jarvis/cogs/admin/purge.py index 3fc6b30..0fef3d1 100644 --- a/jarvis/cogs/admin/purge.py +++ b/jarvis/cogs/admin/purge.py @@ -39,7 +39,7 @@ class PurgeCog(commands.Cog): guild=ctx.guild.id, admin=ctx.author.id, count=amount, - ).insert() + ).save() @cog_ext.cog_subcommand( base="autopurge", @@ -73,17 +73,18 @@ class PurgeCog(commands.Cog): elif delay > 300: await ctx.send("Delay must be < 5 minutes", hidden=True) return - autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id) + autopurge = Autopurge.objects( + guild=ctx.guild.id, channel=channel.id + ).first() if autopurge: await ctx.send("Autopurge already exists.", hidden=True) return - autopurge = Autopurge( + _ = Autopurge( guild=ctx.guild.id, channel=channel.id, admin=ctx.author.id, delay=delay, - ) - autopurge.insert() + ).save() await ctx.send( f"Autopurge set up on {channel.mention}, " + f"delay is {delay} seconds" @@ -104,7 +105,7 @@ class PurgeCog(commands.Cog): ) @admin_or_permissions(manage_messages=True) async def _autopurge_remove(self, ctx: SlashContext, channel: TextChannel): - autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id) + autopurge = Autopurge.objects(guild=ctx.guild.id, channel=channel.id) if not autopurge: await ctx.send("Autopurge does not exist.", hidden=True) return @@ -134,12 +135,12 @@ class PurgeCog(commands.Cog): async def _autopurge_update( self, ctx: SlashContext, channel: TextChannel, delay: int ): - autopurge = Autopurge.get(guild=ctx.guild.id, channel=channel.id) + autopurge = Autopurge.objects(guild=ctx.guild.id, channel=channel.id) if not autopurge: await ctx.send("Autopurge does not exist.", hidden=True) return autopurge.delay = delay - autopurge.update() + autopurge.save() await ctx.send( f"Autopurge delay updated to {delay} seconds on {channel.mention}." ) diff --git a/jarvis/cogs/admin/roleping.py b/jarvis/cogs/admin/roleping.py index 283f3eb..14f8ca4 100644 --- a/jarvis/cogs/admin/roleping.py +++ b/jarvis/cogs/admin/roleping.py @@ -2,12 +2,11 @@ from datetime import datetime, timedelta from ButtonPaginator import Paginator from discord import Member, Role -from discord.ext import commands from discord_slash import SlashContext, cog_ext from discord_slash.model import ButtonStyle from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Roleping +from jarvis.db.models import Roleping from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -33,23 +32,19 @@ class RolepingCog(CacheCog): ) @admin_or_permissions(manage_guild=True) async def _roleping_add(self, ctx: SlashContext, role: Role): - roleping = Roleping.get(guild=ctx.guild.id, role=role.id) - if not roleping: - roleping = Roleping( - role=role.id, - guild=ctx.guild.id, - admin=ctx.author.id, - active=True, - bypass={"roles": [], "users": []}, - ) - - else: + roleping = Roleping.objects(guild=ctx.guild.id, role=role.id).first() + if roleping: await ctx.send( f"Role `{role.name}` already in roleping.", hidden=True ) return - - roleping.insert() + _ = Roleping( + role=role.id, + guild=ctx.guild.id, + admin=ctx.author.id, + active=True, + bypass={"roles": [], "users": []}, + ).save() await ctx.send(f"Role `{role.name}` added to roleping.") @cog_ext.cog_subcommand( @@ -67,7 +62,7 @@ class RolepingCog(CacheCog): ) @admin_or_permissions(manage_guild=True) async def _roleping_remove(self, ctx: SlashContext, role: Role): - roleping = Roleping.get(guild=ctx.guild.id, role=role.id) + roleping = Roleping.objects(guild=ctx.guild.id, role=role.id) if not roleping: await ctx.send("Roleping does not exist", hidden=True) return @@ -91,7 +86,7 @@ class RolepingCog(CacheCog): ) return - rolepings = Roleping.get_many(guild=ctx.guild.id) + rolepings = Roleping.objects(guild=ctx.guild.id) if not rolepings: await ctx.send("No rolepings configured", hidden=True) return @@ -199,7 +194,7 @@ class RolepingCog(CacheCog): async def _roleping_bypass_user( self, ctx: SlashContext, user: Member, rping: Role ): - roleping = Roleping.get(guild=ctx.guild.id, role=rping.id) + roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first() if not roleping: await ctx.send( f"Roleping not configured for {rping.mention}", hidden=True @@ -231,7 +226,7 @@ class RolepingCog(CacheCog): return roleping.bypass["users"].append(user.id) - roleping.update() + roleping.save() await ctx.send( f"{user.nick or user.name} user bypass added for `{rping.name}`" ) @@ -262,7 +257,7 @@ class RolepingCog(CacheCog): async def _roleping_bypass_role( self, ctx: SlashContext, role: Role, rping: Role ): - roleping = Roleping.get(guild=ctx.guild.id, role=rping.id) + roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first() if not roleping: await ctx.send( f"Roleping not configured for {rping.mention}", hidden=True @@ -282,7 +277,7 @@ class RolepingCog(CacheCog): return roleping.bypass["roles"].append(role.id) - roleping.update() + roleping.save() await ctx.send(f"{role.name} role bypass added for `{rping.name}`") @cog_ext.cog_subcommand( @@ -311,7 +306,7 @@ class RolepingCog(CacheCog): async def _roleping_restore_user( self, ctx: SlashContext, user: Member, rping: Role ): - roleping = Roleping.get(guild=ctx.guild.id, role=rping.id) + roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first() if not roleping: await ctx.send( f"Roleping not configured for {rping.mention}", hidden=True @@ -323,7 +318,7 @@ class RolepingCog(CacheCog): return roleping.bypass["users"].delete(user.id) - roleping.update() + roleping.save() await ctx.send( f"{user.nick or user.name} user bypass removed for `{rping.name}`" ) @@ -354,7 +349,7 @@ class RolepingCog(CacheCog): async def _roleping_restore_role( self, ctx: SlashContext, role: Role, rping: Role ): - roleping = Roleping.get(guild=ctx.guild.id, role=rping.id) + roleping = Roleping.objects(guild=ctx.guild.id, role=rping.id).first() if not roleping: await ctx.send( f"Roleping not configured for {rping.mention}", hidden=True @@ -374,5 +369,5 @@ class RolepingCog(CacheCog): return roleping.bypass["roles"].append(role.id) - roleping.update() + roleping.save() await ctx.send(f"{role.name} role bypass added for `{rping.name}`") diff --git a/jarvis/cogs/admin/warning.py b/jarvis/cogs/admin/warning.py index cb424eb..e154ea5 100644 --- a/jarvis/cogs/admin/warning.py +++ b/jarvis/cogs/admin/warning.py @@ -2,12 +2,11 @@ from datetime import datetime, timedelta from ButtonPaginator import Paginator from discord import User -from discord.ext import commands from discord_slash import SlashContext, cog_ext from discord_slash.model import ButtonStyle from discord_slash.utils.manage_commands import create_choice, create_option -from jarvis.db.types import MongoSort, Warning +from jarvis.db.models import Warning from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -63,7 +62,7 @@ class WarningCog(CacheCog): guild=ctx.guild.id, duration=duration, active=True, - ).insert() + ).save() fields = [Field("Reason", reason, False)] embed = build_embed( title="Warning", @@ -112,19 +111,20 @@ class WarningCog(CacheCog): hidden=True, ) return - warnings = Warning.get_many( + warnings = Warning.objects( user=user.id, guild=ctx.guild.id, - sort=MongoSort(direction="desc", key="created_at"), - ) - active_warns = list(filter(lambda x: x.active, warnings)) + ).order_by("-created_at") + active_warns = Warning.objects( + user=user.id, guild=ctx.guild.id, active=False + ).order_by("-created_at") pages = [] if active: - if len(active_warns) == 0: + if active_warns.count() == 0: embed = build_embed( title="Warnings", - description=f"{len(warnings)} total | 0 currently active", + description=f"{warnings.count()} total | 0 currently active", fields=[], ) embed.set_author(name=user.name, icon_url=user.avatar_url) @@ -151,8 +151,8 @@ class WarningCog(CacheCog): for i in range(0, len(fields), 5): embed = build_embed( title="Warnings", - description=f"{len(warnings)} total | " - + f"{len(active_warns)} currently active", + description=f"{warnings.count()} total | " + + f"{active_warns.count()} currently active", fields=fields[i : i + 5], ) embed.set_author( @@ -179,8 +179,8 @@ class WarningCog(CacheCog): for i in range(0, len(fields), 5): embed = build_embed( title="Warnings", - description=f"{len(warnings)} total | " - + f"{len(active_warns)} currently active", + description=f"{warnings.count()} total | " + + f"{active_warns.count()} currently active", fields=fields[i : i + 5], ) embed.set_author( diff --git a/jarvis/cogs/autoreact.py b/jarvis/cogs/autoreact.py index 10aa5b9..44b9e82 100644 --- a/jarvis/cogs/autoreact.py +++ b/jarvis/cogs/autoreact.py @@ -7,7 +7,7 @@ from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option from jarvis.data.unicode import emoji_list -from jarvis.db.types import Autoreact +from jarvis.db.models import Autoreact from jarvis.utils.permissions import admin_or_permissions @@ -34,20 +34,21 @@ class AutoReactCog(commands.Cog): if not isinstance(channel, TextChannel): await ctx.send("Channel must be a text channel", hidden=True) return - exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) + exists = Autoreact.objects( + guild=ctx.guild.id, channel=channel.id + ).first() if exists: await ctx.send( f"Autoreact already exists for {channel.mention}.", hidden=True ) return - autoreact = Autoreact( + _ = Autoreact( guild=ctx.guild.id, channel=channel.id, reactions=[], admin=ctx.author.id, - ) - autoreact.insert() + ).save() await ctx.send(f"Autoreact created for {channel.mention}!") @cog_ext.cog_subcommand( @@ -65,7 +66,9 @@ class AutoReactCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _autoreact_delete(self, ctx, channel: TextChannel): - exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id).delete() + exists = Autoreact.objects( + guild=ctx.guild.id, channel=channel.id + ).delete() if exists: await ctx.send(f"Autoreact removed from {channel.mention}") else: @@ -111,7 +114,9 @@ class AutoReactCog(commands.Cog): "Please use a custom emote from this server.", hidden=True ) return - exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) + exists = Autoreact.objects( + guild=ctx.guild.id, channel=channel.id + ).first() if not exists: await ctx.send( "Please create autoreact first with " @@ -132,7 +137,7 @@ class AutoReactCog(commands.Cog): ) return exists.reactions.append(emote) - exists.update() + exists.save() await ctx.send(f"Added {emote} to {channel.mention} autoreact.") @cog_ext.cog_subcommand( @@ -156,7 +161,9 @@ class AutoReactCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _autoreact_remove(self, ctx, channel: TextChannel, emote: str): - exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) + exists = Autoreact.objects( + guild=ctx.guild.id, channel=channel.id + ).first() if not exists: await ctx.send( "Please create autoreact first with " @@ -171,7 +178,7 @@ class AutoReactCog(commands.Cog): ) return exists.reactions.remove(emote) - exists.update() + exists.save() await ctx.send(f"Removed {emote} from {channel.mention} autoreact.") @cog_ext.cog_subcommand( @@ -189,7 +196,9 @@ class AutoReactCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _autoreact_list(self, ctx, channel: TextChannel): - exists = Autoreact.get(guild=ctx.guild.id, channel=channel.id) + exists = Autoreact.objects( + guild=ctx.guild.id, channel=channel.id + ).first() if not exists: await ctx.send( "Please create autoreact first with " diff --git a/jarvis/cogs/ctc2.py b/jarvis/cogs/ctc2.py index 30dc32e..022484b 100644 --- a/jarvis/cogs/ctc2.py +++ b/jarvis/cogs/ctc2.py @@ -11,8 +11,7 @@ from discord.utils import find from discord_slash import SlashContext, cog_ext from discord_slash.model import ButtonStyle -from jarvis.config import get_config -from jarvis.db import DBManager +from jarvis.db.models import Guess from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -29,8 +28,6 @@ invites = re.compile( class CTCCog(CacheCog): def __init__(self, bot): super().__init__(bot) - mconf = get_config().mongo - self.db = DBManager(mconf["connect"]).mongo self._session = aiohttp.ClientSession() self.url = "https://completethecodetwo.cards/pw" @@ -73,7 +70,7 @@ class CTCCog(CacheCog): hidden=True, ) return - guessed = self.db.ctc2.guesses.find_one({"guess": guess}) + guessed = Guess.objects(guess=guess).first() if guessed: await ctx.send("Already guessed, dipshit.", hidden=True) return @@ -86,9 +83,7 @@ class CTCCog(CacheCog): correct = True else: await ctx.send("Nope.", hidden=True) - self.db.ctc2.guesses.insert_one( - {"guess": guess, "user": ctx.author.id, "correct": correct} - ) + _ = Guess(guess=guess, user=ctx.author.id, correct=correct).save() @cog_ext.cog_subcommand( base="ctc2", @@ -107,17 +102,16 @@ class CTCCog(CacheCog): hidden=True, ) return - guesses = self.db.ctc2.guesses.find().sort( - [("correct", pymongo.DESCENDING), ("_id", pymongo.DESCENDING)] - ) + + guesses = Guess.objects().order_by("-correct", "-id") fields = [] for guess in guesses: user = ctx.guild.get_member(guess["user"]) if not user: - user = self.bot.fetch_user(guess["user"]) + user = await self.bot.fetch_user(guess["user"]) if not user: user = "[redacted]" - if isinstance(user, User) or isinstance(user, Member): + if isinstance(user, (Member, User)): user = user.name + "#" + user.discriminator name = "Correctly" if guess["correct"] else "Incorrectly" name += " guessed by: " + user diff --git a/jarvis/cogs/error.py b/jarvis/cogs/error.py index 9b03007..185ede0 100644 --- a/jarvis/cogs/error.py +++ b/jarvis/cogs/error.py @@ -1,6 +1,8 @@ from discord.ext import commands from discord_slash import SlashContext +from jarvis import slash + class ErrorHandlerCog(commands.Cog): def __init__(self, bot): @@ -40,7 +42,7 @@ class ErrorHandlerCog(commands.Cog): f"Error processing command:\n```{error}```", hidden=True, ) - ctx.slash.commands[ctx.command].reset_cooldown(ctx) + slash.commands[ctx.command].reset_cooldown(ctx) def setup(bot): diff --git a/jarvis/cogs/jokes.py b/jarvis/cogs/jokes.py index 240bcad..d153f4f 100644 --- a/jarvis/cogs/jokes.py +++ b/jarvis/cogs/jokes.py @@ -8,8 +8,7 @@ from discord.ext import commands from discord_slash import cog_ext import jarvis -from jarvis.config import get_config -from jarvis.db import DBManager +from jarvis.db.models import Joke from jarvis.utils import build_embed from jarvis.utils.field import Field @@ -23,10 +22,6 @@ class JokeCog(commands.Cog): def __init__(self, bot): self.bot = bot - config = get_config() - self.db = DBManager(config.mongo["connect"]).mongo[ - config.mongo["database"] - ] # TODO: Make this a command group with subcommands async def _joke(self, ctx, id: str = None): @@ -36,48 +31,35 @@ class JokeCog(commands.Cog): return # TODO: Add this as a parameter that can be passed in threshold = 500 # Minimum score - coll = self.db.jokes result = None if id: - result = coll.find_one({"id": id}) + result = Joke.objects(rid=id).first() else: - result = list( - coll.aggregate( - [ - {"$match": {"score": {"$gt": threshold}}}, - {"$sample": {"size": 1}}, - ] - ) - )[0] - while result["body"] in ["[removed]", "[deleted]"]: - result = list( - coll.aggregate( - [ - {"$match": {"score": {"$gt": threshold}}}, - {"$sample": {"size": 1}}, - ] - ) - )[0] + pipeline = [ + {"$match": {"score": {"$gt": threshold}}}, + {"$sample": {"size": 1}}, + ] + result = Joke.objects().aggregate(pipeline) + while result.body in ["[removed]", "[deleted]"]: + result = Joke.objects().aggregate(pipeline) if result is None: await ctx.send( "Humor module failed. Please try again later.", hidden=True ) return - emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result["body"]) + emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result.body) for match in emotes: - result["body"] = result["body"].replace( - match, html.unescape(match) - ) - emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result["title"]) + result.body = result.body.replace(match, html.unescape(match)) + emotes = re.findall(r"(&#x[a-fA-F0-9]*;)", result.title) for match in emotes: - result["title"] = result["title"].replace( + result.title = result.title.replace( match, html.unescape(match) ) body_chunks = [] body = "" - for word in result["body"].split(" "): + for word in result.body.split(" "): if len(body) + 1 + len(word) > 1024: body_chunks.append(Field("​", body, False)) body = "" @@ -89,7 +71,7 @@ class JokeCog(commands.Cog): body += " " + word desc = "" - title = result["title"] + title = result.title if len(title) > 256: new_title = "" limit = False @@ -106,18 +88,18 @@ class JokeCog(commands.Cog): body_chunks.append(Field("​", body, False)) fields = body_chunks - fields.append(Field("Score", result["score"])) + fields.append(Field("Score", result.score)) # Field( # "Created At", # str(datetime.fromtimestamp(result["created_utc"])), # ), - fields.append(Field("ID", result["id"])) + fields.append(Field("ID", result.id)) embed = build_embed( title=title, description=desc, fields=fields, - url=f"https://reddit.com/r/jokes/comments/{result['id']}", - timestamp=datetime.fromtimestamp(result["created_utc"]), + url=f"https://reddit.com/r/jokes/comments/{result.rid}", + timestamp=datetime.fromtimestamp(result.created_utc), ) await ctx.send(embed=embed) except Exception: diff --git a/jarvis/cogs/modlog/member.py b/jarvis/cogs/modlog/member.py index 150d853..9472f52 100644 --- a/jarvis/cogs/modlog/member.py +++ b/jarvis/cogs/modlog/member.py @@ -7,7 +7,7 @@ from discord.utils import find from jarvis.cogs.modlog.utils import get_latest_log, modlog_embed from jarvis.config import get_config -from jarvis.db.types import Ban, Kick, MongoSort, Mute, Setting +from jarvis.db.models import Ban, Kick, Mute, Setting, Unban from jarvis.utils import build_embed from jarvis.utils.field import Field @@ -18,7 +18,7 @@ class ModlogMemberCog(commands.Cog): @commands.Cog.listener() async def on_member_ban(self, guild: discord.Guild, user: discord.User): - modlog = Setting.get(guild=guild.id, setting="modlog") + modlog = Setting.objects(guild=guild.id, setting="modlog").first() if modlog: channel = guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log @@ -31,11 +31,14 @@ class ModlogMemberCog(commands.Cog): log: discord.AuditLogEntry = get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - ban = Ban.get( - guild=guild.id, - user=user.id, - active=True, - sort=MongoSort(key="created_at", direction="desc"), + ban = ( + Ban.objects( + guild=guild.id, + user=user.id, + active=True, + ) + .sort("-created_at") + .first() ) admin = guild.get_member(ban.admin) embed = modlog_embed( @@ -50,7 +53,7 @@ class ModlogMemberCog(commands.Cog): @commands.Cog.listener() async def on_member_unban(self, guild: discord.Guild, user: discord.User): - modlog = Setting.get(guild=guild.id, setting="modlog") + modlog = Setting.objects(guild=guild.id, setting="modlog").first() if modlog: channel = guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log @@ -63,13 +66,15 @@ class ModlogMemberCog(commands.Cog): log: discord.AuditLogEntry = get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - ban = Ban.get( - guild=guild.id, - user=user.id, - active=True, - sort=MongoSort(key="created_at", direction="desc"), + unban = ( + Unban.objects( + guild=guild.id, + user=user.id, + ) + .sort("-created_at") + .first() ) - admin = guild.get_member(ban.admin) + admin = guild.get_member(unban.admin) embed = modlog_embed( user, admin, @@ -82,7 +87,7 @@ class ModlogMemberCog(commands.Cog): @commands.Cog.listener() async def on_member_remove(self, user: discord.User): - modlog = Setting.get(guild=user.guild.id, setting="modlog") + modlog = Setting.objects(guild=user.guild.id, setting="modlog").first() if modlog: channel = user.guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log @@ -95,10 +100,13 @@ class ModlogMemberCog(commands.Cog): log: discord.AuditLogEntry = get_latest_log(auditlog, user) admin: discord.User = log.user if admin.id == get_config().client_id: - kick = Kick.get( - guild=user.guild.id, - user=user.id, - sort=MongoSort(key="created_at", direction="desc"), + kick = ( + Kick.objects( + guild=user.guild.id, + user=user.id, + ) + .sort("-created_at") + .first() ) admin = user.guild.get_member(kick.admin) embed = modlog_embed( @@ -122,11 +130,14 @@ class ModlogMemberCog(commands.Cog): log: discord.AuditLogEntry = get_latest_log(auditlog, before) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = Mute.get( - guild=before.guild.id, - user=before.id, - active=True, - sort=MongoSort(key="created_at", direction="desc"), + mute = ( + Mute.objects( + guild=before.guild.id, + user=before.id, + active=True, + ) + .sort("-created_at") + .first() ) admin = before.guild.get_member(mute.admin) return modlog_embed( @@ -148,13 +159,15 @@ class ModlogMemberCog(commands.Cog): log: discord.AuditLogEntry = get_latest_log(auditlog, before) admin: discord.User = log.user if admin.id == get_config().client_id: - mute = Mute.get( - guild=before.guild.id, - user=before.id, - active=True, - sort=MongoSort(key="created_at", direction="desc"), + mute = ( + Mute.objects( + guild=before.guild.id, + user=before.id, + active=True, + ) + .sort("-created_at") + .first() ) - mute = Mute(**mute) admin = before.guild.get_member(mute.admin) return modlog_embed( member=before, @@ -214,13 +227,19 @@ class ModlogMemberCog(commands.Cog): async def on_member_update( self, before: discord.User, after: discord.User ): - modlog = Setting.get(guild=before.guild.id, setting="modlog") + modlog = Setting.objects( + guild=before.guild.id, setting="modlog" + ).first() if modlog: channel = after.guild.get_channel(modlog.value) await asyncio.sleep(0.5) # Need to wait for audit log embed = None - mute = Setting.get(guild=before.guild.id, setting="mute") - verified = Setting.get(guild=before.guild.id, setting="verified") + mute = Setting.objects( + guild=before.guild.id, setting="mute" + ).first() + verified = Setting.objects( + guild=before.guild.id, setting="verified" + ).first() if mute and before.guild.get_role(mute.value) in after.roles: embed = await self.process_mute(before, after) elif mute and before.guild.get_role(mute.value) in before.roles: diff --git a/jarvis/cogs/modlog/message.py b/jarvis/cogs/modlog/message.py index 5034041..52b06c1 100644 --- a/jarvis/cogs/modlog/message.py +++ b/jarvis/cogs/modlog/message.py @@ -1,7 +1,7 @@ import discord from discord.ext import commands -from jarvis.db.types import Setting +from jarvis.db.models import Setting from jarvis.utils import build_embed from jarvis.utils.field import Field @@ -15,7 +15,9 @@ class ModlogMessageCog(commands.Cog): self, before: discord.Message, after: discord.Message ): if not before.author.bot: - modlog = Setting.get(guild=after.guild.id, setting="modlog") + modlog = Setting.objects( + guild=after.guild.id, setting="modlog" + ).first() if modlog: if before.content == after.content or before.content is None: return @@ -53,7 +55,9 @@ class ModlogMessageCog(commands.Cog): @commands.Cog.listener() async def on_message_delete(self, message: discord.Message): - modlog = Setting.get(guild=message.guild.id, setting="modlog") + modlog = Setting.objects( + guild=message.guild.id, setting="modlog" + ).first() if modlog: fields = [ Field("Original Message", message.content or "N/A", False) diff --git a/jarvis/cogs/owner.py b/jarvis/cogs/owner.py index c828339..28c2971 100644 --- a/jarvis/cogs/owner.py +++ b/jarvis/cogs/owner.py @@ -10,7 +10,7 @@ from discord.ext import commands import jarvis from jarvis.config import reload_config -from jarvis.db.types import Config +from jarvis.db.models import Config from jarvis.utils import update from jarvis.utils.permissions import user_is_bot_admin @@ -24,7 +24,7 @@ class OwnerCog(commands.Cog): def __init__(self, bot): self.bot = bot - self.admins = Config.get(key="admins") + self.admins = Config.objects(key="admins").first() @commands.command(name="load", hidden=True) @user_is_bot_admin() @@ -32,7 +32,7 @@ class OwnerCog(commands.Cog): info = await self.bot.application_info() if ( ctx.message.author == info.owner - or ctx.message.author.id in self.admins + or ctx.message.author.id in self.admins.value ): try: if "jarvis.cogs." not in cog: @@ -58,7 +58,7 @@ class OwnerCog(commands.Cog): info = await self.bot.application_info() if ( ctx.message.author == info.owner - or ctx.message.author.id in self.admins + or ctx.message.author.id in self.admins.value ): try: if "jarvis.cogs." not in cog: @@ -84,7 +84,7 @@ class OwnerCog(commands.Cog): info = await self.bot.application_info() if ( ctx.message.author == info.owner - or ctx.message.author.id in self.admins + or ctx.message.author.id in self.admins.value ): try: if "jarvis.cogs." not in cog: @@ -118,7 +118,7 @@ class OwnerCog(commands.Cog): info = await self.bot.application_info() if ( ctx.message.author == info.owner - or ctx.message.author.id in self.admins + or ctx.message.author.id in self.admins.value ): await ctx.send("Restarting core systems...") if isinstance(ctx.channel, discord.channel.DMChannel): @@ -141,7 +141,7 @@ class OwnerCog(commands.Cog): info = await self.bot.application_info() if ( ctx.message.author == info.owner - or ctx.message.author.id in self.admins + or ctx.message.author.id in self.admins.value ): await ctx.send("Updating core systems...") status = update() @@ -189,7 +189,7 @@ class OwnerCog(commands.Cog): await ctx.send(f"{user.mention} is already an admin.") return self.admins.value.append(user.id) - self.admins.update() + self.admins.save() reload_config() await ctx.send( f"{user.mention} is now an admin. Use this power carefully." @@ -202,7 +202,7 @@ class OwnerCog(commands.Cog): await ctx.send(f"{user.mention} is not an admin.") return self.admins.value.remove(user.id) - self.admins.update() + self.admins.save() reload_config() await ctx.send(f"{user.mention} is no longer an admin.") diff --git a/jarvis/cogs/remindme.py b/jarvis/cogs/remindme.py index 550b90e..cb770fc 100644 --- a/jarvis/cogs/remindme.py +++ b/jarvis/cogs/remindme.py @@ -16,7 +16,7 @@ from discord_slash.utils.manage_components import ( wait_for_component, ) -from jarvis.db.types import Reminder +from jarvis.db.models import Reminder from jarvis.utils import build_embed from jarvis.utils.cachecog import CacheCog from jarvis.utils.field import Field @@ -124,8 +124,8 @@ class RemindmeCog(CacheCog): ) return - reminders = Reminder.get_active(user=ctx.author.id) - if len(reminders) >= 5: + reminders = Reminder.objects(user=ctx.author.id, active=True).count() + if reminders >= 5: await ctx.send( "You already have 5 (or more) active reminders. " + "Please either remove an old one, or wait for one to pass", @@ -147,7 +147,7 @@ class RemindmeCog(CacheCog): message=message, remind_at=remind_at, active=True, - ).insert() + ).save() embed = build_embed( title="Reminder Set", @@ -210,7 +210,7 @@ class RemindmeCog(CacheCog): hidden=True, ) return - reminders = Reminder.get_active(user=ctx.author.id) + reminders = Reminder.objects(user=ctx.author.id, active=True) if not reminders: await ctx.send("You have no reminders set.", hidden=True) return @@ -225,7 +225,7 @@ class RemindmeCog(CacheCog): description="Delete a reminder", ) async def _delete(self, ctx: SlashContext): - reminders = Reminder.get_active(user=ctx.author.id) + reminders = Reminder.objects(user=ctx.author.id, active=True) if not reminders: await ctx.send("You have no reminders set", hidden=True) return @@ -263,8 +263,8 @@ class RemindmeCog(CacheCog): timeout=60 * 5, ) for to_delete in context.selected_options: - _ = Reminder.get( - user=ctx.author.id, _id=ObjectId(to_delete) + _ = Reminder.objects( + user=ctx.author.id, id=ObjectId(to_delete) ).delete() for row in components: @@ -307,8 +307,8 @@ class RemindmeCog(CacheCog): @loop(seconds=15) async def _remind(self): - reminders = Reminder.get_active( - remind_at={"$lt": datetime.utcnow() + timedelta(seconds=30)} + reminders = Reminder.objects( + remind_at__lte=datetime.utcnow() + timedelta(seconds=30) ) for reminder in reminders: if reminder.remind_at <= datetime.utcnow(): diff --git a/jarvis/cogs/rolegiver.py b/jarvis/cogs/rolegiver.py index 01c9b72..a0e4c59 100644 --- a/jarvis/cogs/rolegiver.py +++ b/jarvis/cogs/rolegiver.py @@ -3,7 +3,7 @@ from discord.ext import commands from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Setting +from jarvis.db.models import Setting from jarvis.utils import build_embed from jarvis.utils.field import Field from jarvis.utils.permissions import admin_or_permissions @@ -28,7 +28,9 @@ class RolegiverCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _rolegiver_add(self, ctx: SlashContext, role: Role): - setting = Setting.get(guild=ctx.guild.id, setting="rolegiver") + setting = Setting.objects( + guild=ctx.guild.id, setting="rolegiver" + ).first() if setting and role.id in setting.value: await ctx.send("Role already in rolegiver", hidden=True) return @@ -51,7 +53,7 @@ class RolegiverCog(commands.Cog): Field(name="Existing Role(s)", value=value), ] setting.value.append(role.id) - setting.update() + setting.save() embed = build_embed( title="Rolegiver Updated", @@ -87,7 +89,9 @@ class RolegiverCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _rolegiver_remove(self, ctx: SlashContext, role: Role): - setting = Setting.get(guild=ctx.guild.id, setting="rolegiver") + setting = Setting.objects( + guild=ctx.guild.id, setting="rolegiver" + ).first() if not setting or (setting and not setting.value): await ctx.send("Rolegiver has no roles", hidden=True) return @@ -96,7 +100,7 @@ class RolegiverCog(commands.Cog): return setting.value.remove(role.id) - setting.update() + setting.save() roles = [] for role_id in setting.value: @@ -137,7 +141,9 @@ class RolegiverCog(commands.Cog): description="List roles rolegiver", ) async def _rolegiver_list(self, ctx: SlashContext): - setting = Setting.get(guild=ctx.guild.id, setting="rolegiver") + setting = Setting.objects( + guild=ctx.guild.id, setting="rolegiver" + ).first() if not setting or (setting and not setting.value): await ctx.send("Rolegiver has no roles", hidden=True) return @@ -186,7 +192,9 @@ class RolegiverCog(commands.Cog): ) @commands.cooldown(1, 10, commands.BucketType.user) async def _role_get(self, ctx: SlashContext, role: Role): - setting = Setting.get(guild=ctx.guild.id, setting="rolegiver") + setting = Setting.objects( + guild=ctx.guild.id, setting="rolegiver" + ).first() if not setting or (setting and not setting.value): await ctx.send("Rolegiver has no roles", hidden=True) return @@ -244,7 +252,9 @@ class RolegiverCog(commands.Cog): ) @commands.cooldown(1, 10, commands.BucketType.user) async def _role_forfeit(self, ctx: SlashContext, role: Role): - setting = Setting.get(guild=ctx.guild.id, setting="rolegiver") + setting = Setting.objects( + guild=ctx.guild.id, setting="rolegiver" + ).first() if not setting or (setting and not setting.value): await ctx.send("Rolegiver has no roles", hidden=True) return diff --git a/jarvis/cogs/settings.py b/jarvis/cogs/settings.py index cce2208..8355721 100644 --- a/jarvis/cogs/settings.py +++ b/jarvis/cogs/settings.py @@ -4,7 +4,7 @@ from discord.utils import find from discord_slash import SlashContext, cog_ext from discord_slash.utils.manage_commands import create_option -from jarvis.db.types import Setting +from jarvis.db.models import Setting from jarvis.utils import build_embed from jarvis.utils.field import Field from jarvis.utils.permissions import admin_or_permissions @@ -15,13 +15,15 @@ class SettingsCog(commands.Cog): self.bot = bot def update_settings(self, setting, value, guild): - setting = Setting(setting=setting, value=value, guild=guild) - updated = setting.update() + setting = Setting.objects(setting=setting, guild=guild).first() + if not setting: + setting = Setting(setting=setting, guild=guild, value=value) + updated = setting.save() return updated is not None def delete_settings(self, setting, guild): - return Setting.delete_many(setting=setting, guild=guild) + return Setting.objects(setting=setting, guild=guild).delete() @cog_ext.cog_subcommand( base="settings", @@ -231,7 +233,7 @@ class SettingsCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _view(self, ctx: SlashContext): - settings = Setting.get_many(guild=ctx.guild.id) + settings = Setting.objects(guild=ctx.guild.id) fields = [] for setting in settings: @@ -261,8 +263,8 @@ class SettingsCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _clear(self, ctx: SlashContext): - deleted = Setting.delete_many(guild=ctx.guild.id) - await ctx.send(f"Guild settings cleared: `{deleted}`") + deleted = Setting.objects(guild=ctx.guild.id).delete() + await ctx.send(f"Guild settings cleared: `{deleted is not None}`") def setup(bot): diff --git a/jarvis/cogs/starboard.py b/jarvis/cogs/starboard.py index 69fe4ae..f855338 100644 --- a/jarvis/cogs/starboard.py +++ b/jarvis/cogs/starboard.py @@ -10,7 +10,7 @@ from discord_slash.utils.manage_components import ( wait_for_component, ) -from jarvis.db.types import Star, Starboard +from jarvis.db.models import Star, Starboard from jarvis.utils import build_embed from jarvis.utils.permissions import admin_or_permissions @@ -34,7 +34,7 @@ class StarboardCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _list(self, ctx): - starboards = Starboard.get_many(guild=ctx.guild.id) + starboards = Starboard.objects(guild=ctx.guild.id) if starboards != []: message = "Available Starboards:\n" for s in starboards: @@ -68,15 +68,17 @@ class StarboardCog(commands.Cog): await ctx.send("Channel must be a TextChannel", hidden=True) return - exists = Starboard.get(channel=channel.id, guild=ctx.guild.id) + exists = Starboard.objects( + channel=channel.id, guild=ctx.guild.id + ).first() if exists: await ctx.send( f"Starboard already exists at {channel.mention}.", hidden=True ) return - count = Starboard.get_many(guild=ctx.guild.id) - if count and len(count) >= 25: + count = Starboard.objects(guild=ctx.guild.id).count() + if count >= 25: await ctx.send("25 starboard limit reached", hidden=True) return @@ -84,7 +86,7 @@ class StarboardCog(commands.Cog): guild=ctx.guild.id, channel=channel.id, admin=ctx.author.id, - ).insert() + ).save() await ctx.send( f"Starboard created. Check it out at {channel.mention}." ) @@ -104,11 +106,11 @@ class StarboardCog(commands.Cog): ) @admin_or_permissions(manage_guild=True) async def _delete(self, ctx, channel: TextChannel): - deleted = Starboard.get( + deleted = Starboard.objects( channel=channel.id, guild=ctx.guild.id ).delete() if deleted: - _ = Star.delete_many(starboard=channel.id) + _ = Star.objects(starboard=channel.id).delete() await ctx.send( f"Starboard deleted from {channel.mention}.", hidden=True ) @@ -146,7 +148,7 @@ class StarboardCog(commands.Cog): ): if not channel: channel = ctx.channel - starboards = Starboard.get_many(guild=ctx.guild.id) + starboards = Starboard.objects(guild=ctx.guild.id) if not starboards: await ctx.send("No starboards exist.", hidden=True) return @@ -189,12 +191,12 @@ class StarboardCog(commands.Cog): message = await channel.fetch_message(int(message)) - exists = Star.get( + exists = Star.objects( message=message.id, channel=message.channel.id, guild=message.guild.id, starboard=starboard.id, - ) + ).first() if exists: await ctx.send( @@ -247,7 +249,7 @@ class StarboardCog(commands.Cog): admin=ctx.author.id, star=star.id, active=True, - ).insert() + ).save() components[0]["components"][0]["disabled"] = True @@ -286,7 +288,9 @@ class StarboardCog(commands.Cog): if not isinstance(starboard, TextChannel): await ctx.send("Channel must be a TextChannel", hidden=True) return - exists = Starboard.get(channel=starboard.id, guild=ctx.guild.id) + exists = Starboard.objects( + channel=starboard.id, guild=ctx.guild.id + ).first() if not exists: await ctx.send( f"Starboard does not exist in {starboard.mention}. " @@ -295,12 +299,12 @@ class StarboardCog(commands.Cog): ) return - star = Star.get( + star = Star.objects( starboard=starboard.id, index=id, guild=ctx.guild.id, active=True, - ) + ).first() if not star: await ctx.send(f"No star exists with id {id}", hidden=True) return @@ -310,7 +314,7 @@ class StarboardCog(commands.Cog): await message.delete() star.active = False - star.update() + star.save() await ctx.send(f"Star {id} deleted") diff --git a/jarvis/cogs/verify.py b/jarvis/cogs/verify.py index 2465b4a..ac62b58 100644 --- a/jarvis/cogs/verify.py +++ b/jarvis/cogs/verify.py @@ -5,7 +5,7 @@ from discord_slash import ComponentContext, SlashContext, cog_ext from discord_slash.model import ButtonStyle from discord_slash.utils import manage_components -from jarvis.db.types import Setting +from jarvis.db.models import Setting def create_layout(): @@ -37,7 +37,7 @@ class VerifyCog(commands.Cog): @commands.cooldown(1, 15, commands.BucketType.user) async def _verify(self, ctx: SlashContext): await ctx.defer() - role = Setting.get(guild=ctx.guild.id, setting="verified") + role = Setting.objects(guild=ctx.guild.id, setting="verified").first() if not role: await ctx.send( "This guild has not enabled verification", delete_after=5 @@ -68,10 +68,14 @@ class VerifyCog(commands.Cog): for c in components: for c2 in c["components"]: c2["disabled"] = True - setting = Setting.get(guild=ctx.guild.id, setting="verified") + setting = Setting.objects( + guild=ctx.guild.id, setting="verified" + ).first() role = ctx.guild.get_role(setting.value) await ctx.author.add_roles(role, reason="Verification passed") - setting = Setting.get(guild=ctx.guild.id, setting="unverified") + setting = Setting.objects( + guild=ctx.guild.id, setting="unverified" + ).first() if setting: role = ctx.guild.get_role(setting.value) await ctx.author.remove_roles( diff --git a/jarvis/config.py b/jarvis/config.py index 5f984d3..dc0dfaa 100644 --- a/jarvis/config.py +++ b/jarvis/config.py @@ -1,6 +1,6 @@ from yaml import load -from jarvis.db import DBManager +from jarvis.db.models import Config as DBConfig try: from yaml import CLoader as Loader @@ -38,10 +38,11 @@ class Config(object): self.events = events self.max_messages = max_messages self.gitlab_token = gitlab_token - db = DBManager(config=mongo["connect"]).mongo[mongo["database"]].config - db_config = db.find() - for item in db_config: - setattr(self, item["key"], item["value"]) + + def get_db_config(self): + config = DBConfig.objects() + for item in config: + setattr(self, item.key, item.value) @classmethod def from_yaml(cls, y): diff --git a/jarvis/db/__init__.py b/jarvis/db/__init__.py deleted file mode 100644 index d6c564c..0000000 --- a/jarvis/db/__init__.py +++ /dev/null @@ -1,14 +0,0 @@ -from pymongo import MongoClient - - -class DBManager(object): - def __new__(cls, *args, **kwargs): - it = cls.__dict__.get("it") - if it is not None: - return it - cls.__it__ = it = object.__new__(cls) - it.init(*args, **kwargs) - return it - - def init(self, config: dict): - self.mongo = MongoClient(**config) diff --git a/jarvis/db/models.py b/jarvis/db/models.py new file mode 100644 index 0000000..db81f3e --- /dev/null +++ b/jarvis/db/models.py @@ -0,0 +1,197 @@ +from datetime import datetime + +from mongoengine import Document +from mongoengine.fields import ( + BooleanField, + DateTimeField, + DictField, + DynamicField, + IntField, + ListField, + LongField, + StringField, +) + + +class SnowflakeField(LongField): + pass + + +class Autopurge(Document): + guild = SnowflakeField(required=True) + channel = SnowflakeField(required=True) + delay = IntField(min_value=1, max_value=300, default=30) + admin = SnowflakeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Autoreact(Document): + guild = SnowflakeField(required=True) + channel = SnowflakeField(required=True) + reactions = ListField(field=StringField()) + admin = SnowflakeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Ban(Document): + active = BooleanField(default=True) + admin = SnowflakeField(required=True) + user = SnowflakeField(required=True) + username = StringField(required=True) + discrim = IntField(min_value=1, max_value=9999, required=True) + duration = IntField(min_value=1, max_value=744, required=False) + guild = SnowflakeField(required=True) + type = StringField(default="perm", max_length=4, required=True) + reason = StringField(max_length=100, required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Config(Document): + key = StringField(required=True) + value = DynamicField(required=True) + + meta = {"db_alias": "main"} + + +class Guess(Document): + correct = BooleanField(default=False) + guess = StringField(max_length=800, required=True) + user = SnowflakeField(required=True) + + meta = {"db_alias": "ctc2"} + + +class Joke(Document): + rid = StringField() + body = StringField() + title = StringField() + created_utc = DateTimeField() + over_18 = BooleanField() + score = IntField() + + meta = {"db_alias": "main"} + + +class Kick(Document): + admin = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + reason = StringField(max_length=100, required=True) + user = SnowflakeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Lock(Document): + active = BooleanField(default=True) + admin = SnowflakeField(required=True) + channel = SnowflakeField(required=True) + duration = IntField(min_value=1, max_value=300, default=10) + guild = SnowflakeField(required=True) + reason = StringField(max_length=100, required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Mute(Document): + active = BooleanField(default=True) + user = SnowflakeField(required=True) + admin = SnowflakeField(required=True) + duration = IntField(min_value=1, max_value=300, default=10) + guild = SnowflakeField(required=True) + reason = StringField(max_length=100, required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Purge(Document): + admin = SnowflakeField(required=True) + channel = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + count = IntField(min_value=1, default=10) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Reminder(Document): + active = BooleanField(default=True) + user = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + message = StringField(max_length=100, required=True) + remind_at = DateTimeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Roleping(Document): + active = BooleanField(default=True) + role = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + admin = SnowflakeField(required=True) + bypass = DictField() + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Setting(Document): + guild = SnowflakeField(required=True) + setting = StringField(required=True) + value = DynamicField() + + meta = {"db_alias": "main"} + + +class Star(Document): + index = IntField(required=True) + message = SnowflakeField(required=True) + starboard = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + admin = SnowflakeField(required=True) + star = SnowflakeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Starboard(Document): + channel = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + admin = SnowflakeField(required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Unban(Document): + user = SnowflakeField(required=True) + username = StringField(required=True) + discrim = IntField(min_value=1, max_value=9999, required=True) + guild = SnowflakeField(required=True) + admin = SnowflakeField(required=True) + reason = StringField(max_length=100, required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} + + +class Warning(Document): + active = BooleanField(default=True) + admin = SnowflakeField(required=True) + user = SnowflakeField(required=True) + guild = SnowflakeField(required=True) + duration = IntField(min_value=1, max_value=120, default=24) + reason = StringField(max_length=100, required=True) + created_at = DateTimeField(default=datetime.utcnow) + + meta = {"db_alias": "main"} diff --git a/jarvis/db/types.py b/jarvis/db/types.py deleted file mode 100644 index e2e9c4d..0000000 --- a/jarvis/db/types.py +++ /dev/null @@ -1,594 +0,0 @@ -import logging -from dataclasses import asdict, dataclass, field -from datetime import datetime -from typing import Any, Optional - -from bson import ObjectId -from pymongo import ASCENDING, DESCENDING - -from jarvis.config import get_config -from jarvis.db import DBManager - -logger = logging.getLogger("mongodb") - -sort_lookup = { - "asc": ASCENDING, - "ascending": ASCENDING, - "desc": DESCENDING, - "descending": DESCENDING, -} - -coll_lookup = { - "Autoreact": "autoreact", - "Autopurge": "autopurges", - "Ban": "bans", - "Config": "config", - "Lock": "locks", - "Kick": "kicks", - "Mute": "mutes", - "Purge": "purges", - "Reminder": "reminders", - "Roleping": "rolepings", - "Setting": "settings", - "Starboard": "starboard", - "Star": "stars", - "Unban": "unbans", - "Warning": "warns", -} - -m_config = get_config().mongo - -db_instance = DBManager(m_config["connect"]).mongo[m_config["database"]] - - -################# -# Core Classes # -################# - - -@dataclass -class MongoSort: - direction: int - key: Any - - def __init__(self, direction, key): - if type(direction) is str: - direction = sort_lookup[direction] - self.direction = direction - self.key = key - - def as_tuple(self): - return (self.key, self.direction) - - -@dataclass -class MongoObject: - """ - A MongoDB object - - :param _id: MongoDB ObjectId, not initialized by default - """ - - _id: ObjectId = field(default=None, init=False) - - def __post_init__(self): - self.coll = db_instance[coll_lookup[type(self).__name__]] - - def to_dict(self): - return asdict(self) - - def insert(self) -> ObjectId: - """ - Inserts the object into the database - - :param collection: Collection to insert object into - :return: Inserted object ObjectId - """ - to_insert = self.to_dict() - if not to_insert["_id"]: - _ = to_insert.pop("_id") - try: - result = self.coll.insert_one(to_insert) - except Exception as e: - logger.error(f"Failed to insert {type(self).__name__}", e) - else: - id = result.inserted_id - self._id = id - return id - - def delete(self) -> bool: - """ - Deletes an object from the database - - :param collection: Collection to delete object from - :return: If delete was successful - """ - to_delete = self.to_dict() - search = {"_id": to_delete["_id"]} - try: - result = self.coll.delete_one(search) - except Exception as e: - logger.error(f"Failed to delete {type(self).__name__}", e) - else: - return result.deleted_count > 0 - - @classmethod - def delete_many(cls, **kwargs) -> bool: - try: - coll = db_instance[coll_lookup[cls.__name__]] - result = coll.delete_many(kwargs) - except Exception as e: - logger.error(f"Failed to delete {cls.__name__}s", e) - else: - return result.deleted_count > 0 - - def update(self) -> bool: - """ - Updates an object in the database - - :param collection: Collection to update - :return: If update was successful - """ - to_update = self.to_dict() - if self._id: - search = {"_id": to_update["_id"]} - try: - result = self.coll.update_one( - search, {"$set": to_update}, upsert=True - ) - except Exception as e: - logger.error(f"Failed to update {type(self).__name__}", e) - else: - return result.modified_count > 0 - else: - self.insert() - - @classmethod - def get(cls, **kwargs) -> Optional[object]: - """ - Get an object from MongoDB - - :param collection: Collection to query - :return: Optional object - """ - try: - sort = None - if "sort" in kwargs: - sort = kwargs.pop("sort") - if type(sort) is list: - sort = [x.as_tuple() for x in sort] - else: - sort = [sort.as_tuple()] - coll = db_instance[coll_lookup[cls.__name__]] - result = coll.find_one(kwargs, sort=sort) - except Exception as e: - logger.error(f"Failed to get {cls.__name__}", e) - else: - if result: - _id = result.pop("_id") - r = cls(**result) - r._id = _id - return r - return None - - @classmethod - def get_many(cls, **kwargs) -> Optional[list]: - """ - Gets objects from MongoDB - - :param collection: Collection to query - :return: Optional object - """ - global db_instance - try: - sort = None - if "sort" in kwargs: - sort = kwargs.pop("sort") - if type(sort) is list: - sort = [x.as_tuple() for x in sort] - else: - sort = [sort.as_tuple()] - coll = db_instance[coll_lookup[cls.__name__]] - result = coll.find(kwargs, sort=sort) - except Exception as e: - logger.error(f"Failed to get {cls.__name__}", e) - else: - if result: - r_list = [] - for r_item in result: - _id = r_item.pop("_id") - r = cls(**r_item) - r._id = _id - r_list.append(r) - return r_list - return [] - - -@dataclass -class ActiveObject: - """ - A type of Mongo object that can be active - - :param active: If object is active - """ - - active: bool - - @classmethod - def get_active(cls, **kwargs) -> list: - """ - Gets a list of active objects - - :param collection: Collection to query - :return: List of active objects - """ - kwargs.update({"active": True}) - try: - sort = None - if "sort" in kwargs: - sort = kwargs.pop("sort") - if type(sort) is list: - sort = [x.as_tuple() for x in sort] - else: - sort = [sort.as_tuple()] - coll = db_instance[coll_lookup[cls.__name__]] - result = coll.find(kwargs, sort=sort) - except Exception as e: - logger.error(f"Failed to get {cls.__name__}", e) - else: - if result: - r_list = [] - for r_item in result: - _id = r_item.pop("_id") - r = cls(**r_item) - r._id = _id - r_list.append(r) - return r_list - return [] - - -#################### -# Database Objects # -#################### - - -@dataclass -class Autopurge(MongoObject): - """ - Channel Autoreact object - - :param _id: MongoDB ID - :param guild: ID of guild - :param channel: ID of channel - :param delay: Purge delay - :param admin: ID of admin who added autoreact - :param created_at: Time the autoreact was created - """ - - guild: int - channel: int - delay: int - admin: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Autoreact(MongoObject): - """ - Channel Autoreact object - - :param _id: MongoDB ID - :param guild: ID of guild - :param channel: ID of channel - :param reactions: List of reactions - :param admin: ID of admin who added autoreact - :param created_at: Time the autoreact was created - """ - - guild: int - channel: int - reactions: list - admin: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Ban(MongoObject, ActiveObject): - """ - User Ban object - - :param _id: MongoDB ID - :param active: If the ban is active - :param admin: ID of admin who banned the user - :param user: ID of banned user - :param username: Username of banned user - :param discrim: Discriminator of banned user - :param duration: Duration of ban - :param guild: ID of guild that user belonged to - :param type: Type of ban - :param reason: Reason for the ban - :param created_at: Time the ban happened - """ - - admin: int - user: int - username: str - discrim: int - duration: int - guild: int - type: str - reason: str - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Config(MongoObject): - """ - J.A.R.V.I.S. Config object - - :param _id: MongoDB ID - :param key: Config key - :param value: Config value - """ - - key: str - value: any - - -@dataclass -class Joke(MongoObject): - """ - Joke object - - :param _id: MongoDB ID - :param id: Reddit ID - :param body: Joke Body - :param title: Joke Title - :param created_utc: Created at - :param over_18: Is NSFW - :param score: Reddit Score - """ - - id: str - body: str - title: str - created_utc: datetime - over_18: bool - score: int - - -@dataclass -class Kick(MongoObject): - """ - User Kick object - - :param _id: MongoDB ID - :param user: Kicked User ID - :param reason: Kick reason - :param admin: ID of admin who kicked user - :param created_at: Time user was kicked - :param guild: Guild ID - """ - - admin: int - guild: int - reason: str - user: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Lock(MongoObject, ActiveObject): - """ - Channel Lock object - - :param _id: MongoDB ID - :param active: If the lock is active - :param admin: ID of admin who locked channel - :param channel: ID of locked channel - :param duration: Duration of lock - :param guild: ID of guild that channel belongs to - :param reason: Reason for the lock - :param created_at: Time the lock happened - """ - - admin: int - channel: int - duration: int - guild: int - reason: str - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Mute(MongoObject, ActiveObject): - """ - User Mute object - - :param _id: MongoDB ID - :param active: If the mute is active - :param admin: ID of admin who muted the user - :param user: ID of muted user - :param duration: Duration of mute - :param guild: ID of guild that user belongs to - :param reason: Reason for the mute - :param created_at: Time the mute happened - """ - - admin: int - user: int - duration: int - guild: int - reason: str - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Purge(MongoObject): - """ - Channel Purge object - - :param _id: MongoDB ID - :param admin: ID of admin who purged messages - :param channel: ID of purged channel - :param guild: ID of guild that channel belongs to - :param count: Number of purged messages - :param created_at: Time the purge happened - """ - - admin: int - channel: int - guild: int - count: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Reminder(MongoObject, ActiveObject): - """ - User Reminder object - - :param _id: MongoDB ID - :param user: User who wanted reminded - :param channel: ID of origin channel - :param guild: ID of origin guild - :param remind_at: Datetime delay to send message - :param active: If reminder is active - :param message: Reminder message - :param created_at: Time the reminder was created - """ - - user: int - channel: int - guild: int - remind_at: datetime - message: str - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Roleping(MongoObject, ActiveObject): - """ - Guild Roleping object - - :param _id: MongoDB ID - :param role: Blocked role - :param guild: ID of origin guild - :param admin: Admin who added roleping - :param bypass: Roles/users who may bypass this roleping - :param active: If the roleping is disabled - :param created_at: Time the roleping was created - """ - - role: int - guild: int - admin: int - bypass: dict - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Setting(MongoObject): - """ - Guild Setting object - - :param _id: MongoDB ID - :param guild: ID of guild - :param setting: Setting key - :param value: Setting value - """ - - guild: int - setting: str - value: any - - -@dataclass -class Star(MongoObject, ActiveObject): - """ - Starboard Star object - - :param _id: MongoDB ID - :param index: Starboard star index - :param message: Star Message ID - :param channel: Star Channel ID - :param starboard: Starboard Channel ID - :param guild: Starboard Guild ID - :param admin: ID of admin who created star - :param star: ID of the star message - :param created_at: Time created - """ - - index: int - message: int - channel: int - starboard: int - guild: int - admin: int - star: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Starboard(MongoObject): - """ - Channel Starboard object - - :param _id: MongoDB ID - :param channel: Starboard Channel ID - :param guild: Starboard Guild ID - :param admin: ID of admin who created starboard - :param created_at: Time created - """ - - channel: int - guild: int - admin: int - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Unban(MongoObject): - """ - Guild Unban object - - :param _id: MongoDB ID - :param user: User ID - :param username: User Username - :param discrim: User Discriminator - :param guild: Guild ID - :param admin: Admin who unbanned user - :param reason: Reason for unban - :param created_at: Time created - """ - - user: int - username: str - discrim: int - guild: int - admin: int - reason: str - created_at: datetime = field(default_factory=datetime.utcnow) - - -@dataclass -class Warning(MongoObject, ActiveObject): - """ - User Warning object - - :param _id: MongoDB ID - :param active: If the warning is active - :param admin: ID of admin who warned the user - :param created_at: Time the warning happened - :param duration: Duration of warning - :param guild: ID of guild that user belongs to - :param reason: Reason for the warning - :param user: ID of warned user - """ - - admin: int - user: int - duration: int - guild: int - reason: str - created_at: datetime = field(default_factory=datetime.utcnow) diff --git a/jarvis/events/guild.py b/jarvis/events/guild.py index eb5ccb9..3b4b79c 100644 --- a/jarvis/events/guild.py +++ b/jarvis/events/guild.py @@ -2,7 +2,7 @@ import asyncio from discord.utils import find -from jarvis.db.types import Setting +from jarvis.db.models import Setting class GuildEventHandler(object): @@ -26,7 +26,6 @@ class GuildEventHandler(object): ) # Set some default settings - setting = Setting(guild=guild.id, setting="massmention", value=5) - setting.insert() + _ = Setting(guild=guild.id, setting="massmention", value=5).save() await general.send("Systems are now fully operational") diff --git a/jarvis/events/member.py b/jarvis/events/member.py index d26170e..43ea280 100644 --- a/jarvis/events/member.py +++ b/jarvis/events/member.py @@ -1,6 +1,6 @@ from discord import Member -from jarvis.db.types import Mute, Setting +from jarvis.db.models import Mute, Setting class MemberEventHandler(object): @@ -10,14 +10,16 @@ class MemberEventHandler(object): async def on_member_join(self, user: Member): guild = user.guild - mutes = Mute.get_active(guild=guild.id) - if mutes and len(mutes) >= 1: - mute_role = Setting.get(guild=guild.id, setting="mute") + mute = Mute.objects(guild=guild.id, user=user.id, active=True).first() + if mute: + mute_role = Setting.objects(guild=guild.id, setting="mute").first() role = guild.get_role(mute_role.value) await user.add_roles( role, reason="User is muted still muted from prior mute" ) - unverified = Setting.get(guild=guild.id, setting="unverified") + unverified = Setting.objects( + guild=guild.id, setting="unverified" + ).first() if unverified: role = guild.get_role(unverified.value) if role not in user.roles: diff --git a/jarvis/events/message.py b/jarvis/events/message.py index a47af51..3b73959 100644 --- a/jarvis/events/message.py +++ b/jarvis/events/message.py @@ -4,7 +4,7 @@ from discord import DMChannel, Message from discord.utils import find from jarvis.config import get_config -from jarvis.db.types import Autopurge, Autoreact, Roleping, Setting, Warning +from jarvis.db.models import Autopurge, Autoreact, Roleping, Setting, Warning from jarvis.utils import build_embed from jarvis.utils.field import Field @@ -21,17 +21,17 @@ class MessageEventHandler(object): self.bot.add_listener(self.on_message_edit) async def autopurge(self, message: Message): - autopurge = Autopurge.get( + autopurge = Autopurge.objects( guild=message.guild.id, channel=message.channel.id - ) + ).first() if autopurge: await message.delete(delay=autopurge.delay) async def autoreact(self, message: Message): - autoreact = Autoreact.get( + autoreact = Autoreact.objects( guild=message.guild.id, channel=message.channel.id, - ) + ).first() if autoreact: for reaction in autoreact.reactions: await message.add_reaction(reaction) @@ -57,15 +57,14 @@ class MessageEventHandler(object): ] if match.group(1) not in allowed: await message.delete() - warning = Warning( + _ = Warning( active=True, admin=get_config().client_id, duration=24, guild=message.guild.id, reason="Sent an invite link", user=message.author.id, - ) - warning.insert() + ).save() fields = [ Field( "Reason", @@ -92,10 +91,10 @@ class MessageEventHandler(object): await message.channel.send(embed=embed) async def massmention(self, message: Message): - massmention = Setting.get( + massmention = Setting.objects( guild=message.guild.id, setting="massmention", - ) + ).first() if ( massmention and massmention.value > 0 @@ -103,15 +102,14 @@ class MessageEventHandler(object): - (1 if message.author in message.mentions else 0) > massmention.value ): - warning = Warning( + _ = Warning( active=True, admin=get_config().client_id, duration=24, guild=message.guild.id, reason="Mass Mention", user=message.author.id, - ) - warning.insert() + ).save() fields = [Field("Reason", "Mass Mention", False)] embed = build_embed( title="Warning", @@ -131,7 +129,7 @@ class MessageEventHandler(object): await message.channel.send(embed=embed) async def roleping(self, message: Message): - rolepings = Roleping.get_active(guild=message.guild.id) + rolepings = Roleping.objects(guild=message.guild.id, active=True) if not rolepings: return @@ -180,15 +178,14 @@ class MessageEventHandler(object): and not user_is_admin and not user_has_bypass ): - warning = Warning( + _ = Warning( active=True, admin=get_config().client_id, duration=24, guild=message.guild.id, reason="Pinged a blocked role/user with a blocked role", user=message.author.id, - ) - warning.insert() + ).save() fields = [ Field( "Reason", diff --git a/jarvis/tasks/unban.py b/jarvis/tasks/unban.py index 597edca..49008bd 100644 --- a/jarvis/tasks/unban.py +++ b/jarvis/tasks/unban.py @@ -1,16 +1,18 @@ from datetime import datetime, timedelta -import pymongo from discord.ext.tasks import loop import jarvis -from jarvis.db.types import Ban +from jarvis.config import get_config +from jarvis.db.models import Ban, Unban + +jarvis_id = get_config().client_id @loop(minutes=10) async def unban(): - bans = Ban.get_active(type="temp") - updates = [] + bans = Ban.objects(type="temp", active=True) + unbans = [] for ban in bans: if ban.created_at + timedelta( hours=ban.duration @@ -19,16 +21,17 @@ async def unban(): user = await jarvis.jarvis.fetch_user(ban.user) if user: guild.unban(user) - updates.append( - pymongo.UpdateOne( - { - "user": user.id, - "guild": guild.id, - "created_at": ban.created_at, - "type": "temp", - }, - {"$set": {"active": False}}, + ban.active = False + ban.save() + unbans.append( + Unban( + user=user.id, + guild=guild.id, + username=user.name, + discrim=user.discriminator, + admin=jarvis_id, + reason="Ban expired", ) ) - if updates: - jarvis.jarvis_db.bans.bulk_write(updates) + if unbans: + Ban.objects().insert(unbans) diff --git a/jarvis/tasks/unlock.py b/jarvis/tasks/unlock.py index 26078a9..037ae90 100644 --- a/jarvis/tasks/unlock.py +++ b/jarvis/tasks/unlock.py @@ -1,16 +1,14 @@ from datetime import datetime, timedelta -import pymongo from discord.ext.tasks import loop import jarvis -from jarvis.db.types import Lock +from jarvis.db.models import Lock @loop(minutes=1) async def unlock(): - locks = Lock.get_active() - updates = [] + locks = Lock.objects(active=True) for lock in locks: if ( lock.created_at + timedelta(minutes=lock.duration) @@ -26,15 +24,5 @@ async def unlock(): await channel.set_permissions( role, overwrite=overrides, reason="Lock expired" ) - updates.append( - pymongo.UpdateOne( - { - "channel": channel.id, - "guild": guild.id, - "created_at": lock.created_at, - }, - {"$set": {"active": False}}, - ) - ) - if updates: - jarvis.jarvis_db.locks.bulk_write(updates) + lock.active = False + lock.save() diff --git a/jarvis/tasks/unmute.py b/jarvis/tasks/unmute.py index 064112d..3a92c29 100644 --- a/jarvis/tasks/unmute.py +++ b/jarvis/tasks/unmute.py @@ -1,17 +1,15 @@ from datetime import datetime, timedelta -import pymongo from discord.ext.tasks import loop import jarvis -from jarvis.db.types import Mute, Setting +from jarvis.db.models import Mute, Setting @loop(minutes=1) async def unmute(): - mutes = Mute.get_active(duration={"$gt": 0}) - mute_roles = Setting.get_many(setting="mute") - updates = [] + mutes = Mute.objects(duration__gt=0, active=True) + mute_roles = Setting.objects(setting="mute") for mute in mutes: if ( mute.created_at + timedelta(minutes=mute.duration) @@ -28,15 +26,5 @@ async def unmute(): await user.remove_roles(role, reason="Mute expired") # Objects can't handle bulk_write, so handle it via raw methods - updates.append( - pymongo.UpdateOne( - { - "user": user.id, - "guild": guild.id, - "created_at": mute.created_at, - }, - {"$set": {"active": False}}, - ) - ) - if updates: - jarvis.jarvis_db.mutes.bulk_write(updates) + mute.active = False + mute.save diff --git a/jarvis/tasks/unwarn.py b/jarvis/tasks/unwarn.py index 423e43f..f4455c6 100644 --- a/jarvis/tasks/unwarn.py +++ b/jarvis/tasks/unwarn.py @@ -1,25 +1,18 @@ from datetime import datetime, timedelta -import pymongo from discord.ext.tasks import loop import jarvis -from jarvis.db.types import Warning +from jarvis.db.models import Warning @loop(hours=1) async def unwarn(): - warns = Warning.get_active() - updates = [] + warns = Warning.objects(active=True) for warn in warns: if ( warn.created_at + timedelta(hours=warn.duration) < datetime.utcnow() ): - updates.append( - pymongo.UpdateOne( - {"_id": warn._id}, {"$set": {"active": False}} - ) - ) - if updates: - jarvis.jarvis_db.warns.bulk_write(updates) + warn.active = False + warn.save() diff --git a/requirements.txt b/requirements.txt index 1916f9a..f44f4d0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -9,3 +9,4 @@ ButtonPaginator>=0.0.3 Pillow>=8.2.0, <9 python-gitlab>=2.9.0, <3 ulid-py>=1.1.0, <2 +mongoengine>=0.23, <1