From 1f62a5e008a618f0557a38f63b5add35cc2c02f0 Mon Sep 17 00:00:00 2001 From: CC Date: Tue, 1 Sep 2026 22:51:06 +0100 Subject: [PATCH] functions --- main.py | 46 ++++++++++++++++++++++++++++------------------ 1 file changed, 28 insertions(+), 18 deletions(-) diff --git a/main.py b/main.py index 1140115..fd7fc0a 100644 --- a/main.py +++ b/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( - drivername="postgresql+psycopg2", - username=os.getenv("DATABASE_USERNAME"), - password=os.getenv("DATABASE_PASSWORD"), - host=os.getenv("DATABASE_IP"), - port=int(os.getenv("DATABASE_PORT")), - database=os.getenv("DATABASE_NAME"), -) +def get_engine() -> Engine: + """Load env vars and return a SQLAlchemy engine connected to the database.""" + load_dotenv(find_dotenv()) -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) -print(data.head()) +def run_query(engine: Engine, query: str) -> pd.DataFrame: + return pd.read_sql_query(query, engine) -# cursor.execute(Query) -# row = cursor.fetchone() -# while row is not None: -# print(row) -# row = cursor.fetchone() +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) + +def main(): + engine = get_engine() + data = fetch_titles_and_genres(engine) + print(data.head()) + + +if __name__ == "__main__": + main() \ No newline at end of file