From 40afc9332db70a3a7e3a7874aac56755bfaad778 Mon Sep 17 00:00:00 2001
From: Matthew Truscott <102960299+m-truscott@users.noreply.github.com>
Date: Mon, 23 Mar 2026 16:13:48 +0000
Subject: [PATCH] feat: add album query and join methods to SqlQuery and sync
 with main

---
 src/llm_benchmark/sql/query.py | 80 ++++++++++++++++------------------
 1 file changed, 38 insertions(+), 42 deletions(-)

diff --git a/src/llm_benchmark/sql/query.py b/src/llm_benchmark/sql/query.py
index 53f6885..f2a42f9 100644
--- a/src/llm_benchmark/sql/query.py
+++ b/src/llm_benchmark/sql/query.py
@@ -1,4 +1,5 @@
 import sqlite3
+from contextlib import closing
 from textwrap import dedent
 
 
@@ -13,11 +14,11 @@ def query_album(name: str) -> bool:
         Returns:
             bool: True if the album exists, False otherwise
         """
-        conn = sqlite3.connect("data/chinook.db")
-        cur = conn.cursor()
+        with closing(sqlite3.connect("data/chinook.db")) as conn:
+            cur = conn.cursor()
 
-        cur.execute(f"SELECT * FROM Album WHERE Title = '{name}'")
-        return len(cur.fetchall()) > 0
+            cur.execute("SELECT 1 FROM Album WHERE Title = ?", (name,))
+            return cur.fetchone() is not None
 
     @staticmethod
     def join_albums() -> list:
@@ -26,30 +27,24 @@ def join_albums() -> list:
         Returns:
             list:
         """
-        conn = sqlite3.connect("data/chinook.db")
-        cur = conn.cursor()
+        with closing(sqlite3.connect("data/chinook.db")) as conn:
+            cur = conn.cursor()
 
-        cur.execute(
-            dedent(
-                """\
-                SELECT 
-                    t.Name AS TrackName, (
-                        SELECT a2.Title 
-                        FROM Album a2 
-                        WHERE a2.AlbumId = t.AlbumId
-                    ) AS AlbumName, 
-                    (
-                        SELECT ar.Name 
-                        FROM Artist ar
-                        JOIN Album a3 ON a3.ArtistId = ar.ArtistId
-                        WHERE a3.AlbumId = t.AlbumId
-                    ) AS ArtistName
-                FROM 
-                    Track t
-                """
+            cur.execute(
+                dedent(
+                    """\
+                    SELECT 
+                        t.Name AS TrackName, 
+                        a.Title AS AlbumName, 
+                        ar.Name AS ArtistName
+                    FROM 
+                        Track t
+                    LEFT JOIN Album a ON a.AlbumId = t.AlbumId
+                    LEFT JOIN Artist ar ON ar.ArtistId = a.ArtistId
+                    """
+                )
             )
-        )
-        return cur.fetchall()
+            return cur.fetchall()
 
     @staticmethod
     def top_invoices() -> list:
@@ -58,21 +53,22 @@ def top_invoices() -> list:
         Returns:
             list: List of tuples
         """
-        conn = sqlite3.connect("data/chinook.db")
-        cur = conn.cursor()
+        with closing(sqlite3.connect("data/chinook.db")) as conn:
+            cur = conn.cursor()
 
-        cur.execute(
-            dedent(
-                """\
-                SELECT 
-                    i.InvoiceId, 
-                    c.FirstName || ' ' || c.LastName AS CustomerName, 
-                    i.Total
-                FROM 
-                    Invoice i
-                JOIN Customer c ON c.CustomerId = i.CustomerId
-                ORDER BY i.Total DESC
-                """
+            cur.execute(
+                dedent(
+                    """\
+                    SELECT 
+                        i.InvoiceId, 
+                        c.FirstName || ' ' || c.LastName AS CustomerName, 
+                        i.Total
+                    FROM 
+                        Invoice i
+                    JOIN Customer c ON c.CustomerId = i.CustomerId
+                    ORDER BY i.Total DESC
+                    LIMIT 10
+                    """
+                )
             )
-        )
-        return cur.fetchall()[:10]
+            return cur.fetchall()
\ No newline at end of file