Update app.py

app.py

@@ -32,6 +32,10 @@ import shutil
 from utils.dataops import auto_split_upscale
 from typing import List, Optional
 from datetime import datetime
+
+import asyncio
+from bson.objectid import ObjectId
+# -------------------
 # Writable directories (Spaces often allow /tmp). Override via env if needed.
 facex_misc.ROOT_DIR = os.environ["FACEXLIB_CACHE_DIR"]
 OUTPUT_DIR = os.environ.get("OUTPUT_DIR", "/tmp/face_upscale/output")
@@ -80,12 +84,23 @@ try:
     _mongo_uri = os.getenv("MONGODB_URI")
     _mongo_db = os.getenv("MONGODB_DB", "face_upscale")
     _mongo_col = os.getenv("MONGODB_COLLECTION", "submits")
-    if _mongo_uri:
-        _mongo_client = MongoClient(_mongo_uri, connect=False)
-        _mongo_collection = _mongo_client[_mongo_db][_mongo_col]
+    _admin_mongo_client = None # NEW client for media clicks
+    _media_clicks_col = None   # NEW collection for media clicks
+    # 2. Media Clicks Connection (using MONGODB_ADMIN_URI, as requested)
+    _admin_mongo_uri = os.getenv("MONGODB_ADMIN_URI")
+    _admin_mongo_db = os.getenv("MONGODB_ADMIN_DB", "face_upscale_admin") 
+    _media_clicks_col_name = "media_clicks" 
+    
+    if _admin_mongo_uri:
+        _admin_mongo_client = MongoClient(_admin_mongo_uri, connect=False)
+        _media_clicks_col = _admin_mongo_client[_admin_mongo_db][_media_clicks_col_name]
+    
 except Exception:
+    # If any MongoDB connection fails, ensure all related variables are None
     _mongo_client = None
     _mongo_collection = None
+    _admin_mongo_client = None
+    _media_clicks_col = None
 
 input_images_limit = 5
 
@@ -851,6 +866,60 @@ class Upscale:
         is_success, im_buf_arr = cv2.imencode(extension, image)
         if is_success:
             im_buf_arr.tofile(save_path)
+            
+    def sync_log_media_click(user_id_str: str, category_id_str: str):
+    """
+    Synchronously logs a click event to the media_clicks collection.
+    This function should be run using asyncio.to_thread().
+    """
+    if _media_clicks_col is None:
+        return
+
+    try:
+        user_oid = ObjectId(user_id_str.strip())
+        category_oid = ObjectId(category_id_str.strip())
+        now = datetime.utcnow()
+
+        # 1. Try updating an existing category entry within the user's document
+        update_result = _media_clicks_col.update_one(
+            {
+                "userId": user_oid,
+                "categories.categoryId": category_oid
+            },
+            {
+                "$set": {
+                    "updatedAt": now,
+                    "categories.$.lastClickedAt": now
+                },
+                "$inc": {
+                    "categories.$.click_count": 1
+                }
+            }
+        )
+
+        # 2. If no category entry exists (matched_count == 0) -> push a new one (or create user doc)
+        if update_result.matched_count == 0:
+            _media_clicks_col.update_one(
+                { "userId": user_oid },
+                {
+                    "$setOnInsert": { "createdAt": now },
+                    "$set": { "updatedAt": now },
+                    "$push": {
+                        "categories": {
+                            "categoryId": category_oid,
+                            "click_count": 1,
+                            "lastClickedAt": now
+                        }
+                    }
+                },
+                upsert=True
+            )
+        
+        # logger.info(f"Media click logged for User {user_id_str} on Category {category_id_str}")
+        
+    except Exception as media_err:
+        # logger.error(f"MEDIA_CLICK LOGGING ERROR: {media_err}")
+        pass # Handle logging error gracefully
 
 class Timer:
     def __init__(self):
@@ -888,6 +957,8 @@ class Timer:
         self.start_time = time.perf_counter()
         self.checkpoints = [("Start", self.start_time)]
 
+    
+
 if __name__ == "__main__":
     if torch.cuda.is_available():
         torch.cuda.set_per_process_memory_fraction(0.900, device='cuda:0')
@@ -1011,6 +1082,10 @@ async def submit(
     face_detection_threshold: float = Form(DEFAULT_FACE_DET_THRESHOLD),
     face_detection_only_center: bool = Form(DEFAULT_ONLY_CENTER),
     with_model_name: bool = Form(DEFAULT_WITH_MODEL_NAME),
+    # --- NEW OPTIONAL FIELDS ---
+    category_id: Optional[str] = Form(None), 
+    user_id: Optional[str] = Form(None),
+    # ---------------------------
     _: bool = Depends(_verify_bearer_token)
 ):
     start_time = datetime.utcnow()  # Start timer
@@ -1043,6 +1118,26 @@ async def submit(
             save_as_png=True,
             progress=progress_stub,
         )
+        # -----------------------------------------------------------------
+        # NEW: MEDIA CLICK LOGGING (AFTER SUCCESSFUL CORE PROCESSING)
+        # -----------------------------------------------------------------
+        # Only log if both optional fields are provided and look like valid ObjectIds
+        if user_id and category_id:
+            try:
+                # Validate the format before passing to the synchronous function
+                ObjectId(user_id.strip())
+                ObjectId(category_id.strip())
+                
+                # Run the blocking log operation in a separate thread safely
+                asyncio.to_thread(
+                    sync_log_media_click, 
+                    user_id, 
+                    category_id
+                )
+            except Exception as oid_err:
+                # Log a warning if the IDs were provided but invalid
+                # logger.warning(f"Skipping media click log due to invalid ID format: {oid_err}")
+                pass
 
         # --- Filter only final outputs ---
         final_images = [