finance-watcher / scripts / seed_database.py
seed_database.py
Raw
from finance_watcher_dataclasses.enums import (
    AccountTypeEnum as AccountTypeEnum, BusinessTypeEnum
)
from finance_watcher_dataclasses.import_dataclasses import AccountRow
from database.database_classes import Transaction, Business, User
from database.finance_database import FinanceDatabase
import random
import string
import datetime
from typing import Optional, List

DEFAULT_COMPANY_LENGTH = 10
DEFAULT_ACCOUNT_LENGTH = 20


def create_filler_strings(length: int = 128)->str:
    return ''.join(random.choices(string.ascii_lowercase + string.digits+ string.ascii_uppercase, k=length))

def create_random_number(start:int = 0, limit: int= 1000) -> int:
    return random.randint(start, limit)

def create_random_date(weeks_ago:int = 24) -> datetime:
    start_weeks_ago = datetime.datetime.now() - datetime.timedelta(weeks=weeks_ago)
    return start_weeks_ago + datetime.timedelta(
        seconds=random.randint(0, int((datetime.datetime.now()-start_weeks_ago).total_seconds()))
    )

def insert_default_types(db: FinanceDatabase):
    db.insert_default_business_types()
    db.insert_account_types()


def create_dashboard_user(db: FinanceDatabase) -> Optional[User]:
    return db.insert_dashboard_user(
        create_filler_strings(10), f'{create_filler_strings(10)}@{create_filler_strings(8)}.com', create_filler_strings(30)
    )

def create_businesses(
    db: FinanceDatabase, 
    num_businesses_per_type: int = 5
) -> Optional[List[Business]]:
    names_business_type_id = {}
    for type in BusinessTypeEnum:
        business_type = db.get_business_type(type)
        for _ in range(0, num_businesses_per_type):
            names_business_type_id[create_filler_strings(DEFAULT_COMPANY_LENGTH)] = business_type.id
        
    return db.insert_businesses(names_business_type_id)

def create_accounts_and_totals(
    db: FinanceDatabase,
    user_id: int,
    totals_per_account: int = 100
) -> Optional[bool]:
    initial_date =  datetime.datetime.now() - datetime.timedelta(weeks=200)
    account_total_rows: List[AccountRow] = []
    for enum in AccountTypeEnum:
        account_type = db.get_account_type(enum)
        if account_type is None:
            return None
        account = db.insert_account(
            name=create_filler_strings(DEFAULT_ACCOUNT_LENGTH),
            account_type_id=account_type.id,
            dashboard_user_id=user_id
        )
        if account is None:
            return None
        date_to_amount = {}
        date = initial_date
        amount = create_random_number(0, 1000)
        for _ in range(0,totals_per_account):
            date = date + datetime.timedelta(days=1)   
            if account_type.is_expense:
                if random.randint(0,1) == 1:
                    amount = (amount + create_random_number(0,2000)) * -1
                else:
                    amount = (amount - create_random_number(0,2000)) * -1
                    if amount < 0:
                        amount = 0
                
                date_to_amount[date] = amount
            else:
                if random.randint(0,1) == 1:
                    amount = amount + create_random_number(0,3000)
                else:
                    amount = amount - create_random_number(1000,2500)
                    if amount < 0:
                        amount = 0
                date_to_amount[date] = amount
        
        account_total_rows.append(
            AccountRow(
                account_id=account.id,
                date_to_amount=date_to_amount
            )
        )
    return db.insert_account_totals(
        account_rows=account_total_rows
    )

def create_transactions(
    db: FinanceDatabase, 
    business_ids: List[int], 
    user_id: int,
    num_transactions: int = 1000
) -> None:
    transactions = [
        Transaction(
            id=0,
            dashboard_user_id=user_id, 
            amount=create_random_number(),
            business_id=business_ids[create_random_number(0, len(business_ids)-1)],
            is_expense=True,
            created_date=create_random_date()
        ) for _ in range(0,num_transactions)
    ]
    db.insert_transactions(transactions)


def seed_db():
    finance_database = FinanceDatabase()
    user = create_dashboard_user(finance_database)
    if user is None:
        print("Failed ot create a user. Seeding the database has been stopped.")
        return 1
    insert_default_types(finance_database)
    businesses = create_businesses(finance_database)
    if businesses is None:
        print(
            "Failed to create businesses. Seeding the database has been stopped."
        )
        return 1
    create_transactions(
        finance_database, [business.id for business in businesses], user.id
    )
    if create_accounts_and_totals(finance_database, user.id):
        print("Sucessfully seeded the database.")
        return 0
    else:
        print("Failed to seed the accounts and totals.")
        return 1

seed_db()