seed_data.py

import sqlite3
import random
from datetime import datetime
from faker import Faker
from pathlib import Path
import create_tables

DB_PATH = "retail.db"
NUM_PRODUCTS = 100
NUM_STORES = 10
NUM_CUSTOMERS = 300
NUM_ORDERS = 2000  

REGIONS = ["North", "South", "East", "West"]
CATEGORIES = ["Electronics", "Home", "Garden", "Clothing", "Toys", "Sports"]
STATUSES = ["completed", "returned"]

SEED = 42


def _connect(db_path=DB_PATH):
    conn = sqlite3.connect(db_path)
    conn.execute("PRAGMA foreign_keys = ON;")
    return conn


def seed(db_path: str = DB_PATH,
         num_products=NUM_PRODUCTS,
         num_stores=NUM_STORES,
         num_customers=NUM_CUSTOMERS,
         num_orders=NUM_ORDERS):
    fake = Faker()
    Faker.seed(SEED)
    random.seed(SEED)

    db_file = Path(db_path)

    if create_tables is not None:
        try:
            create_tables.create_tables(db_path)
        except Exception:
            pass

    conn = _connect(db_path)
    cur = conn.cursor()

    products = []
    for i in range(1, num_products + 1):
        sku = f"SKU{i:04d}"
        name = f"Product {i}"
        category = random.choice(CATEGORIES)
        price = round(random.uniform(5, 500), 2)
        products.append((sku, name, category, price))

    cur.executemany(
        "INSERT OR IGNORE INTO products (sku, name, category, price) VALUES (?, ?, ?, ?);",
        products,
    )

    stores = []
    for i in range(1, num_stores + 1):
        name = f"Store {i}"
        city = fake.city()
        state = fake.state()
        region = random.choice(REGIONS)
        stores.append((name, city, state, region))

    cur.executemany(
        "INSERT OR IGNORE INTO stores (name, city, state, region) VALUES (?, ?, ?, ?);",
        stores,
    )

    customers = []
    for _ in range(num_customers):
        name = fake.name()
        email = fake.email()
        city = fake.city()
        state = fake.state()
        signup_date = fake.date_between(start_date='-3y', end_date='today').isoformat()
        customers.append((name, email, city, state, signup_date))

    cur.executemany(
        "INSERT INTO customers (name, email, city, state, signup_date) VALUES (?, ?, ?, ?, ?);",
        customers,
    )

    conn.commit()

    cur.execute("SELECT id, price FROM products;")
    prod_rows = cur.fetchall()  

    cur.execute("SELECT id FROM stores;")
    store_ids = [r[0] for r in cur.fetchall()]

    cur.execute("SELECT id FROM customers;")
    cust_ids = [r[0] for r in cur.fetchall()]

    if not prod_rows or not store_ids or not cust_ids:
        raise RuntimeError("Products, stores or customers missing after seeding - aborting orders creation.")

    orders_created = 0
    items_created = 0
    for _ in range(num_orders):
        customer_id = random.choice(cust_ids)
        store_id = random.choice(store_ids)
        order_date = fake.date_between(start_date='-2y', end_date='today').isoformat()

        item_count = random.randint(1, 4)
        chosen = random.sample(prod_rows, k=item_count)

        order_items = []  
        total_amount = 0.0
        for pid, price in chosen:
            qty = random.randint(1, 5)
            unit_price = price
            total_amount += round(qty * unit_price, 2)
            order_items.append((pid, qty, unit_price))

        status = random.choices(STATUSES, weights=(0.95, 0.05))[0]

        cur.execute(
            "INSERT INTO orders (customer_id, store_id, order_date, total_amount, status) VALUES (?, ?, ?, ?, ?);",
            (customer_id, store_id, order_date, round(total_amount, 2), status),
        )
        order_id = cur.lastrowid
        orders_created += 1

        order_item_rows = [(order_id, pid, qty, unit_price) for (pid, qty, unit_price) in order_items]

        cur.executemany(
            "INSERT INTO order_items (order_id, product_id, quantity, unit_price) VALUES (?, ?, ?, ?);",
            order_item_rows,
        )
        items_created += len(order_item_rows)

        if orders_created % 500 == 0:
            conn.commit()

    conn.commit()

    print(f"Seeded: {len(prod_rows)} products, {len(store_ids)} stores, {len(cust_ids)} customers")
    print(f"Created {orders_created} orders and {items_created} order_items in {db_file.resolve()}")

    conn.close()


if __name__ == "__main__":
    seed(DB_PATH)