PATCH: Closed database connections in util to avoid data leaks

web_server/utils/recommendation_utils.py

@@ -10,7 +10,8 @@ def user_recommendation_category(user_id: int) -> Optional[int]:
 
     data = db.fetchone(
         "SELECT category_id FROM user_preferences WHERE user_id = ? ORDER BY favourability DESC LIMIT 1", (user_id,))
-    return data[0]
+    db.close_connection()
+    return data
 
 def followed_categories_recommendations(user_id: int):
     """
@@ -25,6 +26,7 @@ def followed_categories_recommendations(user_id: int):
                              WHERE category_id IN (SELECT category_id FROM categories WHERE user_id = ?)
                              ORDER BY num_viewers DESC
                              LIMIT 25;  """, (user_id,))
+    db.close_connection()
     return categories
 
 def recommendations_based_on_category(category_id: int) -> Optional[List[Tuple[int, str, int]]]:
@@ -43,6 +45,7 @@ def recommendations_based_on_category(category_id: int) -> Optional[List[Tuple[i
         WHERE categories.category_id = ? 
         ORDER BY num_viewers DESC 
         LIMIT 25""", (category_id,))
+    db.close_connection()
     return data
 
 def default_recommendations():
@@ -60,5 +63,6 @@ def default_recommendations():
         ORDER BY num_viewers DESC 
         LIMIT 25
     """)
+    db.close_connection()
     return data
 

web_server/utils/stream_utils.py

@@ -9,7 +9,9 @@ def streamer_live_status(user_id: int) -> bool:
     """
     db = Database()
     db.create_connection()
-    return bool(db.fetchone("SELECT 1 FROM streams WHERE user_id = ? AND isLive = 1 ORDER BY stream_id DESC", (user_id,)))
+    is_live = bool(db.fetchone("SELECT 1 FROM streams WHERE user_id = ? AND isLive = 1 ORDER BY stream_id DESC", (user_id,)))
+    db.close_connection()
+    return is_live
 
 def followed_live_streams(user_id: int) -> list[dict]:
     """
@@ -25,6 +27,7 @@ def followed_live_streams(user_id: int) -> list[dict]:
         AND stream_id = (SELECT MAX(stream_id) FROM streams WHERE user_id = streams.user_id)
         AND isLive = 1;
         """, (user_id,))
+    db.close_connection()
 
     return live_streams
 
@@ -41,6 +44,7 @@ def followed_streamers(user_id: int) -> list[dict]:
         WHERE user_id IN (SELECT followed_id FROM follows WHERE user_id = ?);
         """, (user_id,))
 
+    db.close_connection()
     return followed_streamers
 
 def streamer_most_recent_stream(user_id: int) -> dict:
@@ -53,6 +57,7 @@ def streamer_most_recent_stream(user_id: int) -> dict:
                                         user_id = ? AND 
                                         stream_id = (SELECT MAX(stream_id) FROM
                                         streams WHERE user_id = ?)""", (user_id, user_id))
+    db.close_connection()
     return most_recent_stream
 
 def user_stream(user_id: int, stream_id: int) -> dict:
@@ -62,5 +67,6 @@ def user_stream(user_id: int, stream_id: int) -> dict:
     db = Database()
     db.create_connection()
     stream = db.fetchone("SELECT * FROM streams WHERE user_id = ? AND stream_id = ?", (user_id,stream_id))
+    db.close_connection()
 
     return stream

web_server/utils/user_utils.py

@@ -25,6 +25,8 @@ def get_user_id(username: str) -> Optional[int]:
     except Exception as e:
         print(f"Error: {e}")
         return None
+    finally:
+        db.close_connection()
 
 def get_username(user_id: str) -> Optional[str]:
     """
@@ -42,6 +44,8 @@ def get_username(user_id: str) -> Optional[str]:
     except Exception as e:
         print(f"Error: {e}")
         return None
+    finally:
+        db.close_connection()
     
 def is_user_partner(user_id: int) -> bool:
     """
@@ -59,6 +63,8 @@ def is_user_partner(user_id: int) -> bool:
     except Exception as e:
         print(f"Error: {e}")
         return False
+    finally:
+        db.close_connection()
 
 def is_subscribed(user_id: int, streamer_id: int) -> bool:
     """
@@ -76,6 +82,8 @@ def is_subscribed(user_id: int, streamer_id: int) -> bool:
     except Exception as e:
         print(f"Error: {e}")
         return False
+    finally:
+        db.close_connection()
 
 def is_following(user_id: int, followed_id: int) -> bool:
     db = Database()
@@ -90,6 +98,8 @@ def is_following(user_id: int, followed_id: int) -> bool:
     except Exception as e:
         print(f"Error: {e}")
         return False
+    finally:
+        db.close_connection()
 
 def subscription_expiration(user_id: int, subscribed_id: int) -> int:
     """
@@ -102,11 +112,13 @@ def subscription_expiration(user_id: int, subscribed_id: int) -> int:
         data = db.fetchone(
             "SELECT expires from subscriptions WHERE user_id = ? AND subscribed_id = ? AND expires > since", (user_id,subscribed_id))
         if data:
-            expiration_date = data[0]
+            expiration_date = data["expires"]
 
             remaining_time = (expiration_date - datetime.now()).seconds
     except Exception as e:
         print(f"Error: {e}")
+    finally:
+        db.close_connection()
 
     return remaining_time
 
@@ -134,3 +146,5 @@ def reset_password(new_password: str, email: str):
     except Exception as e:
         print(f"Error: {e}")
         return False
+    finally:
+        db.close_connection()