From 4427d35c020de915ebcb1f68a80808c9edbb4de5 Mon Sep 17 00:00:00 2001 From: TurinTech Bot Date: Mon, 7 Oct 2024 16:51:04 +0000 Subject: [PATCH] Artemis Changes --- src/llm_benchmark/sql/query.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) 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