Merge pull request #18 from TogetherCrew/feat/add-shared-codes

Feat: add shared codes

docker-compose.test.yml

@@ -19,6 +19,9 @@ services:
       - POSTGRES_PORT=5432
       - SENTRY_DSN=sample_dsn
       - SENTRY_ENV=local
+      - REDIS_HOST=redis
+      - REDIS_PORT=6379
+      - REDIS_PASSWORD=
       - CHUNK_SIZE=512
       - EMBEDDING_DIM=1024
       - QDRANT_HOST=qdrant
@@ -33,6 +36,8 @@ services:
         condition: service_healthy
       qdrant-healthcheck:
         condition: service_healthy
+      redis:
+        condition: service_healthy
   postgres:
     image: "ankane/pgvector"
     environment:
@@ -69,6 +74,17 @@ services:
       interval: 10s
       timeout: 2s
       retries: 5
+  redis:
+    image: redis:latest
+    ports:
+      - "6379:6379"
+    networks:
+      - python_service_network
+    healthcheck:
+      test: ["CMD", "redis-cli", "ping"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
 
 
 networks:

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name="tc-hivemind-backend",
-    version="1.3.0",
+    version="1.4.0",
     author="Mohammad Amin Dadgar, TogetherCrew",
     maintainer="Mohammad Amin Dadgar",
     maintainer_email="[email protected]",

tc_hivemind_backend/db/modules_base.py

@@ -0,0 +1,133 @@
+from bson import ObjectId
+
+from .mongo import MongoSingleton
+
+
+class ModulesBase:
+    def __init__(self) -> None:
+        pass
+
+    def query(self, platform: str, **kwargs) -> list[dict]:
+        """
+        query the modules database for to get platforms' metadata
+
+        Parameters
+        -----------
+        platform : str
+            the platform to choose
+            it can be `github`, `discourse`, `discord` or etc
+        **kwargs : dict
+            projection : dict[str, int]
+                feature projection on query
+
+        Returns
+        ---------
+        modules_docs : list[dict]
+            all the module documents that have the `platform` within them
+        """
+        client = MongoSingleton.get_instance().get_client()
+        projection = kwargs.get("projection", {})
+
+        cursor = client["Core"]["modules"].find(
+            {
+                "options.platforms.name": platform,
+                "name": "hivemind",
+            },
+            projection,
+        )
+        modules_docs = list(cursor)
+        return modules_docs
+
+    def get_platform_community_ids(self, platform_name: str) -> list[str]:
+        """
+        get all community ids that a platform has
+
+        Parameters
+        ------------
+        platform_name : str
+            the platform having community id and available for hivemind module
+
+        Returns
+        --------
+        community_ids : list[str]
+            id of communities that has discord platform and hivemind module enabled
+
+        """
+        modules = self.query(platform=platform_name, projection={"community": 1})
+        community_ids = list(map(lambda x: str(x["community"]), modules))
+
+        return community_ids
+
+    def get_token(self, platform_id: ObjectId, token_type: str) -> str:
+        """
+        get a specific type of token for a platform
+        This method is called when we needed a token for modules to extract its data
+
+        Parameters
+        ------------
+        platform_id : ObjectId
+            the platform id that we want their token
+        token_type : str
+            the type of token. i.e. `google_refresh`
+
+        Returns
+        --------
+        token : str
+            the token that was required for module's ETL process
+        """
+        client = MongoSingleton.get_instance().get_client()
+
+        user_id = self.get_platform_metadata(platform_id, "userId")
+        user_id = ObjectId(user_id)
+        token_doc = client["Core"]["tokens"].find_one(
+            {
+                "user": user_id,
+                "type": token_type,
+            },
+            {
+                "token": 1,
+            },
+            sort=[("createdAt", -1)],
+        )
+        if token_doc is None:
+            raise ValueError(
+                f"No Token for the given user {user_id} "
+                "in tokens collection of the Core database!"
+            )
+        token = token_doc["token"]
+        return token
+
+    def get_platform_metadata(
+        self, platform_id: ObjectId, metadata_name: str
+    ) -> str | dict | list:
+        """
+        get the userid that belongs to a platform
+
+        Parameters
+        -----------
+        platform_id : bson.ObjectId
+            the platform id we need their owner user id
+        metadata_name : str
+            a specific field of metadata that we want
+
+        Returns
+        ---------
+        metadata_value : Any
+            the values that the metadata belongs to
+        """
+        client = MongoSingleton.get_instance().get_client()
+
+        platform = client["Core"]["platforms"].find_one(
+            {
+                "_id": platform_id,
+                "disconnectedAt": None,
+            },
+            {
+                f"metadata.{metadata_name}": 1,
+            },
+        )
+        if platform is None:
+            raise ValueError(f"No platform available given platform id: {platform_id}")
+
+        metadata_field = platform["metadata"][metadata_name]
+        return metadata_field

tc_hivemind_backend/db/mongo.py

@@ -1,12 +1,13 @@
 import logging
+from typing import Optional
 
 from pymongo import MongoClient
 
 from .credentials import Credentials
 
 
 class MongoSingleton:
-    __instance: "MongoSingleton" | None = None
+    __instance: Optional["MongoSingleton"] = None
 
     def __init__(self):
         if MongoSingleton.__instance is not None:

tc_hivemind_backend/tests/integration/test_modules_base_query_token.py

@@ -0,0 +1,121 @@
+from datetime import datetime, timedelta
+from unittest import TestCase
+
+from bson import ObjectId
+from tc_hivemind_backend.db.modules_base import ModulesBase
+from tc_hivemind_backend.db.mongo import MongoSingleton
+
+
+class TestModulesBaseQueryToken(TestCase):
+    def setUp(self) -> None:
+        self.client = MongoSingleton.get_instance().get_client()
+        self.client["Core"].drop_collection("tokens")
+        self.client["Core"].drop_collection("platforms")
+
+    def test_one_token(self):
+        sample_user = ObjectId("5d7baf326c8a2e2400000000")
+        community_id = ObjectId("6579c364f1120850414e0dc5")
+        sample_token_type = "type1"
+        sample_token_value = "tokenid12345"
+        platform_id = ObjectId("6579c364f1120850414e0dc6")
+
+        self.client["Core"]["platforms"].insert_one(
+            {
+                "_id": platform_id,
+                "name": "platform_name",
+                "metadata": {
+                    "id": "113445975232201081511",
+                    "userId": str(sample_user),
+                },
+                "community": community_id,
+                "disconnectedAt": None,
+                "connectedAt": datetime.now(),
+                "createdAt": datetime.now(),
+                "updatedAt": datetime.now(),
+            }
+        )
+
+        sample_token_doc = {
+            "token": sample_token_value,
+            "user": sample_user,
+            "type": sample_token_type,
+            "expires": datetime.now() + timedelta(days=1),
+            "blacklisted": False,
+            "createdAt": datetime.now() - timedelta(days=1),
+            "updatedAt": datetime.now() - timedelta(days=1),
+        }
+        self.client["Core"]["tokens"].insert_one(sample_token_doc)
+        token = ModulesBase().get_token(
+            platform_id=platform_id, token_type=sample_token_type
+        )
+
+        self.assertEqual(token, sample_token_value)
+
+    def test_empty_tokens_collection(self):
+        platform_id = ObjectId("6579c364f1120850414e0dc6")
+        sample_token_type = "type1"
+        with self.assertRaises(ValueError):
+            _ = ModulesBase().get_token(
+                platform_id=platform_id, token_type=sample_token_type
+            )
+
+    def test_no_platform(self):
+        sample_user = ObjectId("5d7baf326c8a2e2400000000")
+        platform_id = ObjectId("6579c364f1120850414e0dc6")
+        sample_token_type = "type1"
+        sample_token_value = "tokenid12345"
+
+        sample_token_doc = {
+            "token": sample_token_value,
+            "user": sample_user,
+            "type": sample_token_type,
+            "expires": datetime.now() + timedelta(days=1),
+            "blacklisted": False,
+            "createdAt": datetime.now() - timedelta(days=1),
+            "updatedAt": datetime.now() - timedelta(days=1),
+        }
+        self.client["Core"]["tokens"].insert_one(sample_token_doc)
+        with self.assertRaises(ValueError):
+            _ = ModulesBase().get_token(
+                platform_id=platform_id, token_type=sample_token_type
+            )
+
+    def test_no_token(self):
+        sample_user = ObjectId("5d7baf326c8a2e2400000000")
+        sample_user_with_no_token = ObjectId("5d7baf326c8a2e2400000001")
+
+        platform_id = ObjectId("6579c364f1120850414e0dc6")
+        sample_token_type = "type1"
+        sample_token_value = "tokenid12345"
+        community_id = ObjectId("6579c364f1120850414e0dc5")
+
+        self.client["Core"]["platforms"].insert_one(
+            {
+                "_id": platform_id,
+                "name": "platform_name",
+                "metadata": {
+                    "id": "113445975232201081511",
+                    "userId": str(sample_user_with_no_token),
+                },
+                "community": community_id,
+                "disconnectedAt": None,
+                "connectedAt": datetime.now(),
+                "createdAt": datetime.now(),
+                "updatedAt": datetime.now(),
+            }
+        )
+
+        sample_token_doc = {
+            "token": sample_token_value,
+            "user": sample_user,
+            "type": sample_token_type,
+            "expires": datetime.now() + timedelta(days=1),
+            "blacklisted": False,
+            "createdAt": datetime.now() - timedelta(days=1),
+            "updatedAt": datetime.now() - timedelta(days=1),
+        }
+        self.client["Core"]["tokens"].insert_one(sample_token_doc)
+        with self.assertRaises(ValueError):
+            _ = ModulesBase().get_token(
+                platform_id=platform_id, token_type=sample_token_type
+            )

tc_hivemind_backend/tests/integration/test_mongo_singleton.py

@@ -0,0 +1,26 @@
+import unittest
+
+from pymongo.mongo_client import MongoClient
+from tc_hivemind_backend.db.mongo import MongoSingleton
+
+
+class TestMongoSingletonIntegration(unittest.TestCase):
+    """Separate test class for integration tests"""
+
+    def setUp(self):
+        """Reset the singleton instance before each test"""
+        MongoSingleton._MongoSingleton__instance = None
+
+    def test_real_connection(self):
+        """Integration test for real MongoDB connection"""
+        instance = MongoSingleton.get_instance()
+        client = instance.get_client()
+        self.assertIsInstance(client, MongoClient)
+
+        # Test actual connection
+        try:
+            info = client.server_info()
+            self.assertIsInstance(info, dict)
+            self.assertIn("version", info)
+        except Exception as e:
+            self.fail(f"Failed to connect to MongoDB: {e}")