refactor random rows function to take in conditionHEADmaster

1 files changed, 12 insertions, 2 deletions

diff --git a/api/database.py b/api/database.py
index a5d70e6..1c0715c 100644
--- a/api/database.py
+++ b/api/database.py
@@ -84,12 +84,22 @@ class PostgresHandler:
             print(e)
             return False
 
-    def get_random_rows(self, table_name: str, count: int, condition: str = None):
+    def get_random_rows(self, table_name: str, count: int, condition: str = None, unique_only: bool = False, exclude_ids: list = None):
         if condition is None:
             condition = "1 = 1"
+
+        # Add exclusion condition if exclude_ids is provided
+        if exclude_ids and len(exclude_ids) > 0:
+            id_list = ','.join(str(id) for id in exclude_ids)
+            condition = f"{condition} AND id NOT IN ({id_list})"
+
         try:
             cursor = self._connection.cursor()
-            query = f"SELECT * FROM {table_name} WHERE {condition} ORDER BY RANDOM() LIMIT {str(count)}"
+            if unique_only:
+                # Use subquery to handle DISTINCT with ORDER BY RANDOM()
+                query = f"SELECT * FROM (SELECT DISTINCT * FROM {table_name} WHERE {condition}) AS distinct_rows ORDER BY RANDOM() LIMIT {str(count)}"
+            else:
+                query = f"SELECT * FROM {table_name} WHERE {condition} ORDER BY RANDOM() LIMIT {str(count)}"
             cursor.execute(query)
             result = cursor.fetchall()
             return result