rename

DB_utls.py → DB_utils.py

@@ -1,4 +1,4 @@
-#from lambda_mongo_utils import send_data_to_mongodb
+# from lambda_mongo_utils import send_data_to_mongodb
 from pymongo import MongoClient
 from datetime import datetime
 from prefect import task
@@ -15,59 +15,54 @@ connection_string = blinded_connection_string.replace("<db_password>", MONGODB_P
 @task
 def generate_empty_well():
     dbclient = MongoClient(connection_string)
-    db = dbclient["LCM-OT-2-SLD"]  
-    collection = db["wells"] 
-    rows = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H']
+    db = dbclient["LCM-OT-2-SLD"]
+    collection = db["wells"]
+    rows = ["A", "B", "C", "D", "E", "F", "G", "H"]
     columns = [str(i) for i in range(1, 13)]
     for row in rows:
         for col in columns:
             well = f"{row}{col}"
-            metadata = {
-                "well": well,
-                "status": "empty",
-                "project": "OT2"
-            }
+            metadata = {"well": well, "status": "empty", "project": "OT2"}
 
-            #send_data_to_mongodb(collection="wells", data=metadata)
+            # send_data_to_mongodb(collection="wells", data=metadata)
             query = {"well": well}
-            update_data = {"$set": metadata}  
+            update_data = {"$set": metadata}
             result = collection.update_one(query, update_data, upsert=True)
-            
+
     # close connection
     dbclient.close()
 
-@task 
+
+@task
 def update_used_wells(used_wells):
     dbclient = MongoClient(connection_string)
-    db = dbclient["LCM-OT-2-SLD"]  
-    collection = db["wells"] 
+    db = dbclient["LCM-OT-2-SLD"]
+    collection = db["wells"]
 
     for well in used_wells:
-        metadata = {
-            "well": well,
-            "status": "used",
-            "project": "OT2"
-        }
-        #send_data_to_mongodb(collection="wells", data=metadata)
+        metadata = {"well": well, "status": "used", "project": "OT2"}
+        # send_data_to_mongodb(collection="wells", data=metadata)
         query = {"well": well}
-        update_data = {"$set": metadata}  
+        update_data = {"$set": metadata}
         result = collection.update_one(query, update_data, upsert=True)
-        
+
     # close connection
     dbclient.close()
 
+
 @task
 def find_unused_wells():
     dbclient = MongoClient(connection_string)
-    db = dbclient["LCM-OT-2-SLD"]  
-    collection = db["wells"] 
+    db = dbclient["LCM-OT-2-SLD"]
+    collection = db["wells"]
     query = {"status": "empty"}
-    response = list(collection.find(query))  
+    response = list(collection.find(query))
     df = pd.DataFrame(response)
 
     # Extract the "well" column as a list
     if "well" not in df.columns:
         raise ValueError("No available wells.")
+
     # Sort obtained list
     def well_sort_key(well):
         row = well[0]  # A, B, C, ...
@@ -75,43 +70,46 @@ def find_unused_wells():
         return (row, col)
 
     empty_wells = sorted(df["well"].tolist(), key=well_sort_key)
-    #print(empty_wells)
-    
+    # print(empty_wells)
+
     # close connection
     dbclient.close()
 
     # Check if there are any empty wells
     if len(empty_wells) == 0:
         raise ValueError("No empty wells found")
-    #print(empty_wells)
+    # print(empty_wells)
     return empty_wells
-    
-@task 
+
+
+@task
 def save_result(result_data):
     dbclient = MongoClient(connection_string)
-    db = dbclient["LCM-OT-2-SLD"]  
-    collection = db["MSE403_result"] 
-    #collection = db["test_result"] 
+    db = dbclient["LCM-OT-2-SLD"]
+    collection = db["MSE403_result"]
+    # collection = db["test_result"]
     result_data["timestamp"] = datetime.utcnow()  # UTC time
     insert_result = collection.insert_one(result_data)
     inserted_id = insert_result.inserted_id
     # close connection
     dbclient.close()
-    return inserted_id    
+    return inserted_id
+
 
 def get_student_quota(student_id):
     with MongoClient(connection_string) as client:
-        db = client["LCM-OT-2-SLD"]  
-        collection = db["student"] 
+        db = client["LCM-OT-2-SLD"]
+        collection = db["student"]
         student = collection.find_one({"student_id": student_id})
         if student is None:
             raise ValueError(f"Student ID '{student_id}' not found in the database.")
-        return student.get("quota", 0) 
-  
+        return student.get("quota", 0)
+
+
 def decrement_student_quota(student_id):
     dbclient = MongoClient(connection_string)
-    db = dbclient["LCM-OT-2-SLD"]  
-    collection = db["student"] 
+    db = dbclient["LCM-OT-2-SLD"]
+    collection = db["student"]
 
     student = collection.find_one({"student_id": student_id})
     if not student:
@@ -119,17 +117,16 @@ def decrement_student_quota(student_id):
     if student.get("quota", 0) <= 0:
         return f"Student ID {student_id} has no remaining quota."
 
-    
     result = collection.update_one(
-            {"student_id": student_id, "quota": {"$gt": 0}},  
-            {"$inc": {"quota": -1}} 
-        )
+        {"student_id": student_id, "quota": {"$gt": 0}}, {"$inc": {"quota": -1}}
+    )
 
     if result.modified_count > 0:
         return f"Student ID {student_id}'s quota update successfully."
     else:
         return f"Quota update failed for Student ID {student_id}."
-    
+
+
 def add_student_quota(student_id, quota):
     """
     Adds a new student with a given quota.
@@ -140,15 +137,15 @@ def add_student_quota(student_id, quota):
     db = dbclient["LCM-OT-2-SLD"]
     collection = db["student"]
     student_data = {"student_id": student_id, "quota": quota}
-    collection.update_one({"student_id": student_id}, {"$set": student_data}, upsert=True)
+    collection.update_one(
+        {"student_id": student_id}, {"$set": student_data}, upsert=True
+    )
     dbclient.close()
-    
+
 
 if __name__ == "__main__":
     generate_empty_well()
-    #find_unused_wells()
-    #test_id = "test"
-    #quota = 999
-    #add_student_quota(test_id, quota)
-       
-    
+    # find_unused_wells()
+    # test_id = "test"
+    # quota = 999
+    # add_student_quota(test_id, quota)