diff options
| -rw-r--r-- | api/app.py | 15 | ||||
| -rw-r--r-- | api/database.py | 14 |
2 files changed, 23 insertions, 6 deletions
@@ -1,4 +1,4 @@ -from flask import Flask, render_template, jsonify, request, redirect +from flask import Flask, jsonify, request, redirect from flask_cors import CORS from .database import PostgresHandler import os @@ -48,12 +48,19 @@ def generate_organization_captcha(org): return jsonify({"error": "Database Connection Failed. Dynamic Affiliation Endpoint requires a PostgreSQL Connection"}), 500 if server.check_row_exists("vtuber_data", "affiliation", org) is False: return jsonify({"error": "Organization " + org + " was not found in the database" }), 404 - num_correct = random.randint(5, 9) + num_correct = random.randint(4, 8) num_wrong = 16 - num_correct correct_answers= server.get_random_rows('vtuber_data', num_correct, "affiliation = '"+org+"'") - random_answers = server.get_random_rows('vtuber_data', num_wrong) + wrong_condition = "affiliation != '" + org + "'" + random_answers = server.get_random_rows('vtuber_data', num_wrong, wrong_condition) server.close_connection() - question_data = [{"image": question[3], "name": question[1], "affiliation": question[2], "id": question[0] } for question in correct_answers + random_answers] + if correct_answers is False or random_answers is False: + return jsonify({"error": "Failed to fetch VTuber data from database"}), 500 + + # Combine and scramble the questions + all_questions = correct_answers + random_answers + random.shuffle(all_questions) + question_data = [{"image": question[3], "name": question[1], "affiliation": question[2], "id": question[0] } for question in all_questions] if create_session: server = create_database_connection() session_id = secrets.token_urlsafe(16) 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 |