Fix: Ensure atomic DB operations and consistency in global reclustering (#562)

Author: Hemil36

backend/app/database/albums.py

@@ -1,42 +1,7 @@
 import sqlite3
 import bcrypt
-import time
-from contextlib import contextmanager
 from app.config.settings import DATABASE_PATH
-
-
-@contextmanager
-def get_db_connection():
-    """Context manager for database connections with proper error handling and retries"""
-    max_retries = 3
-    retry_delay = 0.1
-
-    for attempt in range(max_retries):
-        try:
-            conn = sqlite3.connect(DATABASE_PATH, timeout=30.0)
-            # Enable WAL mode for better concurrency
-            conn.execute("PRAGMA journal_mode=WAL")
-            # Set busy timeout
-            conn.execute("PRAGMA busy_timeout=30000")
-            # Enable foreign keys
-            conn.execute("PRAGMA foreign_keys=ON")
-
-            try:
-                yield conn
-                conn.commit()
-            except Exception as e:
-                conn.rollback()
-                raise e
-            finally:
-                conn.close()
-            break
-
-        except sqlite3.OperationalError as e:
-            if "database is locked" in str(e).lower() and attempt < max_retries - 1:
-                time.sleep(retry_delay * (2**attempt))  # Exponential backoff
-                continue
-            else:
-                raise e
+from app.database.connection import get_db_connection
 
 
 def db_create_albums_table() -> None:

backend/app/database/connection.py

@@ -0,0 +1,32 @@
+import sqlite3
+from contextlib import contextmanager
+from typing import Generator
+from app.config.settings import DATABASE_PATH
+
+
+@contextmanager
+def get_db_connection() -> Generator[sqlite3.Connection, None, None]:
+    """
+    SQLite connection context manager with all integrity constraints enforced.
+
+    - Enables all major relational integrity PRAGMAs
+    - Works for both single and multi-step transactions
+    - Automatically commits on success or rolls back on failure
+    """
+    conn = sqlite3.connect(DATABASE_PATH)
+
+    # --- Strict enforcement of all relational and logical rules ---
+    conn.execute("PRAGMA foreign_keys = ON;")  # Enforce FK constraints
+    conn.execute("PRAGMA ignore_check_constraints = OFF;")  # Enforce CHECK constraints
+    conn.execute("PRAGMA recursive_triggers = ON;")  # Allow nested triggers
+    conn.execute("PRAGMA defer_foreign_keys = OFF;")  # Immediate FK checking
+    conn.execute("PRAGMA case_sensitive_like = ON;")  # Make LIKE case-sensitive
+
+    try:
+        yield conn
+        conn.commit()
+    except Exception:
+        conn.rollback()
+        raise
+    finally:
+        conn.close()

backend/app/database/face_clusters.py

@@ -39,22 +39,57 @@ def db_create_clusters_table() -> None:
             conn.close()
 
 
-def db_insert_clusters_batch(clusters: List[ClusterData]) -> List[ClusterId]:
+def db_delete_all_clusters(cursor: Optional[sqlite3.Cursor] = None) -> int:
+    """
+    Delete all clusters from the database.
+
+    Args:
+        cursor: Optional existing database cursor. If None, creates a new connection.
+
+    Returns:
+        Number of deleted clusters
+    """
+    own_connection = cursor is None
+    if own_connection:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
+
+    try:
+        cursor.execute("DELETE FROM face_clusters")
+        deleted_count = cursor.rowcount
+        if own_connection:
+            conn.commit()
+        return deleted_count
+    except Exception:
+        if own_connection:
+            conn.rollback()
+        print("Error deleting all clusters.")
+        raise
+    finally:
+        if own_connection:
+            conn.close()
+
+
+def db_insert_clusters_batch(
+    clusters: List[ClusterData], cursor: Optional[sqlite3.Cursor] = None
+) -> List[ClusterId]:
     """
     Insert multiple clusters into the database in batch.
 
     Args:
         clusters: List of ClusterData objects containing cluster information.
-        cluster_id: should be provided as UUID strings.
+        cursor: Optional existing database cursor. If None, creates a new connection.
 
     Returns:
         List of cluster IDs of the newly created clusters
     """
     if not clusters:
         return []
 
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
+    own_connection = cursor is None
+    if own_connection:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
 
     try:
         cluster_ids = []
@@ -72,14 +107,20 @@ def db_insert_clusters_batch(clusters: List[ClusterData]) -> List[ClusterId]:
             """
             INSERT INTO face_clusters (cluster_id, cluster_name, face_image_base64)
             VALUES (?, ?, ?)
-        """,
+            """,
             insert_data,
         )
 
-        conn.commit()
+        if own_connection:
+            conn.commit()
         return cluster_ids
+    except Exception:
+        if own_connection:
+            conn.rollback()
+        raise
     finally:
-        conn.close()
+        if own_connection:
+            conn.close()
 
 
 def db_get_cluster_by_id(cluster_id: ClusterId) -> Optional[ClusterData]:
@@ -145,18 +186,24 @@ def db_get_all_clusters() -> List[ClusterData]:
 def db_update_cluster(
     cluster_id: ClusterId,
     cluster_name: Optional[ClusterName] = None,
+    conn: Optional[sqlite3.Connection] = None,
 ) -> bool:
     """
     Update an existing cluster.
 
     Args:
         cluster_id: The ID of the cluster to update
         cluster_name: New cluster name (optional)
+        conn: Optional existing database connection. If None, creates a new connection.
 
     Returns:
         True if the cluster was updated, False if not found
     """
-    conn = sqlite3.connect(DATABASE_PATH)
+    # Use provided connection or create a new one
+    own_connection = conn is None
+    if own_connection:
+        conn = sqlite3.connect(DATABASE_PATH)
+
     cursor = conn.cursor()
 
     try:
@@ -185,26 +232,6 @@ def db_update_cluster(
         conn.close()
 
 
-def db_delete_all_clusters() -> int:
-    """
-    Delete all entries from the face_clusters table.
-
-    Returns:
-        Number of clusters deleted
-    """
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
-
-    try:
-        cursor.execute("DELETE FROM face_clusters")
-
-        deleted_count = cursor.rowcount
-        conn.commit()
-        return deleted_count
-    finally:
-        conn.close()
-
-
 def db_get_all_clusters_with_face_counts() -> (
     List[Dict[str, Union[str, Optional[str], int]]]
 ):

backend/app/database/faces.py

@@ -287,13 +287,15 @@ def db_get_all_faces_with_cluster_names() -> (
 
 def db_update_face_cluster_ids_batch(
     face_cluster_mapping: List[Dict[str, Union[FaceId, ClusterId]]],
+    cursor: Optional[sqlite3.Cursor] = None,
 ) -> None:
     """
     Update cluster IDs for multiple faces in batch.
 
     Args:
         face_cluster_mapping: List of dictionaries containing face_id and cluster_id pairs
                              Each dict should have keys: 'face_id' and 'cluster_id'
+        cursor: Optional existing database cursor. If None, creates a new connection.
 
     Example:
         face_cluster_mapping = [
@@ -305,8 +307,10 @@ def db_update_face_cluster_ids_batch(
     if not face_cluster_mapping:
         return
 
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
+    own_connection = cursor is None
+    if own_connection:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
 
     try:
         # Prepare update data as tuples (cluster_id, face_id)
@@ -325,9 +329,16 @@ def db_update_face_cluster_ids_batch(
             update_data,
         )
 
-        conn.commit()
+        if own_connection:
+            conn.commit()
+    except Exception:
+        if own_connection:
+            conn.rollback()
+        print("Error updating face cluster IDs in batch.")
+        raise
     finally:
-        conn.close()
+        if own_connection:
+            conn.close()
 
 
 def db_get_cluster_mean_embeddings() -> List[Dict[str, Union[str, FaceEmbedding]]]:

backend/app/database/metadata.py

@@ -55,18 +55,23 @@ def db_get_metadata() -> Optional[Dict[str, Any]]:
         conn.close()
 
 
-def db_update_metadata(metadata: Dict[str, Any]) -> bool:
+def db_update_metadata(
+    metadata: Dict[str, Any], cursor: Optional[sqlite3.Cursor] = None
+) -> bool:
     """
     Update the metadata in the database.
 
     Args:
         metadata: Dictionary containing metadata to store
+        cursor: Optional existing database cursor. If None, creates a new connection.
 
     Returns:
         True if the metadata was updated, False otherwise
     """
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
+    own_connection = cursor is None
+    if own_connection:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
 
     try:
         metadata_json = json.dumps(metadata)
@@ -75,8 +80,16 @@ def db_update_metadata(metadata: Dict[str, Any]) -> bool:
         cursor.execute("DELETE FROM metadata")
         cursor.execute("INSERT INTO metadata (metadata) VALUES (?)", (metadata_json,))
 
-        updated = cursor.rowcount > 0
-        conn.commit()
-        return updated
+        success = cursor.rowcount > 0
+        if own_connection:
+            conn.commit()
+        return success
+    except Exception as e:
+        if own_connection:
+            conn.rollback()
+
+        print(f"Error updating metadata: {e}")
+        raise
     finally:
-        conn.close()
+        if own_connection:
+            conn.close()