Update main.py

Author: GorkiPower

backend/main.py

@@ -1,5 +1,5 @@
-from datetime import datetime
 
+from datetime import datetime
 from bson import ObjectId
 from database import schemas_collection
 from datamodel import SchemaDefinition, UpdateSchema
@@ -18,7 +18,6 @@
     allow_headers=["*"],
 )
 
-
 @app.get("/schemas")
 async def get_all_schemas():
     """
@@ -29,37 +28,37 @@ async def get_all_schemas():
     """
     # Fetch all schema documents from MongoDB
     schema_documents = list(schemas_collection.find())
-
-    # Verbose loop: clarify what we're iterating over and why
+    # Convert ObjectId to string for each document
     for schema_document in schema_documents:
-        # Each schema_document is a MongoDB record representing a schema definition
-        # Convert its ObjectId to string so it can be returned in JSON
         original_id = schema_document["_id"]
         schema_document["_id"] = str(original_id)
-
-        # Optional: Add logging or debugging info
-        # print(f"Converted ObjectId {original_id} to string for schema: {schema_document}")
-
     return schema_documents
 
-
 @app.post("/schemas")
 async def add_schema(schema: SchemaDefinition):
     """
     Add a new schema to the database.
     Args:
         schema (SchemaDefinition): The schema data to insert.
     Returns:
-        dict: The ID of the inserted schema.
+        dict: A dictionary with key being the ID of the inserted schema,
+              and value the actual inserted schema.
     """
     # Update timestamp before insertion
     schema.updated_at = datetime.utcnow()
 
     # Insert schema into MongoDB
     result = schemas_collection.insert_one(schema.dict())
+    inserted_id = result.inserted_id  # ObjectId
 
-    return {"id": str(result.inserted_id)}
+    # Fetch the stored document so we return what the DB actually has
+    inserted_doc = schemas_collection.find_one({"_id": inserted_id})
+    # Normalize _id for JSON
+    if inserted_doc is not None and "_id" in inserted_doc:
+        inserted_doc["_id"] = str(inserted_doc["_id"])
 
+    # Return as { "<id>": <document> } per your requirement
+    return {str(inserted_id): inserted_doc}
 
 @app.put("/schemas/{id}")
 async def update_schema(id: str, update: UpdateSchema):
@@ -82,13 +81,11 @@ async def update_schema(id: str, update: UpdateSchema):
         {"_id": ObjectId(id)},
         {"$set": update_fields},
     )
-
-    if result.modified_count == 0:
+    if result.modified_count == 0 and result.matched_count == 0:
         raise HTTPException(status_code=404, detail="Schema not found")
 
     return {"message": "Schema updated"}
 
-
 @app.delete("/schemas/{id}")
 async def delete_schema(id: str):
     """
@@ -101,6 +98,6 @@ async def delete_schema(id: str):
         dict: Success message.
     """
     result = schemas_collection.delete_one({"_id": ObjectId(id)})
-
     if result.deleted_count == 0:
         raise HTTPException(status_code=404, detail="Schema not found")
+    return {"message": "Schema deleted"}