populate_sqlite.py

import sqlite3
import pandas as pd
import numpy as np

from time import sleep
from datetime import datetime


def connect_to_sqlite(db_name):
    """
    creates a connect to a local sqlite database
    :param db_name: the database to connect to
    :return: sqlite connection object
    """
    return sqlite3.connect(db_name)


def generate_data(n_rows=100, n_cols=10, p_int=0.5, p_str=0.5):
    """
    generates a dataframe of faux data for testing
    :param n_rows: number of rows in the dataframe
    :param n_cols: number of columns in the dataframe
    :param p_int: percentage of columns that are integers
    :param p_str: percentage of columns that are strings
    """
    column_names = ['col_' + str(n) for n in range(n_cols)]

    float_data = []
    for i in range(int(p_int * n_cols)):
        temp_float_data = np.random.choice(100, n_rows).tolist()
        float_data.append(temp_float_data)

    str_data = []
    for i in range(int(p_str * n_cols)):
        temp_str_data = np.random.choice(['a', 'b', 'c', 'd'], n_rows).tolist()
        str_data.append(temp_str_data)

    df = pd.DataFrame()
    for column, data in zip(column_names, float_data + str_data):
        temp_df = pd.DataFrame({column: data})
        df = pd.concat([df, temp_df], axis=1)

    df['insert_timestamp'] = datetime.now().strftime("%m-%d-%Y %H:%M:%S")
    return df


def dynamically_create_sqlite_table_from_dataframe(df, date_fields, table_name, sqlite_conn):
    """
    generates and executes a create table SQL statement based on a dataframe, using duck-typing to infer column datatypes
    :param df: pandas dataframe to write to sqlite
    :param date_fields: fields in the dataframe we need to ensure are read as dates and not strings
    :param table_name: sqlite table now
    :param sqlite_conn: connection to sqlite database
    """
    create_table_statement = f'create table if not exists {table_name} (id integer not null primary key autoincrement,'

    dtype_mapping = {
        'float': 'float',
        'int': 'integer',
        'date': 'timestamp',
        'object': 'text'
    }

    for field in date_fields:
        df[field] = pd.to_datetime(df[field])

    df_dtypes = df.dtypes
    for index, dtype in enumerate(df_dtypes):
        column_name = df_dtypes.index[index]
        mapped_dtype = [val for key, val in dtype_mapping.items() if key in str(dtype)][0]
        create_statement_addendum = f''' {column_name} {mapped_dtype},'''
        create_table_statement += create_statement_addendum

    create_table_statement = create_table_statement[:-1]
    create_table_statement += ')'

    cursor = sqlite_conn.cursor()
    cursor.execute(create_table_statement)


def write_df_to_sqlite_table(df, table_name, sqlite_conn):
    """
    write a pandas dataframe to a sqlite table
    :param df: pandas dataframe to write to sqlite table
    :param table_name: name of the sqlite table
    :param sqlite_conn: connection to sqlite database
    """
    df.to_sql(table_name, sqlite_conn, if_exists='append', index=False)


if __name__ == "__main__":
    sample_df = generate_data()
    sqlite_connection = connect_to_sqlite('sample_db')
    dynamically_create_sqlite_table_from_dataframe(sample_df, ['insert_timestamp'], 'sample_table', sqlite_connection)
    write_df_to_sqlite_table(sample_df, 'sample_table', sqlite_connection)
    sleep(10)
    new_sample_df = generate_data()
    write_df_to_sqlite_table(new_sample_df, 'sample_table', sqlite_connection)