src/data_models/sql_connection.py

""""Create a connection to the database and return the engine and session objects."""

import os

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

HOST = os.getenv("DB_HOST")
PORT = os.getenv("DB_PORT")
DBNAME = os.getenv("DB_DATABASE")
USERNAME = os.getenv("DB_USERNAME")
PASSWORD = os.getenv("DB_PASSWORD")


def get_db_connection() -> tuple:
    """Create a connection to the database and return the engine and session objects.

    Returns:
        engine: connection to the database
        session: session to interact with the database
    """
    database_url = f"postgresql+psycopg2://{USERNAME}:{PASSWORD}@{HOST}:{PORT}/{DBNAME}"

    engine = create_engine(database_url)
    Session = sessionmaker(bind=engine)
    session = Session()

    return engine, session