Distilled version of SQLAlchemy wrapped up in one class

distilldb.py

'''
Name: Distilled (since it's sqlalchemy with the parts a normal person cares about distilled from the rest in one Database class)

Very basic wrapper around the sqlalchemy orm that tries to replicate the ease of use that you get with r's dbplyr. Namely:
- provide the database connection details (in this case ones that are stored in a config file)
- return a single object from which you can do everything you need to

Similar in spirit to the more developed library: [dataset](https://dataset.readthedocs.io/en/latest/install.html)

Rewrote an old version of this with help from chatgpt

I actually think it's pretty good; and it gives you all the benefits of sqlalchemy. It doesnt try to reinvent the wheel; it just makes sqlalchemy more convenient
'''

# Python Standard Library
import urllib, pickle, os, json

# External Dependencies
from sqlalchemy import create_engine, MetaData, Table, Column, func
from sqlalchemy.orm import sessionmaker
import sqlalchemy.orm, sqlalchemy.schema
from sqlalchemy.types import Integer, String, Float

Types = {
    str: String,
    int: Integer,
    float: Float,
}

func = func # <-- it's just easier to work with

class Database:

    def __init__(self, load_meta = None, config = None, **kwargs):

        # Create the database URI
        if config is not None:
            with open(config, 'r') as file: config = json.load(file)
        else:
            config = {}
        if kwargs.get('dialect', config.get('dialect', '')) == 'sqlite':
            self.uri = f"{kwargs.get('dialect', config.get('dialect', ''))}:///{kwargs.get('database', config.get('database', ''))}"
        else:
            self.uri = f"{kwargs.get('dialect', config.get('dialect', ''))}://{kwargs.get('username', config.get('username', ''))}:{urllib.parse.quote(kwargs.get('password', config.get('password', '')))}@{kwargs.get('host', config.get('host', ''))}:{kwargs.get('port', config.get('port', ''))}/{kwargs.get('database', config.get('database', ''))}"

        # Create the SQLAlchemy engine and session
        self.engine = create_engine(self.uri)
        self.Session = sessionmaker(bind=self.engine)

        # Create the metadata object for table definitions
        if (load_meta is not None) and (os.path.exists(load_meta)):
            with open(load_meta, 'rb') as file:
                self.meta = pickle.load(file)
        else:
            self.meta = MetaData(); self.meta.reflect(self.engine)
            if load_meta is not None:
                self.save_meta(load_meta)

        # Create a class attribute for the ORM query function
        self.Session = sessionmaker(self.engine)
        # self.query = self.Session().query

    def __getitem__(self, table_name):
        # Get the table object from the metadata
        table = self.meta.tables.get(table_name)

        if table is None:
            raise KeyError(f"Table '{table_name}' does not exist.")

        return table

    def save_meta(self, path):
        '''For when you dont want to keep calling the db connection'''
        with open(path, 'wb') as file:
            pickle.dump(self.meta, file)

    def update_meta(self, path):
        self.meta = MetaData(); self.meta.reflect(self.engine)
        with open(path, 'wb') as file:
            pickle.dump(self.meta, file)

    def create_table(self, table_name, columns, drop_existing = False, primary_key = None, autoincrement_primary_key = True):
        table = self.meta.tables.get(table_name)

        if (table is not None) & (drop_existing == False):
            print(f"Table '{table_name}' already exist. Ignoring your request and continuing on like it didn't happen...")
        else:
            if (table is not None) & (drop_existing == True):
                table.drop(bind = self.engine)
                self.meta.remove(table)

            # Create a custom table class dynamically
            table = Table(
                table_name, 
                self.meta,
                *[
                    Column(
                        column_name, 
                        Types[column_type],
                        **({'primary_key': True, 'autoincrement': autoincrement_primary_key} if primary_key == column_name else {})
                    )
                    for column_name, column_type in columns.items()
                ]
            )

            # Create the table in the database
            self.meta.create_all(self.engine)


# Example usage
if __name__ == '__main__':

    # Connect to Database; if using sql, use `database` to specify the file location
    db = Database(
        dialect = 'sqlite',
        database = '.ignore/test.db',
    )

    # # Create Table
    if db.meta.tables.get('table_name') is None: # <-- you dont have to do this check if you dont want to. you can just let it throw an error to avoid overwritting an existing table
        db.create_table('table_name',{
            'col1': str,
            'col2': int,
            'col4': float,
        })
        # ^ eventually we'll probably have to complicate this with something like a dictionary of dictionaries (e.g., {'col1': {'type': str, 'primary_key': False, etc}, etc) 

    # Access table
    table = db['table_name']

    # # Add some data
    with db.Session() as session:
        row = table.insert().values([
            {'col1': 'hey', 'col2': 0, 'col4': .2},
            {'col1': 'yo', 'col2': 1, 'col4': .2},
            {'col1': 'sup', 'col2': 2, 'col4': .992},
        ])
        session.execute(row)
        session.commit()

    # # Query Table
    with db.Session() as session:
        print(
            len(session.query(table).all())
        )