Broke out metricity connection into an

abstraction and added metricity endpoint unit
tests.

1 files changed, 42 insertions, 0 deletions

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]