dkAmulet/sql-query-optimizer

1"""2Database setup module.3 4Creates an in-memory SQLite e-commerce database with deterministic5seeded data used by all three tasks.6"""7import sqlite38import random9 10 11def create_database() -> sqlite3.Connection:12    """13    Build and populate an in-memory SQLite database.14 15    Schema: users, categories, products, orders, order_items16    Data:   1500 users | 10 categories | 300 products17            3000 orders | 8000 order_items18    Seed:   42 (fully deterministic)19    """20    conn = sqlite3.connect(":memory:", check_same_thread=False)21 22    # ------------------------------------------------------------------ schema23    conn.executescript("""24        CREATE TABLE users (25            user_id    INTEGER PRIMARY KEY,26            username   TEXT    NOT NULL,27            email      TEXT    NOT NULL,28            first_name TEXT,29            last_name  TEXT,30            city       TEXT,31            country    TEXT,32            is_active  INTEGER DEFAULT 1,33            created_at TEXT34        );35 36        CREATE TABLE categories (37            category_id        INTEGER PRIMARY KEY,38            name               TEXT NOT NULL,39            parent_category_id INTEGER40        );41 42        CREATE TABLE products (43            product_id   INTEGER PRIMARY KEY,44            name         TEXT    NOT NULL,45            category_id  INTEGER,46            price        REAL,47            sku          TEXT,48            is_available INTEGER DEFAULT 149        );50 51        CREATE TABLE orders (52            order_id     INTEGER PRIMARY KEY,53            user_id      INTEGER REFERENCES users(user_id),54            status       TEXT,55            total_amount REAL,56            created_at   TEXT57        );58 59        CREATE TABLE order_items (60            item_id    INTEGER PRIMARY KEY,61            order_id   INTEGER REFERENCES orders(order_id),62            product_id INTEGER REFERENCES products(product_id),63            quantity   INTEGER,64            unit_price REAL65        );66 67        -- Indexes that an optimal query should exploit68        CREATE INDEX idx_users_active  ON users(is_active);69        CREATE INDEX idx_users_country ON users(country, is_active);70        CREATE INDEX idx_orders_user   ON orders(user_id);71        CREATE INDEX idx_orders_status ON orders(status);72        CREATE INDEX idx_items_order   ON order_items(order_id);73        CREATE INDEX idx_items_product ON order_items(product_id);74        CREATE INDEX idx_products_cat  ON products(category_id);75    """)76 77    # ------------------------------------------------------------------ seed78    rng = random.Random(42)79 80    # categories (2 levels)81    categories = [82        (1, "Electronics", None),83        (2, "Clothing",    None),84        (3, "Books",       None),85        (4, "Smartphones", 1),86        (5, "Laptops",     1),87        (6, "Tablets",     1),88        (7, "T-Shirts",    2),89        (8, "Jeans",       2),90        (9, "Fiction",     3),91        (10, "Non-Fiction", 3),92    ]93    conn.executemany("INSERT INTO categories VALUES (?,?,?)", categories)94 95    # users (1 500 rows)96    countries = ["USA", "USA", "USA", "UK", "Canada", "Germany", "France", "Australia"]97    cities    = ["New York", "Los Angeles", "Chicago", "London",98                 "Toronto", "Berlin", "Paris", "Sydney"]99    users = [100        (101            i, f"user_{i}", f"user{i}@example.com",102            f"First{i}", f"Last{i}",103            rng.choice(cities), rng.choice(countries),104            1 if rng.random() > 0.25 else 0,105            f"202{rng.randint(0,3)}-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",106        )107        for i in range(1, 1501)108    ]109    conn.executemany("INSERT INTO users VALUES (?,?,?,?,?,?,?,?,?)", users)110 111    # products (300 rows)112    leaf_cats = [4, 5, 6, 7, 8, 9, 10]113    products = [114        (115            i, f"Product_{i:04d}", rng.choice(leaf_cats),116            round(rng.uniform(9.99, 1499.99), 2),117            f"SKU-{i:05d}", 1 if rng.random() > 0.05 else 0,118        )119        for i in range(1, 301)120    ]121    conn.executemany("INSERT INTO products VALUES (?,?,?,?,?,?)", products)122 123    # orders (3 000 rows)124    statuses = ["pending", "processing", "shipped", "delivered", "cancelled"]125    weights  = [0.10, 0.10, 0.20, 0.50, 0.10]126    orders = [127        (128            i, rng.randint(1, 1500),129            rng.choices(statuses, weights=weights)[0],130            round(rng.uniform(20.0, 3000.0), 2),131            f"2024-{rng.randint(1,12):02d}-{rng.randint(1,28):02d}",132        )133        for i in range(1, 3001)134    ]135    conn.executemany("INSERT INTO orders VALUES (?,?,?,?,?)", orders)136 137    # order_items (8 000 rows)138    items = [139        (140            i, rng.randint(1, 3000), rng.randint(1, 300),141            rng.randint(1, 5), round(rng.uniform(9.99, 1499.99), 2),142        )143        for i in range(1, 8001)144    ]145    conn.executemany("INSERT INTO order_items VALUES (?,?,?,?,?)", items)146 147    conn.commit()148    return conn149