dbre/database.py

from __future__ import annotations

import os
import random
from typing import Optional

import psycopg2
from faker import Faker


class ConnectionError(Exception):
    pass


class DBREPostgres:

    def __init__(
        self,
        host: Optional[str] = None,
        port: Optional[str] = None,
        database: Optional[str] = None,
        user: Optional[str] = None,
        password: Optional[str] = None,
    ) -> None:
        self.host = host or os.getenv("DB_HOST", "localhost")
        self.port = port or os.getenv("DB_PORT", "5432")
        self.database = database or os.getenv("DB_NAME", "dbre")
        self.user = user or os.getenv("DB_USER", "postgres")
        self.password = password or os.getenv("DB_PASSWORD", "postgres")
        self.conn = None
        self.fake = Faker()

    def connect(self) -> None:
        try:
            self.conn = psycopg2.connect(
                host=self.host, port=self.port, database=self.database,
                user=self.user, password=self.password
            )
            self.conn.autocommit = False
        except psycopg2.Error as e:
            raise ConnectionError(f"Failed to connect: {e}")

    def create_tables(self) -> None:
        if not self.conn:
            self.connect()
        try:
            cur = self.conn.cursor()
            cur.execute("""
                CREATE TABLE IF NOT EXISTS customers (
                    customer_id SERIAL PRIMARY KEY, name TEXT NOT NULL,
                    email TEXT UNIQUE NOT NULL, city TEXT,
                    created_at TIMESTAMP DEFAULT NOW()
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS products (
                    product_id SERIAL PRIMARY KEY, name TEXT NOT NULL,
                    category TEXT NOT NULL, price NUMERIC NOT NULL,
                    stock INTEGER DEFAULT 0
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS orders (
                    order_id SERIAL PRIMARY KEY,
                    customer_id INTEGER REFERENCES customers(customer_id),
                    order_date TIMESTAMP DEFAULT NOW(),
                    status TEXT DEFAULT 'pending'
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS order_items (
                    item_id SERIAL PRIMARY KEY,
                    order_id INTEGER REFERENCES orders(order_id),
                    product_id INTEGER REFERENCES products(product_id),
                    quantity INTEGER NOT NULL,
                    unit_price NUMERIC NOT NULL
                )
            """)
            cur.execute("""
                CREATE TABLE IF NOT EXISTS reviews (
                    review_id SERIAL PRIMARY KEY,
                    customer_id INTEGER REFERENCES customers(customer_id),
                    product_id INTEGER REFERENCES products(product_id),
                    rating INTEGER CHECK(rating BETWEEN 1 AND 5),
                    review_text TEXT,
                    created_at TIMESTAMP DEFAULT NOW()
                )
            """)
            self.conn.commit()
            cur.close()
        except psycopg2.Error as e:
            self.conn.rollback()
            raise ConnectionError(f"Failed to create tables: {e}")

    def seed_data(self) -> None:
        """Seed the database with realistic test data. Drops and rebuilds schema first."""
        if not self.conn:
            self.connect()
        random.seed(42)
        Faker.seed(42)
        try:
            cur = self.conn.cursor()
            cur.execute("DROP TABLE IF EXISTS reviews, order_items, orders, products, customers CASCADE;")
            self.conn.commit()
            cur.close()
        except Exception:
            self.conn.rollback()
        
        self.create_tables()
        
        try:
            cur = self.conn.cursor()
            customer_ids = []
            for _ in range(100):
                name = self.fake.name()
                email = self.fake.unique.email()
                city = self.fake.city()
                cur.execute(
                    "INSERT INTO customers (name, email, city) VALUES (%s, %s, %s) RETURNING customer_id",
                    (name, email, city),
                )
                customer_ids.append(cur.fetchone()[0])

            categories = ["Electronics", "Clothing", "Home", "Sports", "Books"]
            product_ids = []
            for _ in range(50):
                name = self.fake.catch_phrase()
                category = random.choice(categories)
                price = round(random.uniform(10, 500), 2)
                stock = random.randint(0, 100)
                cur.execute(
                    "INSERT INTO products (name, category, price, stock) VALUES (%s, %s, %s, %s) RETURNING product_id",
                    (name, category, price, stock),
                )
                product_ids.append(cur.fetchone()[0])

            order_ids = []
            statuses = ["pending", "processing", "shipped", "delivered", "cancelled"]
            for _ in range(300):
                customer_id = random.choice(customer_ids)
                order_date = self.fake.date_time_between(start_date="-1y", end_date="now")
                status = random.choice(statuses)
                cur.execute(
                    "INSERT INTO orders (customer_id, order_date, status) VALUES (%s, %s, %s) RETURNING order_id",
                    (customer_id, order_date, status),
                )
                order_ids.append(cur.fetchone()[0])

            for _ in range(600):
                order_id = random.choice(order_ids)
                product_id = random.choice(product_ids)
                quantity = random.randint(1, 10)
                unit_price = round(random.uniform(10, 500), 2)
                cur.execute(
                    "INSERT INTO order_items (order_id, product_id, quantity, unit_price) VALUES (%s, %s, %s, %s)",
                    (order_id, product_id, quantity, unit_price),
                )

            for _ in range(200):
                customer_id = random.choice(customer_ids)
                product_id = random.choice(product_ids)
                rating = random.randint(1, 5)
                review_text = self.fake.text(max_nb_chars=200)
                cur.execute(
                    "INSERT INTO reviews (customer_id, product_id, rating, review_text) VALUES (%s, %s, %s, %s)",
                    (customer_id, product_id, rating, review_text),
                )

            self.conn.commit()
            cur.close()
        except psycopg2.Error as e:
            self.conn.rollback()
            raise ConnectionError(f"Failed to seed data: {e}")