tickets_plus.database.layer

  1"""A layer for the database session.
  2
  3This module contains the `OnlineConfig` class.
  4It is used to make the database session easier to use.
  5It is also an async context manager.
  6Generally, you should use this class instead of the session directly.
  7Additionally, of note is the fact that one-on-one relationships are loaded
  8by default. But any relationship that is a one-to-many or many-to-many
  9relationship is not loaded by default.
 10This is due to efficiency concerns.
 11If you need to load a relationship, you can use the options argument, which
 12is a list of `sqlalchemy.sql.base.ExecutableOption`s.
 13For more information on `ExecutableOption`s, see the SQLAlchemy documentation.
 14
 15Typical usage example:
 16    ```py
 17    async with OnlineConfig(bot, session) as db:
 18        guild = await db.get_guild(guild_id)
 19        # Do something with the guild
 20        # and if you changed something
 21        await db.commit()
 22    ```
 23"""
 24# License: EPL-2.0
 25# SPDX-License-Identifier: EPL-2.0
 26# Copyright (c) 2021-present The Tickets+ Contributors
 27# This Source Code may also be made available under the following
 28# Secondary Licenses when the conditions for such availability set forth
 29# in the Eclipse Public License, v. 2.0 are satisfied: GPL-3.0-only OR
 30# If later approved by the Initial Contributor, GPL-3.0-or-later.
 31
 32import datetime
 33import types
 34from typing import Any, Sequence, Tuple, Type
 35
 36import discord
 37from discord import utils
 38from discord.ext import commands
 39from sqlalchemy import sql
 40from sqlalchemy.ext import asyncio as sa_asyncio
 41from sqlalchemy.sql import base
 42
 43from tickets_plus.database import models
 44
 45
 46class OnlineConfig:
 47    """A convenience layer for the database session.
 48
 49    This class is used to make the database session easier to use.
 50    It also handles the 'async with' statement.
 51    Any and all commits have to be done manually.
 52    """
 53
 54    def __init__(self, bot_instance: commands.AutoShardedBot, session: sa_asyncio.AsyncSession) -> None:
 55        """Initialises the database session layer.
 56
 57        Wraps the provided session in an async context manager.
 58        And set the bot instance as a private attribute.
 59
 60        Args:
 61            bot_instance: The bot instance.
 62            session: The database session.
 63        """
 64        self._session = session
 65        self._bot = bot_instance
 66
 67    async def __aenter__(self) -> "OnlineConfig":
 68        """Enter the 'async with' statement.
 69
 70        This method is called when entering the 'async with' statement.
 71
 72        Returns:
 73            OnlineConfig: The database session layer.
 74        """
 75        return self
 76
 77    async def __aexit__(self, exc_type: Type[BaseException] | None, exc_value: BaseException | None,
 78                        traceback: types.TracebackType | None) -> None:
 79        """Exit the 'async with' statement.
 80
 81        This method is called when exiting the 'async with' statement.
 82        Or when an exception is raised.
 83        We roll back the session if an exception is raised.
 84        We then close the session regardless.
 85
 86        Args:
 87            exc_type: The exception type.
 88            exc_value: The exception value.
 89            traceback: The traceback.
 90        """
 91        if exc_type:
 92            await self.rollback()
 93        await self.close()
 94
 95    async def close(self) -> None:
 96        """Close the database session.
 97
 98        This method closes the underlying SQLAlchemy session.
 99        """
