feat: use jsonable_encoder with custom ObjectId encoder for clean REST API responses

Author: GorkiPower

backend/main.py

@@ -1,103 +1,133 @@
 
 from datetime import datetime
+from typing import Any, Dict, List
+
 from bson import ObjectId
-from database import schemas_collection
-from datamodel import SchemaDefinition, UpdateSchema
 from fastapi import FastAPI, HTTPException
+from fastapi.encoders import jsonable_encoder
 from fastapi.middleware.cors import CORSMiddleware
 
-# Initialize FastAPI application
+from database import schemas_collection
+from datamodel import SchemaDefinition, UpdateSchema
+
+
+# ---- FastAPI app & CORS ----
 app = FastAPI()
 
-# Enable CORS for all origins (useful for frontend integration)
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=["*"],        # adjust in production
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+
+# ---- Routes ----
 @app.get("/schemas")
-async def get_all_schemas():
+async def get_all_schemas() -> List[Dict[str, Any]]:
     """
     Retrieve all schema documents from the database.
-    Converts MongoDB ObjectId to string for JSON serialization.
+
     Returns:
-        list: A list of schema documents with stringified IDs.
+        list: A list of schema documents with JSON-safe '_id' fields (strings).
     """
-    # Fetch all schema documents from MongoDB
-    schema_documents = list(schemas_collection.find())
-    # Convert ObjectId to string for each document
-    for schema_document in schema_documents:
-        original_id = schema_document["_id"]
-        schema_document["_id"] = str(original_id)
-    return schema_documents
+    docs = list(schemas_collection.find())
+    # One-liner conversion: ObjectId -> str (applies even to nested ObjectIds)
+    return jsonable_encoder(docs, custom_encoder={ObjectId: str})
+
 
 @app.post("/schemas")
-async def add_schema(schema: SchemaDefinition):
+async def add_schema(schema: SchemaDefinition) -> Dict[str, Any]:
     """
     Add a new schema to the database.
+
     Args:
         schema (SchemaDefinition): The schema data to insert.
+
     Returns:
-        dict: A dictionary with key being the ID of the inserted schema,
-              and value the actual inserted schema.
+        dict: A dictionary with keys:
+              - "id": the ID of the inserted schema (string),
+              - "schema": the actual inserted schema document as stored in the database,
+                          with JSON-safe '_id' and other BSON types.
     """
-    # Update timestamp before insertion
+    # Set server-side timestamp
     schema.updated_at = datetime.utcnow()
 
-    # Insert schema into MongoDB
+    # Insert and obtain new ObjectId
     result = schemas_collection.insert_one(schema.dict())
-    inserted_id = result.inserted_id  # ObjectId
+    inserted_oid = result.inserted_id
+
+    # Fetch the stored document to return exactly what's in DB
+    inserted_doc = schemas_collection.find_one({"_id": inserted_oid})
+
+    # Make the document JSON-safe (ObjectId -> str, etc.)
+    inserted_doc_json = jsonable_encoder(inserted_doc, custom_encoder={ObjectId: str})
 
-    # 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"])
+    # Stable, tooling-friendly shape:
+    return {"id": str(inserted_oid), "schema": inserted_doc_json}
+
+    # If you want ID-as-key instead, return this:
+    # return {str(inserted_oid): inserted_doc_json}
 
-    # 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):
+async def update_schema(id: str, update: UpdateSchema) -> Dict[str, Any]:
     """
     Update an existing schema by ID.
+
     Args:
-        id (str): The schema ID.
+        id (str): The schema ID (string form of ObjectId).
         update (UpdateSchema): Fields to update.
+
     Raises:
         HTTPException: If the schema is not found.
+
     Returns:
-        dict: Success message.
+        dict: Success message and (optionally) the updated document.
     """
     # Prepare update fields (ignore None values)
-    update_fields = {
-        key: value for key, value in update.dict().items() if value is not None
-    }
+    update_fields = {k: v for k, v in update.dict().items() if v is not None}
+
+    # If anything is updated, refresh timestamp
+    if update_fields:
+        update_fields["updated_at"] = datetime.utcnow()
 
     result = schemas_collection.update_one(
         {"_id": ObjectId(id)},
-        {"$set": update_fields},
+        {"$set": update_fields} if update_fields else {}
     )
-    if result.modified_count == 0 and result.matched_count == 0:
+
+    if result.matched_count == 0:
         raise HTTPException(status_code=404, detail="Schema not found")
 
+    # Optionally return the updated document (commented out by default)
+    # updated_doc = schemas_collection.find_one({"_id": ObjectId(id)})
+    # return {
+    #     "message": "Schema updated",
+    #     "schema": jsonable_encoder(updated_doc, custom_encoder={ObjectId: str}),
+    # }
+
     return {"message": "Schema updated"}
 
+
 @app.delete("/schemas/{id}")
-async def delete_schema(id: str):
+async def delete_schema(id: str) -> Dict[str, str]:
     """
     Delete a schema by ID.
+
     Args:
-        id (str): The schema ID.
+        id (str): The schema ID (string form of ObjectId).
+
     Raises:
         HTTPException: If the schema is not found.
+
     Returns:
         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"}