database.py

'''
This file hosts all functionality related to data base operations using sqlalchemy.
- Initialization of the database file and engine
- save data in a new table
- close the database connection after use
'''

import sqlalchemy as db
from sqlalchemy.exc import SQLAlchemyError

class DatabaseHandler:
    '''
    Handles all database operations — creates and closes the connection and moves data into tables
    '''
    def __init__(self, db_path):
        '''
        Upon instantiation of this class, the engine is created
        '''
        try:
            self.engine = db.create_engine(f"sqlite:///{db_path}")
            with self.engine.connect():
                pass
        except SQLAlchemyError as e:
            raise SQLAlchemyError(f"Failed to connect to database: {e}")

    def save_data(self, df, table_name):
        '''
        Save a DataFrame into the specified table
        '''
        try:
            df.to_sql(table_name, self.engine, if_exists='replace', index=False)
        except SQLAlchemyError as e:
            raise SQLAlchemyError(f"Failed to save data to table '{table_name}': {e}")
    
    def close(self):
        '''
        This function decommissions the db engine after usage 
        '''
        try:
            self.engine.dispose()
        except SQLAlchemyError as e:
            print(f"Warning: error while closing database connection: {e}")