functions

main.py

@@ -2,27 +2,37 @@ 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:
+    """Load env vars and return a SQLAlchemy engine connected to the database."""
+    load_dotenv(find_dotenv())
+
+    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")),
+        port=int(os.getenv("DATABASE_PORT", 5432)),
         database=os.getenv("DATABASE_NAME"),
-)
+    )
 
-engine = create_engine(url)
+    return create_engine(url)
 
-query = """SELECT title, genres FROM audiobookshelf.ao3;"""
+def run_query(engine: Engine, query: str) -> pd.DataFrame:
+    return pd.read_sql_query(query, engine)
 
-data = pd.read_sql_query(query, engine)
-print(data.head())
+def fetch_titles_and_genres(engine: Engine) -> pd.DataFrame:
+    """Query title/genres from audiobookshelf.ao3 and return as a DataFrame."""
+    query = """SELECT title, genres FROM audiobookshelf.ao3;"""
+    return run_query(engine, query)
 
-# cursor.execute(Query)                                          
-# row = cursor.fetchone()                                                      
-# while row is not None:
-#     print(row)
-#     row = cursor.fetchone()                                                                               
+def main():
+    engine = get_engine()
+    data = fetch_titles_and_genres(engine)
+    print(data.head())
+
+
+if __name__ == "__main__":
+    main()