diff --git a/src/llm_benchmark/sql/query.py b/src/llm_benchmark/sql/query.py index 53f6885..413b52f 100644 --- a/src/llm_benchmark/sql/query.py +++ b/src/llm_benchmark/sql/query.py @@ -16,8 +16,10 @@ def query_album(name: str) -> bool: conn = sqlite3.connect("data/chinook.db") cur = conn.cursor() - cur.execute(f"SELECT * FROM Album WHERE Title = '{name}'") - return len(cur.fetchall()) > 0 + cur.execute("SELECT EXISTS(SELECT 1 FROM Album WHERE Title = ?)", (name,)) + exists = cur.fetchone()[0] + conn.close() + return exists @staticmethod def join_albums() -> list: @@ -75,4 +77,4 @@ def top_invoices() -> list: """ ) ) - return cur.fetchall()[:10] + return cur.fetchall()[:10] \ No newline at end of file