diff options
Diffstat (limited to 'pydis_site/apps/api/models')
-rw-r--r-- | pydis_site/apps/api/models/bot/infraction.py | 3 | ||||
-rw-r--r-- | pydis_site/apps/api/models/bot/metricity.py | 42 |
2 files changed, 44 insertions, 1 deletions
diff --git a/pydis_site/apps/api/models/bot/infraction.py b/pydis_site/apps/api/models/bot/infraction.py index 7660cbba..60c1e8dd 100644 --- a/pydis_site/apps/api/models/bot/infraction.py +++ b/pydis_site/apps/api/models/bot/infraction.py @@ -15,7 +15,8 @@ class Infraction(ModelReprMixin, models.Model): ("mute", "Mute"), ("kick", "Kick"), ("ban", "Ban"), - ("superstar", "Superstar") + ("superstar", "Superstar"), + ("voice_ban", "Voice Ban"), ) inserted_at = models.DateTimeField( default=timezone.now, diff --git a/pydis_site/apps/api/models/bot/metricity.py b/pydis_site/apps/api/models/bot/metricity.py new file mode 100644 index 00000000..25b42fa2 --- /dev/null +++ b/pydis_site/apps/api/models/bot/metricity.py @@ -0,0 +1,42 @@ +from django.db import connections + + +class NotFound(Exception): + """Raised when an entity cannot be found.""" + + pass + + +class Metricity: + """Abstraction for a connection to the metricity database.""" + + def __init__(self): + self.cursor = connections['metricity'].cursor() + + def __enter__(self): + return self + + def __exit__(self, *_): + self.cursor.close() + + def user(self, user_id: str) -> dict: + """Query a user's data.""" + columns = ["verified_at"] + query = f"SELECT {','.join(columns)} FROM users WHERE id = '%s'" + self.cursor.execute(query, [user_id]) + values = self.cursor.fetchone() + + if not values: + raise NotFound() + + return dict(zip(columns, values)) + + def total_messages(self, user_id: str) -> int: + """Query total number of messages for a user.""" + self.cursor.execute("SELECT COUNT(*) FROM messages WHERE author_id = '%s'", [user_id]) + values = self.cursor.fetchone() + + if not values: + raise NotFound() + + return values[0] |