Merge pull request #547 from Aditya30ag/fix/database-connection-leaks-464

Fix: Critical database connection leaks in all database functions

Author: rahulharpal1603

backend/app/database/albums.py

@@ -40,69 +40,83 @@ def get_db_connection():
 
 
 def db_create_albums_table() -> None:
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
-    cursor.execute(
-        """
-        CREATE TABLE IF NOT EXISTS albums (
-            album_id TEXT PRIMARY KEY,
-            album_name TEXT UNIQUE,
-            description TEXT,
-            is_hidden BOOLEAN DEFAULT 0,
-            password_hash TEXT
+    conn = None
+    try:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS albums (
+                album_id TEXT PRIMARY KEY,
+                album_name TEXT UNIQUE,
+                description TEXT,
+                is_hidden BOOLEAN DEFAULT 0,
+                password_hash TEXT
+            )
+            """
         )
-        """
-    )
-    conn.commit()
-    conn.close()
+        conn.commit()
+    finally:
+        if conn is not None:
+            conn.close()
 
 
 def db_create_album_images_table() -> None:
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
-    cursor.execute(
-        """
-        CREATE TABLE IF NOT EXISTS album_images (
-            album_id TEXT,
-            image_id TEXT,
-            PRIMARY KEY (album_id, image_id),
-            FOREIGN KEY (album_id) REFERENCES albums(album_id) ON DELETE CASCADE,
-            FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE
+    conn = None
+    try:
+        conn = sqlite3.connect(DATABASE_PATH)
+        cursor = conn.cursor()
+        cursor.execute(
+            """
+            CREATE TABLE IF NOT EXISTS album_images (
+                album_id TEXT,
+                image_id TEXT,
+                PRIMARY KEY (album_id, image_id),
+                FOREIGN KEY (album_id) REFERENCES albums(album_id) ON DELETE CASCADE,
+                FOREIGN KEY (image_id) REFERENCES images(id) ON DELETE CASCADE
+            )
+            """
         )
-        """
-    )
-    conn.commit()
-    conn.close()
+        conn.commit()
+    finally:
+        if conn is not None:
+            conn.close()
 
 
 def db_get_all_albums(show_hidden: bool = False):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    if show_hidden:
-        cursor.execute("SELECT * FROM albums")
-    else:
-        cursor.execute("SELECT * FROM albums WHERE is_hidden = 0")
-    albums = cursor.fetchall()
-    conn.close()
-    return albums
+    try:
+        if show_hidden:
+            cursor.execute("SELECT * FROM albums")
+        else:
+            cursor.execute("SELECT * FROM albums WHERE is_hidden = 0")
+        albums = cursor.fetchall()
+        return albums
+    finally:
+        conn.close()
 
 
 def db_get_album_by_name(name: str):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("SELECT * FROM albums WHERE album_name = ?", (name,))
-    album = cursor.fetchone()
-    conn.close()
-    return album if album else None
+    try:
+        cursor.execute("SELECT * FROM albums WHERE album_name = ?", (name,))
+        album = cursor.fetchone()
+        return album if album else None
+    finally:
+        conn.close()
 
 
 def db_get_album(album_id: str):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("SELECT * FROM albums WHERE album_id = ?", (album_id,))
-    album = cursor.fetchone()
-    conn.close()
-    return album if album else None
+    try:
+        cursor.execute("SELECT * FROM albums WHERE album_id = ?", (album_id,))
+        album = cursor.fetchone()
+        return album if album else None
+    finally:
+        conn.close()
 
 
 def db_insert_album(
@@ -114,20 +128,22 @@ def db_insert_album(
 ):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    password_hash = None
-    if password:
-        password_hash = bcrypt.hashpw(
-            password.encode("utf-8"), bcrypt.gensalt()
-        ).decode("utf-8")
-    cursor.execute(
-        """
-        INSERT INTO albums (album_id, album_name, description, is_hidden, password_hash)
-        VALUES (?, ?, ?, ?, ?)
-        """,
-        (album_id, album_name, description, int(is_hidden), password_hash),
-    )
-    conn.commit()
-    conn.close()
+    try:
+        password_hash = None
+        if password:
+            password_hash = bcrypt.hashpw(
+                password.encode("utf-8"), bcrypt.gensalt()
+            ).decode("utf-8")
+        cursor.execute(
+            """
+            INSERT INTO albums (album_id, album_name, description, is_hidden, password_hash)
+            VALUES (?, ?, ?, ?, ?)
+            """,
+            (album_id, album_name, description, int(is_hidden), password_hash),
+        )
+        conn.commit()
+    finally:
+        conn.close()
 
 
 def db_update_album(
@@ -139,38 +155,52 @@ def db_update_album(
 ):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    password_hash = None
-    if password:
-        password_hash = bcrypt.hashpw(
-            password.encode("utf-8"), bcrypt.gensalt()
-        ).decode("utf-8")
-    cursor.execute(
-        """
-        UPDATE albums
-        SET album_name = ?, description = ?, is_hidden = ?, password_hash = ?
-        WHERE album_id = ?
-        """,
-        (album_name, description, int(is_hidden), password_hash, album_id),
-    )
-    conn.commit()
-    conn.close()
+    try:
+        if password is not None:
+            # Update with new password
+            password_hash = bcrypt.hashpw(
+                password.encode("utf-8"), bcrypt.gensalt()
+            ).decode("utf-8")
+            cursor.execute(
+                """
+                UPDATE albums
+                SET album_name = ?, description = ?, is_hidden = ?, password_hash = ?
+                WHERE album_id = ?
+                """,
+                (album_name, description, int(is_hidden), password_hash, album_id),
+            )
+        else:
+            # Update without changing password
+            cursor.execute(
+                """
+                UPDATE albums
+                SET album_name = ?, description = ?, is_hidden = ?
+                WHERE album_id = ?
+                """,
+                (album_name, description, int(is_hidden), album_id),
+            )
+        conn.commit()
+    finally:
+        conn.close()
 
 
 def db_delete_album(album_id: str):
-    conn = sqlite3.connect(DATABASE_PATH)
-    cursor = conn.cursor()
-    cursor.execute("DELETE FROM albums WHERE album_id = ?", (album_id,))
-    conn.commit()
-    conn.close()
+    with get_db_connection() as conn:
+        cursor = conn.cursor()
+        cursor.execute("DELETE FROM albums WHERE album_id = ?", (album_id,))
 
 
 def db_get_album_images(album_id: str):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("SELECT image_id FROM album_images WHERE album_id = ?", (album_id,))
-    images = cursor.fetchall()
-    conn.close()
-    return [img[0] for img in images]
+    try:
+        cursor.execute(
+            "SELECT image_id FROM album_images WHERE album_id = ?", (album_id,)
+        )
+        images = cursor.fetchall()
+        return [img[0] for img in images]
+    finally:
+        conn.close()
 
 
 def db_add_images_to_album(album_id: str, image_ids: list[str]):
@@ -214,20 +244,26 @@ def db_remove_image_from_album(album_id: str, image_id: str):
 def db_remove_images_from_album(album_id: str, image_ids: list[str]):
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.executemany(
-        "DELETE FROM album_images WHERE album_id = ? AND image_id = ?",
-        [(album_id, img_id) for img_id in image_ids],
-    )
-    conn.commit()
-    conn.close()
+    try:
+        cursor.executemany(
+            "DELETE FROM album_images WHERE album_id = ? AND image_id = ?",
+            [(album_id, img_id) for img_id in image_ids],
+        )
+        conn.commit()
+    finally:
+        conn.close()
 
 
 def verify_album_password(album_id: str, password: str) -> bool:
     conn = sqlite3.connect(DATABASE_PATH)
     cursor = conn.cursor()
-    cursor.execute("SELECT password_hash FROM albums WHERE album_id = ?", (album_id,))
-    row = cursor.fetchone()
-    conn.close()
-    if not row or not row[0]:
-        return False
-    return bcrypt.checkpw(password.encode("utf-8"), row[0].encode("utf-8"))
+    try:
+        cursor.execute(
+            "SELECT password_hash FROM albums WHERE album_id = ?", (album_id,)
+        )
+        row = cursor.fetchone()
+        if not row or not row[0]:
+            return False
+        return bcrypt.checkpw(password.encode("utf-8"), row[0].encode("utf-8"))
+    finally:
+        conn.close()