functions

main.py

@@ -2,7 +2,11 @@ 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
 
+
+def get_engine() -> Engine:
+    """Load env vars and return a SQLAlchemy engine connected to the database."""
     load_dotenv(find_dotenv())
 
     url = URL.create(
@@ -10,19 +14,25 @@ url = URL.create(
         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)
 
+def run_query(engine: Engine, query: str) -> pd.DataFrame:
+    return pd.read_sql_query(query, engine)
+
+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)
 
-data = pd.read_sql_query(query, engine)
+def main():
+    engine = get_engine()
+    data = fetch_titles_and_genres(engine)
     print(data.head())
 
-# cursor.execute(Query)                                          
-# row = cursor.fetchone()                                                      
-# while row is not None:
-#     print(row)
-#     row = cursor.fetchone()                                                                               
+
+if __name__ == "__main__":
+    main()