generate_data.py

import random
import uuid
import argparse
import sys
from datetime import timedelta

import psycopg2
from faker import Faker


fake = Faker()

genres = [
    "Fiction",
    "Fantasy",
    "Mystery",
    "Sci-Fi",
    "Romance",
    "History",
    "Biography",
    "Non-fiction",
    "Thriller",
    "Classic",
]


def unique_email():
    return f"{fake.first_name().lower()}.{uuid.uuid4().hex[:12]}@example.com"


def unique_isbn():
    return f"978-{uuid.uuid4().hex[:12]}"


def generate_book_title():
    words = fake.words(nb=random.randint(2, 4))
    words = [word.capitalize() for word in words]
    title = " ".join(words)

    if random.random() < 0.4:
        title = "The " + title

    return title


def connect():
    return psycopg2.connect(
        user="postgres",
        dbname="library",
    )


def insert_authors(cur, n):
    for _ in range(n):
        cur.execute(
            "INSERT INTO authors (name, country) VALUES (%s, %s)",
            (fake.name(), fake.country()),
        )


def insert_members(cur, n):
    for _ in range(n):
        cur.execute(
            "INSERT INTO members (name, email, joined_at) VALUES (%s, %s, %s)",
            (
                fake.name(),
                unique_email(),
                fake.date_time_between(start_date="-5y", end_date="now"),
            ),
        )


def insert_books(cur, n):
    cur.execute("SELECT id FROM authors")
    author_ids = [row[0] for row in cur.fetchall()]

    if not author_ids:
        raise RuntimeError("No authors found. Insert authors before inserting books.")

    total_copies = 0

    for _ in range(n):
        cur.execute(
            """
            INSERT INTO books (title, author_id, genre, published_year, isbn, description)
            VALUES (%s, %s, %s, %s, %s, %s)
            RETURNING id
            """,
            (
                generate_book_title(),
                random.choice(author_ids),
                random.choice(genres),
                random.randint(1900, 2025),
                unique_isbn(),
                fake.text(max_nb_chars=random.randint(500, 1500)),
            ),
        )

        book_id = cur.fetchone()[0]
        total_copies += insert_copies_for_book(cur, book_id)

    return total_copies


def insert_copies_for_book(cur, book_id):
    used_shelves = set()
    copy_count = 0

    for _ in range(random.randint(1, 6)):
        while True:
            shelf = f"{random.choice('ABCDEFGHIJKLMNOPQRSTUVWXYZ')}{random.randint(1, 80)}"
            if shelf not in used_shelves:
                used_shelves.add(shelf)
                break

        cur.execute(
            """
            INSERT INTO copies (book_id, shelf_location, status)
            VALUES (%s, %s, %s)
            """,
            (
                book_id,
                shelf,
                random.choice(["available", "available", "available", "borrowed", "reserved"]),
            ),
        )
        copy_count += 1

    return copy_count


def insert_loans_bulk(cur, n, batch_size=5000):
    cur.execute("SELECT id FROM members")
    member_ids = [row[0] for row in cur.fetchall()]

    cur.execute("SELECT id FROM copies")
    copy_ids = [row[0] for row in cur.fetchall()]

    if not member_ids:
        raise RuntimeError("No members found. Insert members before inserting loans.")

    if not copy_ids:
        raise RuntimeError("No copies found. Insert books/copies before inserting loans.")

    inserted = 0

    while inserted < n:
        rows = []
        current_batch = min(batch_size, n - inserted)

        for _ in range(current_batch):
            loan_date = fake.date_time_between(start_date="-5y", end_date="now")
            due_date = loan_date + timedelta(days=random.choice([14, 21, 30]))

            if random.random() < 0.12:
                return_date = None
            else:
                return_date = loan_date + timedelta(days=random.randint(1, 60))

            rows.append((
                random.choice(member_ids),
                random.choice(copy_ids),
                loan_date,
                due_date,
                return_date,
            ))

        cur.executemany(
            """
            INSERT INTO loans (member_id, copy_id, loan_date, due_date, return_date)
            VALUES (%s, %s, %s, %s, %s)
            """,
            rows,
        )

        inserted += current_batch

    return inserted


