perf: optimize sql queries and database connections

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()