LogicGoInfotechSpaces commited on
Commit
1ff09f8
·
1 Parent(s): acfa98e

Fix critical bugs: MongoDB connection handling and missing dependencies

Browse files

- 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

Files changed (2) hide show
  1. api/main.py +35 -9
  2. requirements.txt +1 -0
api/main.py CHANGED
@@ -64,11 +64,24 @@ app = FastAPI(title="Photo Object Removal API", version="1.0.0")
64
  file_store: Dict[str, Dict[str, str]] = {}
65
  logs: List[Dict[str, str]] = []
66
 
67
- MONGO_URI = os.environ.get("MONGO_URI")
68
- mongo_client = MongoClient(MONGO_URI)
69
- mongo_db = mongo_client["object_remover"]
70
- mongo_logs = mongo_db["api_logs"]
71
- grid_fs = GridFS(mongo_db)
 
 
 
 
 
 
 
 
 
 
 
 
 
72
 
73
  ADMIN_MONGO_URI = os.environ.get("MONGODB_ADMIN")
74
  DEFAULT_CATEGORY_ID = "69368f722e46bd68ae188984"
@@ -130,6 +143,11 @@ def _admin_logging_status() -> Dict[str, object]:
130
 
131
  def _save_upload_to_gridfs(upload: UploadFile, file_type: str) -> str:
132
  """Store an uploaded file into GridFS and return its ObjectId string."""
 
 
 
 
 
133
  data = upload.file.read()
134
  if not data:
135
  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:
144
 
145
  def _read_gridfs_bytes(file_id: str, expected_type: str) -> bytes:
146
  """Fetch raw bytes from GridFS and validate the stored type metadata."""
 
 
 
 
 
147
  try:
148
  oid = ObjectId(file_id)
149
  except Exception:
@@ -594,10 +617,13 @@ def inpaint(req: InpaintRequest, request: Request, _: None = Depends(bearer_auth
594
  if error_msg:
595
  log_doc["error"] = error_msg
596
 
597
- try:
598
- mongo_logs.insert_one(log_doc)
599
- except Exception as mongo_err:
600
- log.error(f"Mongo log insert failed: {mongo_err}")
 
 
 
601
 
602
  # @app.post("/inpaint")
603
  # def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
 
64
  file_store: Dict[str, Dict[str, str]] = {}
65
  logs: List[Dict[str, str]] = []
66
 
67
+ MONGO_URI = os.environ.get("MONGO_URI") or os.environ.get("MONGODB_URI")
68
+ mongo_client = None
69
+ mongo_db = None
70
+ mongo_logs = None
71
+ grid_fs = None
72
+
73
+ if MONGO_URI:
74
+ try:
75
+ mongo_client = MongoClient(MONGO_URI)
76
+ mongo_db = mongo_client["object_remover"]
77
+ mongo_logs = mongo_db["api_logs"]
78
+ grid_fs = GridFS(mongo_db)
79
+ log.info("MongoDB connection initialized successfully")
80
+ except Exception as err:
81
+ log.error("Failed to initialize MongoDB connection: %s", err)
82
+ log.warning("GridFS operations will be disabled. Set MONGO_URI or MONGODB_URI environment variable.")
83
+ else:
84
+ log.warning("MONGO_URI not set. GridFS operations will be disabled. Upload endpoints will not work.")
85
 
86
  ADMIN_MONGO_URI = os.environ.get("MONGODB_ADMIN")
87
  DEFAULT_CATEGORY_ID = "69368f722e46bd68ae188984"
 
143
 
144
  def _save_upload_to_gridfs(upload: UploadFile, file_type: str) -> str:
145
  """Store an uploaded file into GridFS and return its ObjectId string."""
146
+ if grid_fs is None:
147
+ raise HTTPException(
148
+ status_code=503,
149
+ detail="MongoDB/GridFS not configured. Set MONGO_URI or MONGODB_URI environment variable."
150
+ )
151
  data = upload.file.read()
152
  if not data:
153
  raise HTTPException(status_code=400, detail=f"{file_type} file is empty")
 
162
 
163
  def _read_gridfs_bytes(file_id: str, expected_type: str) -> bytes:
164
  """Fetch raw bytes from GridFS and validate the stored type metadata."""
165
+ if grid_fs is None:
166
+ raise HTTPException(
167
+ status_code=503,
168
+ detail="MongoDB/GridFS not configured. Set MONGO_URI or MONGODB_URI environment variable."
169
+ )
170
  try:
171
  oid = ObjectId(file_id)
172
  except Exception:
 
617
  if error_msg:
618
  log_doc["error"] = error_msg
619
 
620
+ if mongo_logs is not None:
621
+ try:
622
+ mongo_logs.insert_one(log_doc)
623
+ except Exception as mongo_err:
624
+ log.error("Mongo log insert failed: %s", mongo_err)
625
+ else:
626
+ log.warning("MongoDB not configured, skipping log insert")
627
 
628
  # @app.post("/inpaint")
629
  # def inpaint(req: InpaintRequest, _: None = Depends(bearer_auth)) -> Dict[str, str]:
requirements.txt CHANGED
@@ -9,3 +9,4 @@ opencv-python-headless
9
  Pillow
10
  pymongo
11
  streamlit==1.24.1
 
 
9
  Pillow
10
  pymongo
11
  streamlit==1.24.1
12
+ streamlit-drawable-canvas