add db async methods

token-scanner-api/src/server/api.py

@@ -18,10 +18,10 @@ db_name = env_vars['MONGODB_DB_NAME']
 
 @app.on_event("startup")
 def startup_db_client():
-    app.mongodb_client = MongoClient(url)
-    app.database = app.mongodb_client[db_name]
+    app.client = MongoClient(url)
+    app.database = app.client[db_name]
     os_utils.log_info("Connected to the MongoDB database!")
 
 @app.on_event("shutdown")
 def shutdown_db_client():
-    app.mongodb_client.close()
+    app.client.close()

token-scanner-api/src/server/database.py

@@ -3,58 +3,77 @@
 # This class implements the database connection.
 
 
-from pymongo import MongoClient
+import pymongo 
 
 
+env_vars = os_utils.load_config()
+url = env_vars['MONGODB_URL']
+db_name = env_vars['MONGODB_DB_NAME']
+collection = env_vars['MONGODB_COLLECTION_NAME']
 
 
-client = MongoClient(MONGO_DETAILS)
+############################################
+#     Private methods: database            #
+############################################
 
-database = client.balances
+def _get_db_collection():
+    """Connect to the database."""
 
-collection = database.get_collection("balances_fuckkk")
+    client = pymongo.MongoClient(url)
+    database = client.balances
+    return database[collection]
 
 
-print(database.list_collection_names())
-
-
-def wallet_helper(item) -> dict:
+def _wallet_helper(item) -> dict:
     return {
         "wallet": item["wallet"],
     }
 
 
-from bson.objectid import ObjectId
-
-async def retrieve_students():
-    bals = collection.find()
-
-    res = []
-    counter = 0
-    for i in bals:
-        res.append(wallet_helper(i))
-        if counter > 2:
-            break
-        counter += 1
-
-    return res
-
-
-def balancer_helper(item) -> dict:
+def _balancer_helper(item) -> dict:
     return {
         "wallet": item["wallet"],
         "balance": item["balance"],
     }
 
 
-# Retrieve a student with a matching ID
+############################################
+#     Public methods: database             #
+############################################
+
+async def retrieve_top_balances(top_number: int) -> list:
+    """Retrieve top balances from the database."""
+
+    collection = _get_db_collection()
+    top_balances = collection.find()
+
+    result = []
+    counter = 0
+    
+    for balance in top_balances:
+        result.append(_wallet_helper(balance))
+        if counter > top_number:
+            break
+        counter += 1
+
+    return result
+
+
 async def retrieve_balance(wallet: str) -> dict:
+    """Retrieve balance from the database."""
+
+    collection = _get_db_collection()
     balance = collection.find_one({"wallet": wallet})
+
     if balance:
-        return balancer_helper(balance)
+        return _balancer_helper(balance)
+    else:
+        return {}
 
-async def retrieve_top_balances():
+
+async def retrieve_holder_weekly_change(address: str) -> int:
+    """Retrieve weekly change of a given address."""
+
+    collection = _get_db_collection()
+    # todo
     pass
-
-async def retrieve_holder_weekly_change(address: str):
-    pass

token-scanner-api/src/server/routes.py

@@ -32,10 +32,12 @@ async def get_token_balance(address: str) -> dict:
 
 
 @router.get("/top")
-async def get_top_holders() -> dict:
+async def get_top_holders(top_number=None) -> dict:
     """Get top holders of a given token."""
 
-    futures = [retrieve_top_balances()]
+    top_number = top_number or 10
+
+    futures = [retrieve_top_balances(top_number)]
     result = await asyncio.gather(*futures)
     if result:
         return {"top_holders": result}

token-scanner-api/src/utils/db_processing.py

@@ -26,7 +26,7 @@ def populate_db(filepath):
     env_vars = os_utils.load_config()
     url = env_vars['MONGODB_URL']
     db_name = env_vars['MONGODB_DB_NAME']
-    collection_name = env_vars['MONGODB_COLLECTION_NAME']
+    collection = env_vars['MONGODB_COLLECTION_NAME']
     
     client = pymongo.MongoClient(url)
 
@@ -40,7 +40,7 @@ def populate_db(filepath):
         # Create collection for balances
         # (make sure it's empty first)
         ################################
-        wallet_collection = database[collection_name]
+        wallet_collection = database[collection]
         wallet_collection.drop()
 
         ##############################