Skip to content
Snippets Groups Projects
Forked from Michael Bücker / jupyterhub-ai
52 commits behind the upstream repository.
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
01_load_db.py 956 B
import os
from dotenv import load_dotenv
import sqlalchemy
import pandas as pd
from urllib.parse import quote_plus

load_dotenv("/usr/src/app/.env")

class Database:
    def __init__(self, db_name):
        self.conn = db_connect(db_name)

    def __getattr__(self, table_name):
        return pd.read_sql_table(table_name, self.conn)

    def list_tables(self):
        inspector = sqlalchemy.inspect(self.conn)
        table_names = inspector.get_table_names()
        return table_names

def db_connect(db_name):
    hostname=os.getenv("DB_HOST")
    user=os.getenv("DB_USER")
    password=quote_plus(os.getenv("DB_PASSWORD"))
    conn = sqlalchemy.create_engine(f'postgresql+psycopg2://{user}:{password}@{hostname}/{db_name}')
    return conn

def get_table(db_name, table_name):
    conn = db_connect(db_name)
    dat = pd.read_sql_table(table_name, conn)
    return dat

def get_all_tables(db_name):
    db_obj = Database(db_name)
    return db_obj