genre counts in SQL

main.py

@@ -2,27 +2,47 @@ import os
 import pandas as pd
 from dotenv import load_dotenv, find_dotenv
 from sqlalchemy.engine import URL, create_engine
+from sqlalchemy.engine.base import Engine
 
-load_dotenv(find_dotenv())
 
-url = URL.create(
+def get_engine() -> Engine:
-    drivername="postgresql+psycopg2",
+    """Load env vars and return a SQLAlchemy engine connected to the database."""
-    username=os.getenv("DATABASE_USERNAME"),
+    load_dotenv(find_dotenv())
-    password=os.getenv("DATABASE_PASSWORD"),
-    host=os.getenv("DATABASE_IP"),
-    port=int(os.getenv("DATABASE_PORT")),
-    database=os.getenv("DATABASE_NAME"),
-)
 
-engine = create_engine(url)
+    url = URL.create(
+        drivername="postgresql+psycopg2",
+        username=os.getenv("DATABASE_USERNAME"),
+        password=os.getenv("DATABASE_PASSWORD"),
+        host=os.getenv("DATABASE_IP"),
+        port=int(os.getenv("DATABASE_PORT", 5432)),
+        database=os.getenv("DATABASE_NAME"),
+    )
 
-query = """SELECT title, genres FROM audiobookshelf.ao3;"""
+    return create_engine(url)
 
-data = pd.read_sql_query(query, engine)
+def run_query(engine: Engine, query: str) -> pd.DataFrame:
-print(data.head())
+    return pd.read_sql_query(query, engine)
 
-# cursor.execute(Query)                                          
+def fetch_titles_and_genres(engine: Engine) -> pd.DataFrame:
-# row = cursor.fetchone()                                                      
+    """Query title/genres from audiobookshelf.ao3 and return as a DataFrame."""
-# while row is not None:
+    query = """SELECT title, genres FROM audiobookshelf.ao3;"""
-#     print(row)
+    return run_query(engine, query)
-#     row = cursor.fetchone()                                                                               
+
+def count_genre_occurrences(engine: Engine) -> pd.DataFrame:
+    """Count how many times each genre appears across all titles in SQL before returning."""
+    query = """SELECT genre, COUNT(*) 
+                FROM audiobookshelf.ao3, unnest(genres) AS genre
+                GROUP BY genre
+                ORDER BY COUNT(*) DESC;"""
+    return run_query(engine, query)
+
+def main():
+    engine = get_engine()
+    #data = fetch_titles_and_genres(engine)
+    genre_counts = count_genre_occurrences(engine)
+    print(genre_counts)
+    # print(data.head())
+
+
+if __name__ == "__main__":
+    main()