db_operation.py

'''
Abstraction for saving data to db
'''
import pandas as pd
import table_schema as ts
from sqlalchemy import create_engine

db_url = 'sqlite:///instance/local.db'

def _validate_schema(df, schema):
    '''
    validate df has the same columns and data types as schema

    Parameters
    ----------
    df: pd.DataFrame
    schema: dict
        {column_name: data_type}

    Returns
    -------
    bool
        True if df has the same columns and data types as schema
        False otherwise
    '''

    # check if the DataFrame has the same columns as the schema
    if set(df.columns) != set(schema.keys()):
        return False
    # check if the data types of the columns match the schema
    # TODO: ignoring type check for now
    # for col, dtype in schema.items():
    #     if df[col].dtype != dtype:
    #         return False
    return True

# def append_benchmark_profile(df):
#     '''append new entry to benchmark profile table'''
#     with create_engine(db_url).connect() as conn:
#         df.to_sql(ts.BENCHMARK_PROFILE_TABLE, con=conn, if_exists='append', index=False)


def get_most_recent_profile(type):
    table_name = 'benchmark_profile' if type == 'benchmark' else 'portfolio_profile'
    query = f"SELECT * FROM {table_name} WHERE date = (SELECT MAX(date) FROM {table_name})"
    with create_engine(db_url).connect() as conn:
        df = pd.read_sql(query, con=conn)
        # convert date to datetime object
        df['date'] = pd.to_datetime(df['date'])
        return df

def _get_most_recent(table_name, ts_column='date'):
    '''return the most recent entry in the table'''
    query = f"SELECT * FROM {table_name} WHERE {ts_column} = (SELECT MAX({ts_column}) FROM {table_name})"
    with create_engine(db_url).connect() as conn:
        df = pd.read_sql(query, con=conn)
        # convert date to datetime object
        df[ts_column] = pd.to_datetime(df[ts_column])
        return df

def get_most_recent_portfolio_profile():
    df = _get_most_recent(ts.PORTFOLIO_TABLE)
    df['date'] = pd.to_datetime(df['date'])
    return df

def get_most_recent_stocks_price():
    df = _get_most_recent(ts.STOCKS_PRICE_TABLE, ts_column='time')
    df['time'] = pd.to_datetime(df['time'])
    return df

def _append_df_to_db(df, table_name, schema):
    # validation 
    if not _validate_schema(df, schema):
        raise Exception(f'VALIDATION_ERROR: df does not have the same schema as the table {table_name}')
    with create_engine(db_url).connect() as conn:
        df.to_sql(table_name, con=conn, if_exists='append', index=False)

def append_to_stocks_price_table(df):
    '''append new entry to stocks price table'''
    _append_df_to_db(df, ts.STOCKS_PRICE_TABLE, ts.STOCKS_PRICE_TABLE_SCHEMA)

def get_all_stocks():
    with create_engine(db_url).connect() as conn:
        all_stocks = pd.read_sql(ts.STOCKS_DETAILS_TABLE, con=conn)
        return all_stocks