Added total message count to metricity data response.

2 files changed, 24 insertions, 1 deletions

diff --git a/postgres/init.sql b/postgres/init.sql
index fd29ddbc..45ad440c 100644
--- a/postgres/init.sql
+++ b/postgres/init.sql
@@ -32,3 +32,19 @@ INSERT INTO users VALUES (
     '{}',
     NULL
 );
+
+CREATE TABLE messages (
+    id varchar(255),
+    author_id varchar(255) references users(id),
+    primary key(id)
+);
+
+INSERT INTO messages VALUES(
+    0,
+    0
+);
+
+INSERT INTO messages VALUES(
+    1,
+    0
+);
diff --git a/pydis_site/apps/api/viewsets/bot/user.py b/pydis_site/apps/api/viewsets/bot/user.py
index 059bc0f0..b3d880cc 100644
--- a/pydis_site/apps/api/viewsets/bot/user.py
+++ b/pydis_site/apps/api/viewsets/bot/user.py
@@ -76,7 +76,8 @@ class UserViewSet(BulkCreateModelMixin, ModelViewSet):
     ...    "is_guild": True,
     ...    "is_verified": False,
     ...    "public_flags": {},
-    ...    "verified_at": null
+    ...    "verified_at": None,
+    ...    "total_messages": 2
     ...}
 
     #### Status codes
@@ -157,9 +158,15 @@ class UserViewSet(BulkCreateModelMixin, ModelViewSet):
         column_keys = ["id", "name", "avatar_hash", "joined_at", "created_at", "is_staff",
                        "opt_out", "bot", "is_guild", "is_verified", "public_flags", "verified_at"]
         with connections['metricity'].cursor() as cursor:
+            # Get user data
             query = f"SELECT {','.join(column_keys)} FROM users WHERE id = '%s'"
             cursor.execute(query, [user.id])
             values = cursor.fetchone()
             data = dict(zip(column_keys, values))
             data["public_flags"] = json.loads(data["public_flags"])
+
+            # Get message count
+            cursor.execute("SELECT COUNT(*) FROM messages WHERE author_id = '%s'", [user.id])
+            data["total_messages"], = cursor.fetchone()
+
             return Response(data, status=status.HTTP_200_OK)