""" MongoDB database connection and logging utilities, including admin media click logging. """ import os import logging from datetime import datetime, timedelta from typing import Optional, Dict, Any, Union, List from bson import ObjectId from pymongo import MongoClient from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError logger = logging.getLogger(__name__) # MongoDB connection _client: Optional[MongoClient] = None _db = None # Admin MongoDB connection (media_clicks) _admin_client: Optional[MongoClient] = None _admin_db = None def get_mongodb_client() -> Optional[MongoClient]: """Get or create MongoDB client""" global _client if _client is None: mongodb_uri = os.getenv("MONGODB_URI") if not mongodb_uri: logger.warning("MONGODB_URI environment variable not set. MongoDB features will be disabled.") return None try: _client = MongoClient( mongodb_uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=5000 ) # Test connection _client.admin.command('ping') logger.info("MongoDB connection established successfully") except (ConnectionFailure, ServerSelectionTimeoutError) as e: logger.error("Failed to connect to MongoDB: %s", str(e)) _client = None return _client def get_database(): """Get database instance""" global _db if _db is None: client = get_mongodb_client() if client: db_name = os.getenv("MONGODB_DB_NAME", "colorization_db") _db = client[db_name] else: logger.warning("MongoDB client not available") return _db def get_admin_client() -> Optional[MongoClient]: """Get or create admin MongoDB client (for media_clicks collection).""" global _admin_client if _admin_client is None: mongodb_uri = os.getenv("MONGODB_ADMIN") if not mongodb_uri: logger.warning("MONGODB_ADMIN environment variable not set. Admin MongoDB features will be disabled.") return None try: _admin_client = MongoClient( mongodb_uri, serverSelectionTimeoutMS=5000, connectTimeoutMS=5000, ) _admin_client.admin.command("ping") logger.info("Admin MongoDB connection established successfully") except (ConnectionFailure, ServerSelectionTimeoutError) as exc: logger.error("Failed to connect to Admin MongoDB: %s", str(exc)) _admin_client = None return _admin_client def get_admin_database(): """Get admin database instance.""" global _admin_db if _admin_db is None: client = get_admin_client() if client: db_name = os.getenv("MONGODB_ADMIN_DB_NAME", "adminPanel") _admin_db = client[db_name] logger.info("Using admin database: %s", db_name) else: logger.warning("Admin MongoDB client not available") return _admin_db def _normalize_object_id(raw_value: Optional[Union[str, int, ObjectId]]) -> ObjectId: """Normalize user id inputs into a deterministic ObjectId.""" if isinstance(raw_value, ObjectId): return raw_value if raw_value is None: return ObjectId() # Handle empty strings as None if isinstance(raw_value, str) and not raw_value.strip(): return ObjectId() try: if isinstance(raw_value, int) or (isinstance(raw_value, str) and raw_value.strip().lstrip("-").isdigit()): int_value = int(str(raw_value).strip()) hex_str = format(abs(int_value), "x").zfill(24)[-24:] if ObjectId.is_valid(hex_str): return ObjectId(hex_str) except Exception as exc: # pragma: no cover - defensive logger.debug("Numeric user id normalization failed: %s", str(exc)) if isinstance(raw_value, str): candidate = raw_value.strip() if ObjectId.is_valid(candidate): return ObjectId(candidate) return ObjectId() def _objectid_from_any(value: str) -> ObjectId: """Convert arbitrary string into an ObjectId deterministically when possible.""" if ObjectId.is_valid(value): return ObjectId(value) try: hex_str = value.encode("utf-8").hex().zfill(24)[-24:] if ObjectId.is_valid(hex_str): return ObjectId(hex_str) except Exception as exc: # pragma: no cover - defensive logger.debug("Category id normalization failed: %s", str(exc)) return ObjectId() def _resolve_category_id( category_id: Optional[str], endpoint_path: Optional[str], default_category_id: Optional[str], ) -> ObjectId: """Pick category id from explicit value, endpoint default, or fallback.""" # Handle empty strings as None if isinstance(category_id, str) and not category_id.strip(): category_id = None endpoint_map = { "colorization": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), "upload": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), "colorize": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), } normalized_endpoint = None if endpoint_path: normalized_endpoint = endpoint_path.strip("/").split("/")[0].lower() or None chosen = category_id if not chosen and normalized_endpoint and endpoint_map.get(normalized_endpoint): chosen = endpoint_map[normalized_endpoint] if not chosen: chosen = default_category_id or os.getenv("DEFAULT_CATEGORY_FALLBACK", "69368fcd2e46bd68ae1889b2") return _objectid_from_any(chosen) def _get_today_midnight_utc() -> datetime: """Get today's date at midnight UTC (timezone-naive for MongoDB compatibility).""" now = datetime.utcnow() return datetime(now.year, now.month, now.day) def _update_daily_count(collection, user_object_id: ObjectId, today: datetime) -> None: """Update ai_edit_daily_count array: add today with count 1 if missing, fill gaps with count 0.""" # Ensure today is at midnight today_start = today.replace(hour=0, minute=0, second=0, microsecond=0) logger.info("Updating daily count for user %s, today: %s", str(user_object_id), today_start.isoformat()) # Check if today's date already exists user_doc = collection.find_one({"userId": user_object_id}) if not user_doc: # User document doesn't exist yet - this shouldn't happen as we create it first # But handle it gracefully by initializing the field logger.warning("User document not found when updating daily count, initializing") collection.update_one( {"userId": user_object_id}, { "$setOnInsert": { "ai_edit_daily_count": [{ "date": today_start, "count": 1 }] } }, upsert=False ) return # Get existing daily counts existing_counts = user_doc.get("ai_edit_daily_count", []) logger.debug("Existing daily counts: %d entries", len(existing_counts)) # Check if today's date already exists today_exists = False for entry in existing_counts: entry_date = entry.get("date") if entry_date: # Normalize to midnight for comparison if isinstance(entry_date, datetime): normalized_date = entry_date.replace(hour=0, minute=0, second=0, microsecond=0) logger.debug("Comparing entry date %s with today %s", normalized_date.isoformat(), today_start.isoformat()) if normalized_date == today_start: today_exists = True logger.info("Today's date already exists in daily count, leaving unchanged: %s", today_start.isoformat()) break # If today exists, do nothing (leave it as is) if today_exists: logger.debug("Today's date already exists in daily count, leaving unchanged: %s", today_start.isoformat()) return # Today doesn't exist - need to add it and fill missing dates logger.info("Today's date does not exist in daily count, adding new entry") # Find the latest date in existing counts last_date = None if existing_counts: dates = [] for entry in existing_counts: entry_date = entry.get("date") if entry_date: # Normalize to midnight for comparison if isinstance(entry_date, datetime): dates.append(entry_date.replace(hour=0, minute=0, second=0, microsecond=0)) if dates: last_date = max(dates) logger.debug("Last date found: %s", last_date.isoformat()) else: logger.debug("No valid dates found in existing counts") else: logger.debug("No existing counts, will add today as first entry") # Generate missing dates between last_date and today dates_to_add = [] if last_date: # Normalize last_date to midnight last_date = last_date.replace(hour=0, minute=0, second=0, microsecond=0) current_date = last_date + timedelta(days=1) # Fill missing dates with count 0 while current_date < today_start: dates_to_add.append({ "date": current_date, "count": 0 }) current_date += timedelta(days=1) logger.debug("Filled %d gap dates between %s and %s", len(dates_to_add), last_date.isoformat(), today_start.isoformat()) # Add today's entry with count 1 dates_to_add.append({ "date": today_start, "count": 1 }) logger.info("Adding %d new date entries (including today with count 1)", len(dates_to_add)) # Merge existing entries with the new ones, sort by date (oldest first), # and keep only the most recent 32 dates (drop the oldest beyond 32). if dates_to_add: all_entries = list(existing_counts) + dates_to_add def _entry_sort_key(entry: Dict[str, Any]) -> datetime: dt = entry.get("date") if isinstance(dt, datetime): return dt.replace(hour=0, minute=0, second=0, microsecond=0) return datetime.min all_entries.sort(key=_entry_sort_key) if len(all_entries) > 32: removed = len(all_entries) - 32 all_entries = all_entries[-32:] logger.debug("Removed %d oldest entries to maintain 32-entry limit", removed) result = collection.update_one( {"userId": user_object_id}, {"$set": {"ai_edit_daily_count": all_entries}}, ) logger.info( "Updated ai_edit_daily_count with %d entries (oldest first, max 32). Matched: %d, Modified: %d", len(all_entries), result.matched_count, result.modified_count, ) else: logger.warning("No dates to add - this should not happen!") def log_api_call( endpoint: str, method: str, status_code: int = 200, request_data: Optional[Dict[str, Any]] = None, response_data: Optional[Dict[str, Any]] = None, error: Optional[str] = None, user_id: Optional[str] = None, ip_address: Optional[str] = None ) -> bool: """ Log API call to MongoDB Args: endpoint: API endpoint path method: HTTP method (GET, POST, etc.) status_code: HTTP status code request_data: Request data/parameters response_data: Response data error: Error message if any user_id: User ID if authenticated ip_address: Client IP address Returns: True if logged successfully, False otherwise """ try: db = get_database() if db is None: logger.warning("MongoDB not available, skipping API log") return False collection = db["api_calls"] log_entry = { "endpoint": endpoint, "method": method, "status_code": status_code, "timestamp": datetime.utcnow(), "request_data": request_data or {}, "response_data": response_data or {}, "error": error, "user_id": user_id, "ip_address": ip_address } result = collection.insert_one(log_entry) logger.info("API call logged to MongoDB: %s", result.inserted_id) return True except Exception as e: logger.error("Failed to log API call to MongoDB: %s", str(e)) return False def log_image_upload( image_id: str, filename: str, file_size: int, content_type: str, user_id: Optional[str] = None, ip_address: Optional[str] = None ) -> bool: """ Log image upload to MongoDB Args: image_id: Unique image identifier filename: Original filename file_size: File size in bytes content_type: MIME type user_id: User ID if authenticated ip_address: Client IP address Returns: True if logged successfully, False otherwise """ try: db = get_database() if db is None: logger.warning("MongoDB not available, skipping upload log") return False collection = db["image_uploads"] log_entry = { "image_id": image_id, "filename": filename, "file_size": file_size, "content_type": content_type, "uploaded_at": datetime.utcnow(), "user_id": user_id, "ip_address": ip_address } result = collection.insert_one(log_entry) logger.info("Image upload logged to MongoDB: %s", result.inserted_id) return True except Exception as e: logger.error("Failed to log image upload to MongoDB: %s", str(e)) return False def log_colorization( result_id: Optional[str] = None, image_id: Optional[str] = None, prompt: Optional[str] = None, model_type: Optional[str] = None, processing_time: Optional[float] = None, user_id: Optional[str] = None, ip_address: Optional[str] = None, status: str = "success", error: Optional[str] = None ) -> bool: """ Log colorization request to MongoDB (colorization_db -> colorizations collection) Logs both successful and failed API calls. Args: result_id: Unique result identifier (None for failed requests) image_id: Original image identifier prompt: Text prompt used (if any) model_type: Model type used (fastai, pytorch, sdxl, etc.) processing_time: Time taken to process in seconds user_id: User ID if authenticated ip_address: Client IP address status: Status of the request ("success" or "failed") error: Error message if status is "failed" Returns: True if logged successfully, False otherwise """ try: db = get_database() if db is None: logger.warning("MongoDB not available, skipping colorization log") return False collection = db["colorizations"] # Generate result_id if not provided (for failed requests) if not result_id: import uuid result_id = str(uuid.uuid4()) log_entry = { "result_id": result_id, "image_id": image_id, "prompt": prompt, "model_type": model_type, "processing_time": processing_time, "created_at": datetime.utcnow(), "user_id": user_id, "ip_address": ip_address } # Add status and error fields only if provided (for new documents) # Existing documents won't have these fields, which is fine if status: log_entry["status"] = status if error: log_entry["error"] = error result = collection.insert_one(log_entry) logger.info("Colorization logged to MongoDB (status: %s): %s", status, result.inserted_id) return True except Exception as e: logger.error("Failed to log colorization to MongoDB: %s", str(e)) return False def log_media_click( user_id: Optional[Union[str, int, ObjectId]], category_id: Optional[str], *, endpoint_path: Optional[str] = None, default_category_id: Optional[str] = None, ) -> bool: """Log media clicks into the admin MongoDB (media_clicks collection). Only logs if user_id is provided. If user_id is None or empty, returns False without logging. Regular MongoDB logging (api_calls, image_uploads, colorizations) always happens regardless. """ # Only log to media_clicks if user_id is provided if not user_id or (isinstance(user_id, str) and not user_id.strip()): logger.debug("Skipping media click log - user_id not provided") return False try: db = get_admin_database() if db is None: logger.warning("Admin MongoDB not available, skipping media click log") return False collection = db["media_clicks"] # Drop legacy index to avoid duplicate key errors (best effort) try: collection.drop_index("user_id_1_header_1_media_id_1") except Exception as exc: logger.debug("Legacy index drop skipped: %s", str(exc)) user_object_id = _normalize_object_id(user_id) category_object_id = _resolve_category_id(category_id, endpoint_path, default_category_id) now = datetime.utcnow() today = _get_today_midnight_utc() logger.info("Media click - user_id input: %s, normalized: %s, category_id input: %s, normalized: %s", user_id, str(user_object_id), category_id, str(category_object_id)) # Try to update existing category in new schema (categories.categoryId) update_existing = collection.update_one( {"userId": user_object_id, "categories.categoryId": category_object_id}, { "$inc": { "categories.$.click_count": 1, "ai_edit_complete": 1 # Increment usage count }, "$set": { "categories.$.lastClickedAt": now, "updatedAt": now, "ai_edit_last_date": now # Update last used date (MongoDB Date object) }, }, ) if update_existing.matched_count == 0: # Category not found in new schema, check if user exists user_exists = collection.find_one({"userId": user_object_id}) if user_exists: # User exists but category doesn't in new schema - add to categories array collection.update_one( {"userId": user_object_id}, { "$inc": {"ai_edit_complete": 1}, # Increment usage count "$set": { "updatedAt": now, "ai_edit_last_date": now # Update last used date (MongoDB Date object) }, "$push": { "categories": { "categoryId": category_object_id, "click_count": 1, "lastClickedAt": now, } }, }, ) else: # User doesn't exist - create new document with categories array # Default ai_edit_complete = 0, then increment to 1 on first use collection.update_one( {"userId": user_object_id}, { "$setOnInsert": { "createdAt": now, "ai_edit_complete": 0, # Default 0 for new users "ai_edit_daily_count": [] # Initialize empty array }, "$inc": {"ai_edit_complete": 1}, # Increment to 1 on first use "$set": { "updatedAt": now, "ai_edit_last_date": now # Set last used date (MongoDB Date object) }, "$push": { "categories": { "categoryId": category_object_id, "click_count": 1, "lastClickedAt": now, } }, }, upsert=True, ) # Update daily count (after document is created/updated) _update_daily_count(collection, user_object_id, today) logger.info("Media click logged for user %s", str(user_object_id)) return True except Exception as exc: logger.error("Failed to log media click to admin MongoDB: %s", str(exc)) return False def close_connection(): """Close MongoDB connection""" global _client, _db, _admin_client, _admin_db if _client: _client.close() _client = None _db = None logger.info("MongoDB connection closed") if _admin_client: _admin_client.close() _admin_client = None _admin_db = None logger.info("Admin MongoDB connection closed") # """ # MongoDB database connection and logging utilities, including admin media click logging. # """ # import os # import logging # from datetime import datetime, timedelta # from typing import Optional, Dict, Any, Union, List # from bson import ObjectId # from pymongo import MongoClient # from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError # logger = logging.getLogger(__name__) # # MongoDB connection # _client: Optional[MongoClient] = None # _db = None # # Admin MongoDB connection (media_clicks) # _admin_client: Optional[MongoClient] = None # _admin_db = None # def get_mongodb_client() -> Optional[MongoClient]: # """Get or create MongoDB client""" # global _client # if _client is None: # mongodb_uri = os.getenv("MONGODB_URI") # if not mongodb_uri: # logger.warning("MONGODB_URI environment variable not set. MongoDB features will be disabled.") # return None # try: # _client = MongoClient( # mongodb_uri, # serverSelectionTimeoutMS=5000, # connectTimeoutMS=5000 # ) # # Test connection # _client.admin.command('ping') # logger.info("MongoDB connection established successfully") # except (ConnectionFailure, ServerSelectionTimeoutError) as e: # logger.error("Failed to connect to MongoDB: %s", str(e)) # _client = None # return _client # def get_database(): # """Get database instance""" # global _db # if _db is None: # client = get_mongodb_client() # if client: # db_name = os.getenv("MONGODB_DB_NAME", "colorization_db") # _db = client[db_name] # else: # logger.warning("MongoDB client not available") # return _db # def get_admin_client() -> Optional[MongoClient]: # """Get or create admin MongoDB client (for media_clicks collection).""" # global _admin_client # if _admin_client is None: # mongodb_uri = os.getenv("MONGODB_ADMIN") # if not mongodb_uri: # logger.warning("MONGODB_ADMIN environment variable not set. Admin MongoDB features will be disabled.") # return None # try: # _admin_client = MongoClient( # mongodb_uri, # serverSelectionTimeoutMS=5000, # connectTimeoutMS=5000, # ) # _admin_client.admin.command("ping") # logger.info("Admin MongoDB connection established successfully") # except (ConnectionFailure, ServerSelectionTimeoutError) as exc: # logger.error("Failed to connect to Admin MongoDB: %s", str(exc)) # _admin_client = None # return _admin_client # def get_admin_database(): # """Get admin database instance.""" # global _admin_db # if _admin_db is None: # client = get_admin_client() # if client: # db_name = os.getenv("MONGODB_ADMIN_DB_NAME", "adminPanel") # _admin_db = client[db_name] # logger.info("Using admin database: %s", db_name) # else: # logger.warning("Admin MongoDB client not available") # return _admin_db # def _normalize_object_id(raw_value: Optional[Union[str, int, ObjectId]]) -> ObjectId: # """Normalize user id inputs into a deterministic ObjectId.""" # if isinstance(raw_value, ObjectId): # return raw_value # if raw_value is None: # return ObjectId() # # Handle empty strings as None # if isinstance(raw_value, str) and not raw_value.strip(): # return ObjectId() # try: # if isinstance(raw_value, int) or (isinstance(raw_value, str) and raw_value.strip().lstrip("-").isdigit()): # int_value = int(str(raw_value).strip()) # hex_str = format(abs(int_value), "x").zfill(24)[-24:] # if ObjectId.is_valid(hex_str): # return ObjectId(hex_str) # except Exception as exc: # pragma: no cover - defensive # logger.debug("Numeric user id normalization failed: %s", str(exc)) # if isinstance(raw_value, str): # candidate = raw_value.strip() # if ObjectId.is_valid(candidate): # return ObjectId(candidate) # return ObjectId() # def _objectid_from_any(value: str) -> ObjectId: # """Convert arbitrary string into an ObjectId deterministically when possible.""" # if ObjectId.is_valid(value): # return ObjectId(value) # try: # hex_str = value.encode("utf-8").hex().zfill(24)[-24:] # if ObjectId.is_valid(hex_str): # return ObjectId(hex_str) # except Exception as exc: # pragma: no cover - defensive # logger.debug("Category id normalization failed: %s", str(exc)) # return ObjectId() # def _resolve_category_id( # category_id: Optional[str], # endpoint_path: Optional[str], # default_category_id: Optional[str], # ) -> ObjectId: # """Pick category id from explicit value, endpoint default, or fallback.""" # # Handle empty strings as None # if isinstance(category_id, str) and not category_id.strip(): # category_id = None # endpoint_map = { # "colorization": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), # "upload": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), # "colorize": os.getenv("DEFAULT_CATEGORY_COLORIZATION"), # } # normalized_endpoint = None # if endpoint_path: # normalized_endpoint = endpoint_path.strip("/").split("/")[0].lower() or None # chosen = category_id # if not chosen and normalized_endpoint and endpoint_map.get(normalized_endpoint): # chosen = endpoint_map[normalized_endpoint] # if not chosen: # chosen = default_category_id or os.getenv("DEFAULT_CATEGORY_FALLBACK", "69368fcd2e46bd68ae1889b2") # return _objectid_from_any(chosen) # def _get_today_midnight_utc() -> datetime: # """Get today's date at midnight UTC (timezone-naive for MongoDB compatibility).""" # now = datetime.utcnow() # return datetime(now.year, now.month, now.day) # def _update_daily_count(collection, user_object_id: ObjectId, today: datetime) -> None: # """Update ai_edit_daily_count array: add today with count 1 if missing, fill gaps with count 0.""" # # Ensure today is at midnight # today_start = today.replace(hour=0, minute=0, second=0, microsecond=0) # # Check if today's date already exists # user_doc = collection.find_one({"userId": user_object_id}) # if not user_doc: # # User document doesn't exist yet - this shouldn't happen as we create it first # # But handle it gracefully by initializing the field # logger.warning("User document not found when updating daily count, initializing") # collection.update_one( # {"userId": user_object_id}, # { # "$setOnInsert": { # "ai_edit_daily_count": [{ # "date": today_start, # "count": 1 # }] # } # }, # upsert=False # ) # return # # Get existing daily counts # existing_counts = user_doc.get("ai_edit_daily_count", []) # # Check if today's date already exists # today_exists = False # for entry in existing_counts: # entry_date = entry.get("date") # if entry_date: # # Normalize to midnight for comparison # if isinstance(entry_date, datetime): # normalized_date = entry_date.replace(hour=0, minute=0, second=0, microsecond=0) # if normalized_date == today_start: # today_exists = True # break # # If today exists, do nothing (leave it as is) # if today_exists: # logger.debug("Today's date already exists in daily count, leaving unchanged: %s", today_start.isoformat()) # return # # Today doesn't exist - need to add it and fill missing dates # # Find the latest date in existing counts # last_date = None # if existing_counts: # dates = [] # for entry in existing_counts: # entry_date = entry.get("date") # if entry_date: # # Normalize to midnight for comparison # if isinstance(entry_date, datetime): # dates.append(entry_date.replace(hour=0, minute=0, second=0, microsecond=0)) # if dates: # last_date = max(dates) # # Generate missing dates between last_date and today # dates_to_add = [] # if last_date: # # Normalize last_date to midnight # last_date = last_date.replace(hour=0, minute=0, second=0, microsecond=0) # current_date = last_date + timedelta(days=1) # # Fill missing dates with count 0 # while current_date < today_start: # dates_to_add.append({ # "date": current_date, # "count": 0 # }) # current_date += timedelta(days=1) # # Add today's entry with count 1 # dates_to_add.append({ # "date": today_start, # "count": 1 # }) # # Merge existing entries with the new ones, sort by date (oldest first), # # and keep only the most recent 32 dates (drop the oldest beyond 32). # if dates_to_add: # all_entries = list(existing_counts) + dates_to_add # def _entry_sort_key(entry: Dict[str, Any]) -> datetime: # dt = entry.get("date") # if isinstance(dt, datetime): # return dt.replace(hour=0, minute=0, second=0, microsecond=0) # return datetime.min # all_entries.sort(key=_entry_sort_key) # if len(all_entries) > 32: # all_entries = all_entries[-32:] # collection.update_one( # {"userId": user_object_id}, # {"$set": {"ai_edit_daily_count": all_entries}}, # ) # logger.debug( # "Updated ai_edit_daily_count with %d entries (oldest first, max 32)", # len(all_entries), # ) # def log_api_call( # endpoint: str, # method: str, # status_code: int = 200, # request_data: Optional[Dict[str, Any]] = None, # response_data: Optional[Dict[str, Any]] = None, # error: Optional[str] = None, # user_id: Optional[str] = None, # ip_address: Optional[str] = None # ) -> bool: # """ # Log API call to MongoDB # Args: # endpoint: API endpoint path # method: HTTP method (GET, POST, etc.) # status_code: HTTP status code # request_data: Request data/parameters # response_data: Response data # error: Error message if any # user_id: User ID if authenticated # ip_address: Client IP address # Returns: # True if logged successfully, False otherwise # """ # try: # db = get_database() # if db is None: # logger.warning("MongoDB not available, skipping API log") # return False # collection = db["api_calls"] # log_entry = { # "endpoint": endpoint, # "method": method, # "status_code": status_code, # "timestamp": datetime.utcnow(), # "request_data": request_data or {}, # "response_data": response_data or {}, # "error": error, # "user_id": user_id, # "ip_address": ip_address # } # result = collection.insert_one(log_entry) # logger.info("API call logged to MongoDB: %s", result.inserted_id) # return True # except Exception as e: # logger.error("Failed to log API call to MongoDB: %s", str(e)) # return False # def log_image_upload( # image_id: str, # filename: str, # file_size: int, # content_type: str, # user_id: Optional[str] = None, # ip_address: Optional[str] = None # ) -> bool: # """ # Log image upload to MongoDB # Args: # image_id: Unique image identifier # filename: Original filename # file_size: File size in bytes # content_type: MIME type # user_id: User ID if authenticated # ip_address: Client IP address # Returns: # True if logged successfully, False otherwise # """ # try: # db = get_database() # if db is None: # logger.warning("MongoDB not available, skipping upload log") # return False # collection = db["image_uploads"] # log_entry = { # "image_id": image_id, # "filename": filename, # "file_size": file_size, # "content_type": content_type, # "uploaded_at": datetime.utcnow(), # "user_id": user_id, # "ip_address": ip_address # } # result = collection.insert_one(log_entry) # logger.info("Image upload logged to MongoDB: %s", result.inserted_id) # return True # except Exception as e: # logger.error("Failed to log image upload to MongoDB: %s", str(e)) # return False # def log_colorization( # result_id: Optional[str] = None, # image_id: Optional[str] = None, # prompt: Optional[str] = None, # model_type: Optional[str] = None, # processing_time: Optional[float] = None, # user_id: Optional[str] = None, # ip_address: Optional[str] = None, # status: str = "success", # error: Optional[str] = None # ) -> bool: # """ # Log colorization request to MongoDB (colorization_db -> colorizations collection) # Logs both successful and failed API calls. # Args: # result_id: Unique result identifier (None for failed requests) # image_id: Original image identifier # prompt: Text prompt used (if any) # model_type: Model type used (fastai, pytorch, sdxl, etc.) # processing_time: Time taken to process in seconds # user_id: User ID if authenticated # ip_address: Client IP address # status: Status of the request ("success" or "failed") # error: Error message if status is "failed" # Returns: # True if logged successfully, False otherwise # """ # try: # db = get_database() # if db is None: # logger.warning("MongoDB not available, skipping colorization log") # return False # collection = db["colorizations"] # # Generate result_id if not provided (for failed requests) # if not result_id: # import uuid # result_id = str(uuid.uuid4()) # log_entry = { # "result_id": result_id, # "image_id": image_id, # "prompt": prompt, # "model_type": model_type, # "processing_time": processing_time, # "created_at": datetime.utcnow(), # "user_id": user_id, # "ip_address": ip_address # } # # Add status and error fields only if provided (for new documents) # # Existing documents won't have these fields, which is fine # if status: # log_entry["status"] = status # if error: # log_entry["error"] = error # result = collection.insert_one(log_entry) # logger.info("Colorization logged to MongoDB (status: %s): %s", status, result.inserted_id) # return True # except Exception as e: # logger.error("Failed to log colorization to MongoDB: %s", str(e)) # return False # def log_media_click( # user_id: Optional[Union[str, int, ObjectId]], # category_id: Optional[str], # *, # endpoint_path: Optional[str] = None, # default_category_id: Optional[str] = None, # ) -> bool: # """Log media clicks into the admin MongoDB (media_clicks collection). # Only logs if user_id is provided. If user_id is None or empty, returns False without logging. # Regular MongoDB logging (api_calls, image_uploads, colorizations) always happens regardless. # """ # # Only log to media_clicks if user_id is provided # if not user_id or (isinstance(user_id, str) and not user_id.strip()): # logger.debug("Skipping media click log - user_id not provided") # return False # try: # db = get_admin_database() # if db is None: # logger.warning("Admin MongoDB not available, skipping media click log") # return False # collection = db["media_clicks"] # # Drop legacy index to avoid duplicate key errors (best effort) # try: # collection.drop_index("user_id_1_header_1_media_id_1") # except Exception as exc: # logger.debug("Legacy index drop skipped: %s", str(exc)) # user_object_id = _normalize_object_id(user_id) # category_object_id = _resolve_category_id(category_id, endpoint_path, default_category_id) # now = datetime.utcnow() # today = _get_today_midnight_utc() # logger.info("Media click - user_id input: %s, normalized: %s, category_id input: %s, normalized: %s", # user_id, str(user_object_id), category_id, str(category_object_id)) # # Try to update existing category in new schema (categories.categoryId) # update_existing = collection.update_one( # {"userId": user_object_id, "categories.categoryId": category_object_id}, # { # "$inc": { # "categories.$.click_count": 1, # "ai_edit_complete": 1 # Increment usage count # }, # "$set": { # "categories.$.lastClickedAt": now, # "updatedAt": now, # "ai_edit_last_date": now # Update last used date (MongoDB Date object) # }, # }, # ) # if update_existing.matched_count == 0: # # Category not found in new schema, check if user exists # user_exists = collection.find_one({"userId": user_object_id}) # if user_exists: # # User exists but category doesn't in new schema - add to categories array # collection.update_one( # {"userId": user_object_id}, # { # "$inc": {"ai_edit_complete": 1}, # Increment usage count # "$set": { # "updatedAt": now, # "ai_edit_last_date": now # Update last used date (MongoDB Date object) # }, # "$push": { # "categories": { # "categoryId": category_object_id, # "click_count": 1, # "lastClickedAt": now, # } # }, # }, # ) # else: # # User doesn't exist - create new document with categories array # # Default ai_edit_complete = 0, then increment to 1 on first use # collection.update_one( # {"userId": user_object_id}, # { # "$setOnInsert": { # "createdAt": now, # "ai_edit_complete": 0, # Default 0 for new users # "ai_edit_daily_count": [] # Initialize empty array # }, # "$inc": {"ai_edit_complete": 1}, # Increment to 1 on first use # "$set": { # "updatedAt": now, # "ai_edit_last_date": now # Set last used date (MongoDB Date object) # }, # "$push": { # "categories": { # "categoryId": category_object_id, # "click_count": 1, # "lastClickedAt": now, # } # }, # }, # upsert=True, # ) # # Update daily count (after document is created/updated) # _update_daily_count(collection, user_object_id, today) # logger.info("Media click logged for user %s", str(user_object_id)) # return True # except Exception as exc: # logger.error("Failed to log media click to admin MongoDB: %s", str(exc)) # return False # def close_connection(): # """Close MongoDB connection""" # global _client, _db, _admin_client, _admin_db # if _client: # _client.close() # _client = None # _db = None # logger.info("MongoDB connection closed") # if _admin_client: # _admin_client.close() # _admin_client = None # _admin_db = None # logger.info("Admin MongoDB connection closed")