Fix critical bugs: MongoDB connection handling and missing dependencies

- Fix CRITICAL: Add null check for MONGO_URI before creating MongoClient to prevent startup crash
- Fix CRITICAL: Add error handling for GridFS operations when MongoDB is not configured
- Add missing dependency: streamlit-drawable-canvas to requirements.txt
- Add support for both MONGO_URI and MONGODB_URI environment variables
- Add proper error handling for all MongoDB logging operations
- Improve error messages and logging for better debugging

api/main.py

@@ -64,11 +64,24 @@ app = FastAPI(title="Photo Object Removal API", version="1.0.0")
 file_store: Dict[str, Dict[str, str]] = {}
 logs: List[Dict[str, str]] = []
 
-MONGO_URI = os.environ.get("MONGO_URI")
-mongo_client = MongoClient(MONGO_URI)
-mongo_db = mongo_client["object_remover"]
-mongo_logs = mongo_db["api_logs"]
-grid_fs = GridFS(mongo_db)
+MONGO_URI = os.environ.get("MONGO_URI") or os.environ.get("MONGODB_URI")
+mongo_client = None
+mongo_db = None
+mongo_logs = None
+grid_fs = None
+
+if MONGO_URI:
+    try:
+        mongo_client = MongoClient(MONGO_URI)
+        mongo_db = mongo_client["object_remover"]
+        mongo_logs = mongo_db["api_logs"]
+        grid_fs = GridFS(mongo_db)
+        log.info("MongoDB connection initialized successfully")
+    except Exception as err:
+        log.error("Failed to initialize MongoDB connection: %s", err)
+        log.warning("GridFS operations will be disabled. Set MONGO_URI or MONGODB_URI environment variable.")
+else:
+    log.warning("MONGO_URI not set. GridFS operations will be disabled. Upload endpoints will not work.")
 
 ADMIN_MONGO_URI = os.environ.get("MONGODB_ADMIN")
 DEFAULT_CATEGORY_ID = "69368f722e46bd68ae188984"
@@ -130,6 +143,11 @@ def _admin_logging_status() -> Dict[str, object]:
 
 def _save_upload_to_gridfs(upload: UploadFile, file_type: str) -> str:
     """Store an uploaded file into GridFS and return its ObjectId string."""
+    if grid_fs is None:
+        raise HTTPException(
+            status_code=503,
+            detail="MongoDB/GridFS not configured. Set MONGO_URI or MONGODB_URI environment variable."
+        )
     data = upload.file.read()
     if not data:
         raise HTTPException(status_code=400, detail=f"{file_type} file is empty")
@@ -144,6 +162,11 @@ def _save_upload_to_gridfs(upload: UploadFile, file_type: str) -> str:
 
 def _read_gridfs_bytes(file_id: str, expected_type: str) -> bytes:
     """Fetch raw bytes from GridFS and validate the stored type metadata."""
+    if grid_fs is None:
+        raise HTTPException(
+            status_code=503,
+            detail="MongoDB/GridFS not configured. Set MONGO_URI or MONGODB_URI environment variable."
+        )
     try:
         oid = ObjectId(file_id)
     except Exception:
@@ -594,10 +617,13 @@ def inpaint(req: InpaintRequest, request: Request, _: None = Depends(bearer_auth
         if error_msg:
             log_doc["error"] = error_msg
 
-        try:
-            mongo_logs.insert_one(log_doc)
-        except Exception as mongo_err:
-            log.error(f"Mongo log insert failed: {mongo_err}")
+        if mongo_logs is not None:
+            try:
+                mongo_logs.insert_one(log_doc)
+            except Exception as mongo_err:
+                log.error("Mongo log insert failed: %s", mongo_err)
+        else:
+            log.warning("MongoDB not configured, skipping log insert")
 
 # @app.post("/inpaint")
 # def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:

requirements.txt

@@ -9,3 +9,4 @@ opencv-python-headless
 Pillow
 pymongo
 streamlit==1.24.1
+streamlit-drawable-canvas