def insert_reservations(cur, n, batch_size=5000):
    cur.execute("SELECT id FROM books")
    book_ids = [row[0] for row in cur.fetchall()]

    cur.execute("SELECT id FROM members")
    member_ids = [row[0] for row in cur.fetchall()]

    if not book_ids:
        raise RuntimeError("No books found. Insert books before inserting reservations.")

    if not member_ids:
        raise RuntimeError("No members found. Insert members before inserting reservations.")

    inserted = 0

    while inserted < n:
        rows = []
        current_batch = min(batch_size, n - inserted)

        for _ in range(current_batch):
            rows.append((
                random.choice(member_ids),
                random.choice(book_ids),
                fake.date_time_between(start_date="-3y", end_date="now"),
                random.choice(["active", "cancelled", "fulfilled", "fulfilled", "fulfilled"]),
            ))

        cur.executemany(
            """
            INSERT INTO reservations (member_id, book_id, reserved_at, status)
            VALUES (%s, %s, %s, %s)
            """,
            rows,
        )

        inserted += current_batch

    return inserted


def insert_reviews(cur, n, batch_size=5000):
    cur.execute("SELECT id FROM books")
    book_ids = [row[0] for row in cur.fetchall()]

    cur.execute("SELECT id FROM members")
    member_ids = [row[0] for row in cur.fetchall()]

    if not book_ids:
        raise RuntimeError("No books found. Insert books before inserting reviews.")

    if not member_ids:
        raise RuntimeError("No members found. Insert members before inserting reviews.")

    inserted = 0

    while inserted < n:
        rows = []
        current_batch = min(batch_size, n - inserted)

        for _ in range(current_batch):
            rows.append((
                random.choice(member_ids),
                random.choice(book_ids),
                random.randint(1, 5),
                fake.sentence(nb_words=random.randint(6, 14)),
                fake.date_time_between(start_date="-4y", end_date="now"),
                fake.text(max_nb_chars=random.randint(100, 800)),
            ))

        cur.executemany(
            """
            INSERT INTO reviews (member_id, book_id, rating, comment, created_at, review_text)
            VALUES (%s, %s, %s, %s, %s, %s)
            """,
            rows,
        )

        inserted += current_batch

    return inserted


def main(authors, members, books, loans, reservations, reviews):
    conn = None

    try:
        conn = connect()
        cur = conn.cursor()

        insert_authors(cur, authors)
        print(f"Added {authors} authors")

        insert_members(cur, members)
        print(f"Added {members} members")

        total_copies = insert_books(cur, books)
        print(f"Added {books} books and {total_copies} copies")

        inserted_loans = insert_loans_bulk(cur, loans)
        print(f"Added {inserted_loans} loans")

        inserted_reservations = insert_reservations(cur, reservations)
        print(f"Added {inserted_reservations} reservations")

        inserted_reviews = insert_reviews(cur, reviews)
        print(f"Added {inserted_reviews} reviews")

        conn.commit()
        cur.close()
        conn.close()

        print("Data generation completed successfully!")

    except Exception as e:
        if conn:
            conn.rollback()
            conn.close()

        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Generate realistic demo data for the library database.")

    parser.add_argument("--authors", type=int, default=500)
    parser.add_argument("--members", type=int, default=5000)
    parser.add_argument("--books", type=int, default=10000)
    parser.add_argument("--loans", type=int, default=100000)
    parser.add_argument("--reservations", type=int, default=50000)
    parser.add_argument("--reviews", type=int, default=50000)

    args = parser.parse_args()

    main(
        authors=args.authors,
        members=args.members,
        books=args.books,
        loans=args.loans,
        reservations=args.reservations,
        reviews=args.reviews,
    )