100        await self._session.close()
101
102    async def flush(self) -> None:
103        """Flush the database session.
104
105        Moves all the changes to the database transaction.
106        """
107        await self._session.flush()
108
109    async def commit(self) -> None:
110        """Commit the database session.
111
112        Commits the underlying SQLAlchemy session.
113        """
114        await self._session.commit()
115
116    async def rollback(self) -> None:
117        """Rollback the database session.
118
119        Cleans all the non-committed changes.
120        This includes the flushed changes.
121        """
122        await self._session.rollback()
123
124    async def delete(self, obj) -> None:
125        """Delete a row from the database.
126
127        Marks the object for deletion.
128        The deletion will be done on the next flush.
129        """
130        await self._session.delete(obj)
131
132    async def get_guild(self, guild_id: int, options: Sequence[base.ExecutableOption] | None = None) -> models.Guild:
133        """Get or create a guild from the database.
134
135        Fetches a guild from the database.
136        Due to the convenience of the database layer,
137        we guarantee a guild will always be returned.
138        It will be created if it does not exist.
139        However, we do not commit the changes.
140
141        Args:
142            guild_id: The guild ID.
143            options: The options to use when querying the database.
144                Those options are mostly used for relationship loading.
145                As in async, we can't use lazy loading.
146
147        Returns:
148            models.Guild: The guild.
149                With the relationships loaded if options are provided.
150                Otherwise, attempting to access the relationships will
151                result in an error.
152        """
153        if options:
154            guild_conf = await self._session.scalar(
155                sql.select(models.Guild).where(models.Guild.guild_id == guild_id).options(*options))
156        else:
157            guild_conf = await self._session.scalar(sql.select(models.Guild).where(models.Guild.guild_id == guild_id))
158        if guild_conf is None:
159            guild_conf = models.Guild(guild_id=guild_id)
160            self._session.add(guild_conf)
161        return guild_conf
162
163    async def get_user(self, user_id: int, options: Sequence[base.ExecutableOption] | None = None) -> models.User:
164        """Get or create a user from the database.
165
166        Fetches a user from the database.
167        Due to the convenience of the database layer,
168        we guarantee a user will always be returned.
169        It will be created if it does not exist.
170        However, we do not commit the changes.
171
172        Args:
173            user_id: The user ID.
174            options: The options to use when querying the database.
175                Those options are mostly used for relationship loading.
176                As in async, we can't use lazy loading.
177
178        Returns:
179            models.User: The user.
180                With the relationships loaded if options are provided.
181                Otherwise, attempting to access the relationships will
182                result in an error.
183        """
184        if options:
185            user = await self._session.scalar(
186                sql.select(models.User).where(models.User.user_id == user_id).options(*options))
187        else:
188            user = await self._session.scalar(sql.select(models.User).where(models.User.user_id == user_id))
189        if user is None:
190            user = models.User(user_id=user_id)
191            self._session.add(user)
192        return user
193
194    async def get_member(self, user_id: int, guild_id: int) -> models.Member:
195        """Get or create a member from the database.
196
197        Fetches a member from the database.
198        If the member does not exist, it will be created.
199        We also check if the guild exists and create it if it does not.
200        We also check if the user exists and create it if it does not.
201
202        Args:
203            user_id: The user ID.
204            guild_id: The guild ID.
205
206        Returns:
207            models.Member: The member.
208                As members have a one-to-one relationship with guilds,
209                and a one-to-one relationship with users, we automatically
210                load the guild and user relationships.
211        """
212        guild = await self.get_guild(guild_id)
213        user = await self.get_user(user_id)
214        member_conf = await self._session.scalar(
215            sql.select(models.Member).where(models.Member.user == user, models.Member.guild == guild))
216        if member_conf is None:
217            member_conf = models.Member(user=user, guild=guild)
218            self._session.add(member_conf)
219        return member_conf
220
221    async def get_expired_members(self) -> Sequence[models.Member]:
222        """Get expired members from the database.
223
224        Fetches all members with expired status from the database.
225        Used to clean up the roles.
226
227        Returns:
228            Sequence[models.Member]: The members with expired status.
229        """
230        time = datetime.datetime.utcnow()
231        expr_members = await self._session.scalars(sql.select(models.Member).where(models.Member.status_till <= time))
232        return expr_members.all()
233
234    async def get_ticket_bot(self, user_id: int, guild_id: int) -> Tuple[bool, models.TicketBot]:
235        """Get or create a ticket bot from the database.
236
237        Fetches a ticket bot from the database.
238        If the ticket bot does not exist, it will be created.
239        We also check if the guild exists and create it if it does not.
240
241        Args:
242            user_id: The user ID.
243            guild_id: The guild ID.
244
245        Returns:
246            Tuple[bool, models.TicketBot]: A tuple containing a boolean
247                indicating if the ticket bot was created, and the ticket bot.
248                Relationships are loaded.
249        """
250        guild = await self.get_guild(guild_id)
251        ticket_user = await self._session.scalar(
252            sql.select(models.TicketBot).where(models.TicketBot.user_id == user_id, models.TicketBot.guild == guild))
253        new = False
254        if ticket_user is None:
255            new = True
256            ticket_user = models.TicketBot(user_id=user_id, guild=guild)
257            self._session.add(ticket_user)
258        return new, ticket_user
259
260    async def check_ticket_bot(self, user_id: int, guild_id: int) -> bool:
261        """Check if the ticket user exists.
262
263        A more efficient way to check if a ticket user exists.
264        Instead of attempting to create the ticket user,
265        we just check if it exists.
266        This is the only check that uses a guild ID.
267        This is because the user ID is shared across all guilds.
268
269        Args:
270            user_id: The user ID.
271            guild_id: The guild ID.
272
273        Returns:
274            bool: A boolean indicating if the ticket user exists.
275        """
276        ticket_user = await self._session.scalar(
277            sql.select(models.TicketBot).where(models.TicketBot.user_id == user_id,
278                                               models.TicketBot.guild_id == guild_id))
279        return ticket_user is not None
280
281    async def get_ticket_type(self,
282                              guild_id: int,
283                              name: str,
284                              comping: bool = False,
285                              comaccs: bool = False,
286                              strpbuttns: bool = False,
287                              ignore: bool = False) -> Tuple[bool, models.TicketType]:
288        """Get or create a ticket type from the database.
289
290        Fetches a ticket type from the database.
291        If the ticket type does not exist, it will be created.
292        We also check if the guild exists and create it if it does not.
293
294        Args:
295            guild_id: The guild ID.
296            name: The ticket type name.
297            comping: The comping flag.
298            comaccs: The comaccs flag.
299            strpbuttns: The strpbuttns flag.
300            ignore: The ignore flag.
301
302        Returns:
303            Tuple[bool, models.TicketType]: A tuple containing a boolean
304                indicating if the ticket type was created, and the ticket type.
305                Relationships are loaded.
306        """
307        guild = await self.get_guild(guild_id)
308        ticket_type = await self._session.scalar(
309            sql.select(models.TicketType).where(models.TicketType.guild == guild, models.TicketType.prefix == name))
310        new = False
311        if ticket_type is None:
312            new = True
313            ticket_type = models.TicketType(guild=guild,
314                                            prefix=name,
315                                            comping=comping,
316                                            comaccs=comaccs,
317                                            strpbuttns=strpbuttns,
318                                            ignore=ignore)
319            self._session.add(ticket_type)
320        return new, ticket_type
321
322    async def get_ticket_types(self, guild_id: int) -> Sequence[models.TicketType]:
323        """Get ticket types from the database.
324
325        Fetches all ticket types from the database.
326        We also check if the guild exists and create it if it does not.
327
328        Args:
329            guild_id: The guild ID.
330
331        Returns:
332            Sequence[models.TicketType]: The ticket types.
333        """
334        guild = await self.get_guild(guild_id)
335        ticket_types = await self._session.scalars(
336            sql.select(models.TicketType).where(models.TicketType.guild == guild))
337        return ticket_types.all()
338
339    async def fetch_ticket(self, channel_id: int) -> models.Ticket | None:
340        """Fetch a ticket from the database.
341
342        Attempts to fetch a ticket from the database.
343        If the ticket does not exist, None is returned.
344        If you want to create a ticket if it does not exist,
345        use get_ticket instead.
346
347        Args:
348            channel_id: The channel ID.
349
350        Returns:
351            models.Ticket | None: The ticket.
352        """
353        ticket = await self._session.get(models.Ticket, channel_id)
354        return ticket
355
356    async def get_ticket(self,
357                         channel_id: int,
358                         guild_id: int,
359                         user_id: int | None = None,
360                         staff_note: int | None = None) -> Tuple[bool, models.Ticket]:
361        """Get or create a ticket from the database.
362
363        Fetches a ticket from the database.
364        If the ticket does not exist, it will be created.
365        We also check if the guild exists and create it if it does not.
366        If you want to check if a ticket exists, use fetch_ticket instead.
367
368        Args:
369            channel_id: The channel ID.
370            guild_id: The guild ID. Used to create the guild if missing.
371            user_id: The user ID. Only filed if integration is enabled.
372            staff_note: The staff note thread ID. Used to annotate the thread,
373                if the ticket is created.
374
375        Returns:
376            Tuple[bool, models.Ticket]: A tuple containing a boolean
377                indicating if the ticket was created, and the ticket.
378                Relationships are loaded.
379        """
380        guild = await self.get_guild(guild_id)
381        ticket = await self._session.get(models.Ticket, channel_id)
382        new = False
383        if ticket is None:
384            new = True
385            ticket = models.Ticket(channel_id=channel_id, guild=guild, user_id=user_id, staff_note_thread=staff_note)
386            self._session.add(ticket)
387        return new, ticket
388
389    async def get_pending_tickets(self) -> Sequence[models.Ticket]:
390        """Get pending tickets from the database.
391
392        Fetches all pending tickets from the database.
393
394        Returns:
395            Sequence[models.Ticket]: The pending tickets.
396        """
397        tickets = await self._session.scalars(
398            sql.select(models.Ticket).join(models.Guild).filter(
399                models.Guild.warn_autoclose.isnot(None), models.Ticket.notified.isnot(True), models.Ticket.last_response
400                <= models.UTCnow() - models.Guild.warn_autoclose))
401        return tickets.all()
402
403    async def fetch_tag(self, guild_id: int, tag: str) -> discord.Embed | str | None:
404        """Fetch a tag from the database.
405
406        Attempts to fetch a tag from the database.
407        If the tag does not exist, None is returned.
408        If you want to create a tag if it does not exist,
409        use get_tag instead.
410
411        Args:
412            guild_id: The guild ID.
413            tag: The tag.
414
415        Returns:
416            discord.Embed | str | None: The tag.
417        """
418        guild = await self.get_guild(guild_id)
419        embed = await self._session.scalar(
420            sql.select(models.Tag).where(models.Tag.guild == guild, models.Tag.tag_name == tag))
421        if embed is None:
422            return None
423        if embed.title:
424            emdd = vars(embed)
425            emd2 = {}
426            for key, data in emdd.items():
427                if data is not None:
428                    emd2[key] = data
429            if embed.author:
430                emd2["author"] = {"name": embed.author}
431            if embed.footer:
432                emd2["footer"] = {"text": embed.footer}
433            result = discord.Embed.from_dict(emd2)
434            result.timestamp = utils.utcnow()
435            return result
436        return embed.description
437
438    async def get_tag(
439        self,
440        guild_id: int,
441        tag_name: str,
442        description: str,
443        embed_args: dict[str, Any] | None = None,
444    ) -> Tuple[bool, models.Tag]:
445        """Get or create a tag from the database.
446
447        Fetches a tag from the database.
448        If the tag does not exist, it will be created.
449        We also check if the guild exists and create it if it does not.
450        If you want to check if a tag exists, use fetch_tag instead.
451
452        Args:
453            guild_id: The guild ID.
454            tag_name: The tag.
455            description: The description.
456            embed_args: The embed arguments.
457                Basically, the arguments to pass to discord.Embed,
458                when using discord.Embed.from_dict.
459        """
460        guild = await self.get_guild(guild_id)
461        tag = await self._session.scalar(
462            sql.select(models.Tag).where(models.Tag.guild == guild, models.Tag.tag_name == tag_name))
463        new = False
464        if tag is None:
465            new = True
466            if embed_args is None:
467                embed_args = {}
468            tag = models.Tag(guild=guild, tag_name=tag_name, description=description, **embed_args)
469            self._session.add(tag)
470        return new, tag
471
472    async def get_tags(self, guild_id: int) -> Sequence[models.Tag]:
473        """Get tags from the database.
474
475        Fetches all tags from the database.
476        We also check if the guild exists and create it if it does not.
477
478        Args:
479            guild_id: The guild ID.
480
481        Returns:
482            Sequence[models.Tag]: The tags.
483        """
484        guild = await self.get_guild(guild_id)
485        tags = await self._session.scalars(sql.select(models.Tag).where(models.Tag.guild == guild))
486        return tags.all()
487
488    async def get_staff_role(self, role_id: int, guild_id: int) -> Tuple[bool, models.StaffRole]:
489        """Get or create the staff role from the database.
490
491        Fetches a staff role from the database.
492        If the staff role does not exist, it will be created.
493        We also check if the guild exists and create it if it does not.
494
495        Args:
496            role_id: The role ID.
497            guild_id: The guild ID.
498
499        Returns:
500            Tuple[bool, models.StaffRole]: A tuple containing a boolean
501                indicating if the staff role was created, and the staff role.
502                Relationships are loaded.
503        """
504        guild = await self.get_guild(guild_id)
505        staff_role = await self._session.get(models.StaffRole, role_id)
506        new = False
507        if staff_role is None:
508            new = True
509            staff_role = models.StaffRole(role_id=role_id, guild=guild)
510            self._session.add(staff_role)
511        return new, staff_role
512
513    async def get_all_staff_roles(self, guild_id: int) -> Sequence[models.StaffRole]:
514        """Get all staff roles from the database.
515
516        Fetches all staff roles from the database.
517        If the staff roles do not exist, an empty list is returned.
518
519        Args:
520            guild_id: The guild ID.
521
522        Returns:
523            Sequence[models.StaffRole]: A list of staff roles.
524                Relationships are loaded.
525        """
526        guild = await self.get_guild(guild_id)
527        staff_roles = await self._session.scalars(sql.select(models.StaffRole).where(models.StaffRole.guild == guild))
528        return staff_roles.all()
529
530    async def check_staff_role(self, role_id: int) -> bool:
531        """Check if the staff role exists.
532
533        A more efficient way to check if a staff role exists.
534        Instead of attempting to create the staff role,
535        we just check if it exists.
536
537        Args:
538            role_id: The role ID.
539
540        Returns:
541            bool: A boolean indicating if the staff role exists.
542        """
543        staff_role = await self._session.get(models.StaffRole, role_id)
544        return staff_role is not None
545
546    async def get_observers_role(self, role_id: int, guild_id: int) -> Tuple[bool, models.ObserversRole]:
547        """Get or create the observer role from the database.
548
549        Fetches an observer role from the database.
550        If the observer role does not exist, it will be created.
551        We also check if the guild exists and create it if it does not.
552
553        Args:
554            role_id: The role ID.
555            guild_id: The guild ID.
556
557        Returns:
558            Tuple[bool, models.ObserversRole]: A tuple containing a boolean
559                indicating if the observer role was created, and the observers'
560                role. Relationships are loaded.
561        """
562        guild = await self.get_guild(guild_id)
563        observers_role = await self._session.get(models.ObserversRole, role_id)
564        new = False
565        if observers_role is None:
566            new = True
567            observers_role = models.ObserversRole(role_id=role_id, guild=guild)
568            self._session.add(observers_role)
569        return new, observers_role
570
571    async def get_all_observers_roles(self, guild_id: int) -> Sequence[models.ObserversRole]:
572        """Get all observer roles from the database.
573
574        Fetches all observer roles from the database.
575        If the observer roles do not exist, an empty list is returned.
576
577        Args:
578            guild_id: The guild ID.
579
580        Returns:
581            Sequence[models.ObserversRole]: A list of observer roles.
582                Relationships are loaded.
583        """
584        guild = await self.get_guild(guild_id)
585        observers_roles = await self._session.scalars(
586            sql.select(models.ObserversRole).where(models.ObserversRole.guild == guild))
587        return observers_roles.all()
588
589    async def check_observers_role(self, role_id: int) -> bool:
590        """Check if the observer role exists.
591
592        A more efficient way to check if an observer role exists.
593        Instead of attempting to create the observer role,
594        we just check if it exists.
595
596        Args:
597            role_id: The role ID.
598
599        Returns:
600            bool: A boolean indicating if the observer role exists.
601        """
602        observers_role = await self._session.get(models.ObserversRole, role_id)
603        return observers_role is not None
604
605    async def get_community_role(self, role_id: int, guild_id: int) -> Tuple[bool, models.CommunityRole]:
606        """Get or create the community role from the database.
607
608        Fetches a community role from the database.
609        If the community role does not exist, it will be created.
610        We also check if the guild exists and create it if it does not.
611
612        Args:
613            role_id: The role ID.
614            guild_id: The guild ID.
615
616        Returns:
617            Tuple[bool, models.CommunityRole]: A tuple containing a boolean
618                indicating if the community role was created, and the community
619                role. Relationships are loaded.
620        """
621        guild = await self.get_guild(guild_id)
622        community_role = await self._session.get(models.CommunityRole, role_id)
623        new = False
624        if community_role is None:
625            new = True
626            community_role = models.CommunityRole(role_id=role_id, guild=guild)
627            self._session.add(community_role)
628        return new, community_role
629
630    async def get_all_community_roles(self, guild_id: int) -> Sequence[models.CommunityRole]:
631        """Get all community roles from the database.
632
633        Fetches all community roles from the database.
634        If the community roles do not exist, an empty list is returned.
635
636        Args:
637            guild_id: The guild ID.
638
639        Returns:
640            Sequence[models.CommunityRole]: A list of community roles.
641                Relationships are loaded.
642        """
643        guild = await self.get_guild(guild_id)
644        community_roles = await self._session.scalars(
645            sql.select(models.CommunityRole).where(models.CommunityRole.guild == guild))
646        return community_roles.all()
647
648    async def check_community_role(self, role_id: int) -> bool:
649        """Check if the community role exists.
650
651        A more efficient way to check if a community role exists.
652        Instead of attempting to create the community role,
653        we just check if it exists.
654
655        Args:
656            role_id: The role ID.
657
658        Returns:
659            bool: A boolean indicating if the community role exists.
660        """
661        community_role = await self._session.get(models.CommunityRole, role_id)
662        return community_role is not None
663
664    async def get_community_ping(self, role_id: int, guild_id: int) -> Tuple[bool, models.CommunityPing]:
665        """Get or create the community ping from the database.
666
667        Fetches a community ping from the database.
668        If the community ping does not exist, it will be created.
669        We also check if the guild exists and create it if it does not.
670
671        Args:
672            role_id: The role ID.
673            guild_id: The guild ID.
674
675        Returns:
676            Tuple[bool, models.CommunityPing]: A tuple containing a boolean
677                indicating if the community ping was created, and the community
678                pings. Relationships are loaded.
679        """
680        guild = await self.get_guild(guild_id)
681        community_ping = await self._session.get(models.CommunityPing, role_id)
682        new = False
683        if community_ping is None:
684            new = True
685            community_ping = models.CommunityPing(role_id=role_id, guild=guild)
686            self._session.add(community_ping)
687        return new, community_ping
688
689    async def get_all_community_pings(self, guild_id: int) -> Sequence[models.CommunityPing]:
690        """Get all community pings from the database.
691
692        Fetches all community pings from the database.
693        If the community pings do not exist, an empty list is returned.
694
695        Args:
696            guild_id: The guild ID.
697
698        Returns:
699            Sequence[models.CommunityPing]: A list of community pings.
700                Relationships are loaded.
701        """
702        guild = await self.get_guild(guild_id)
703        community_pings = await self._session.scalars(
704            sql.select(models.CommunityPing).where(models.CommunityPing.guild == guild))
705        return community_pings.all()
706
707    async def check_community_ping(self, role_id: int) -> bool:
708        """Check if the community ping exists.
709
710        A more efficient way to check if a community ping exists.
711        Instead of attempting to create the community ping,
712        we just check if it exists.
713
714        Args:
715            role_id: The role ID.
716
717        Returns:
718            bool: A boolean indicating if the community ping exists.
719        """
720        community_ping = await self._session.get(models.CommunityPing, role_id)
721        return community_ping is not None