Add Mongodb indexes

db.py

@@ -0,0 +1,19 @@
+import os
+
+from dotenv import load_dotenv
+from pymongo import MongoClient
+
+load_dotenv()
+
+
+class MongoDBConnection:
+    def __init__(self, connection_string=os.environ.get("MONGODB_URI")):
+        self.client = MongoClient(connection_string)
+        self.db = self.client[os.environ.get("MONGODB_DB_NAME")]
+
+    @property
+    def database(self):
+        return self.db
+
+    def close_connection(self):
+        self.client.close()

update_cars.py

@@ -2,6 +2,7 @@
 from typing import List
 
 import updater
+from db import MongoDBConnection
 
 
 async def main():
@@ -12,6 +13,15 @@ async def main():
     )
     key_fields: List[str] = ["month"]
 
+    db = MongoDBConnection().database
+    collection = db[collection_name]
+    collection.create_index({"month": 1, "make": 1});
+    collection.create_index({"make": 1});
+    collection.create_index({"fuel_type": 1});
+    collection.create_index({"make": 1, "fuel_type": 1});
+    collection.create_index({"number": 1});
+    db.client.close()
+
     response = await updater.main(
         collection_name=collection_name,
         zip_url=zip_url,

update_coe.py

@@ -2,6 +2,7 @@
 from typing import List
 
 import updater
+from db import MongoDBConnection
 
 
 async def main():
@@ -12,6 +13,15 @@ async def main():
     )
     key_fields: List[str] = ["month", "bidding_no"]
 
+    db = MongoDBConnection().database
+    collection = db[collection_name]
+    collection.create_index({"month": 1, "vehicle_class": 1});
+    collection.create_index({"vehicle_class": 1});
+    collection.create_index({"month": 1, "bidding_no": 1});
+    collection.create_index({"premium": 1});
+    collection.create_index({"bids_success": 1, "bids_received": 1});
+    db.client.close()
+
     response = await updater.main(
         collection_name=collection_name,
         zip_url=zip_url,

updater.py

@@ -6,8 +6,8 @@
 from typing import List, Dict, Any
 
 from dotenv import load_dotenv
-from pymongo import MongoClient
 
+from db import MongoDBConnection
 from download_file import download_file
 from utils.create_unique_key import create_unique_key
 from utils.extract_zip_file import extract_zip_file
@@ -18,8 +18,9 @@
 async def updater(
         collection_name: str, zip_file_name: str, zip_url: str, key_fields: List[str]
 ) -> str:
-    client = MongoClient(os.environ.get("MONGODB_URI"))
-    db = client[os.environ.get("MONGODB_DB_NAME")]
+    # client = MongoClient(os.environ.get("MONGODB_URI"))
+    # db = client[os.environ.get("MONGODB_DB_NAME")]
+    db = MongoDBConnection().database
     collection = db[collection_name]
 
     try:
@@ -53,6 +54,7 @@ async def updater(
                 start = time.time()
                 result = collection.insert_many(new_data_to_insert)
                 end = time.time()
+                db.client.close()
                 message = f"{len(result.inserted_ids)} document(s) inserted in {round((end - start) * 1000)}ms"
             else:
                 message = "No new data to insert. The provided data matches the existing records."