Upload DB_utls.py

DB_utls.py

@@ -0,0 +1,154 @@
+#from lambda_mongo_utils import send_data_to_mongodb
+from pymongo import MongoClient
+from datetime import datetime
+from prefect import task
+import pandas as pd
+import os
+
+MONGODB_PASSWORD = os.getenv("MONGODB_PASSWORD")
+
+blinded_connection_string = os.getenv("blinded_connection_string")
+
+connection_string = blinded_connection_string.replace("<db_password>", MONGODB_PASSWORD)
+
+
+@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']
+    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"
+            }
+
+            #send_data_to_mongodb(collection="wells", data=metadata)
+            query = {"well": well}
+            update_data = {"$set": metadata}  
+            result = collection.update_one(query, update_data, upsert=True)
+            
+    # close connection
+    dbclient.close()
+
+@task 
+def update_used_wells(used_wells):
+    dbclient = MongoClient(connection_string)
+    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)
+        query = {"well": well}
+        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"] 
+    query = {"status": "empty"}
+    response = list(collection.find(query))  
+    df = pd.DataFrame(response)
+
+    # Extract the "well" column as a list
+    if "well" not in df.columns:
+        raise ValueError("The returned data does not contain the 'well' field.")
+    # Sort obtained list
+    def well_sort_key(well):
+        row = well[0]  # A, B, C, ...
+        col = int(well[1:])  # 1, 2, ..., 12
+        return (row, col)
+
+    empty_wells = sorted(df["well"].tolist(), key=well_sort_key)
+    #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)
+    return empty_wells
+    
+@task 
+def save_result(result_data):
+    dbclient = MongoClient(connection_string)
+    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    
+
+def get_student_quota(student_id):
+    with MongoClient(connection_string) as client:
+        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) 
+  
+def decrement_student_quota(student_id):
+    dbclient = MongoClient(connection_string)
+    db = dbclient["LCM-OT-2-SLD"]  
+    collection = db["student"] 
+
+    student = collection.find_one({"student_id": student_id})
+    if not student:
+        return f"Student ID {student_id} not found."
+    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}} 
+        )
+
+    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.
+    :param student_id: The ID of the student.
+    :param quota: The initial quota for the student.
+    """
+    dbclient = MongoClient(connection_string)
+    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)
+    dbclient.close()
+    
+
+if __name__ == "__main__":
+    generate_empty_well()
+    #find_unused_wells()
+    #test_id = "test"
+    #quota = 999
+    #add_student_quota(test_id, quota)
+       
+