seed_data.py - master - Askoline/Homework03

seed_data.py
172 строки6 KB
Askoline
first_commit
05 авг 2026, 14:29
05 авг 2026, 14:29a97374b
# -*- coding: utf-8 -*-
# seed_data.py
# Script to populate database with test data

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from datetime import datetime, timedelta
import random

from models import Publisher, Book, Shop, Stock, Sale, Base


def get_engine(db_url="sqlite:///bookstore.db"):
    return create_engine(db_url, echo=False)


def get_session(engine):
    Session = sessionmaker(bind=engine)
    return Session()


def seed_database():
    """Populate database with test data"""
    print("Seeding database with test data...")
    
    engine = get_engine()
    session = get_session(engine)
    
    # Create tables
    Base.metadata.create_all(engine)
    
    try:
        # Clear existing data (optional)
        session.query(Sale).delete()
        session.query(Stock).delete()
        session.query(Book).delete()
        session.query(Shop).delete()
        session.query(Publisher).delete()
        session.commit()
        
        # Create publishers
        publishers = [
            Publisher(name="Pushkin"),
            Publisher(name="Tolstoy"),
            Publisher(name="Dostoevsky"),
            Publisher(name="Chekhov"),
            Publisher(name="Gogol"),
        ]
        session.add_all(publishers)
        session.commit()
        
        # Create shops
        shops = [
            Shop(name="Bookvoed"),
            Shop(name="Labyrinth"),
            Shop(name="Book House"),
            Shop(name="Read City"),
            Shop(name="Book World"),
        ]
        session.add_all(shops)
        session.commit()
        
        # Create books for each publisher
        books_data = [
            # Pushkin
            {"title": "Captain's Daughter", "publisher": "Pushkin"},
            {"title": "Ruslan and Ludmila", "publisher": "Pushkin"},
            {"title": "Eugene Onegin", "publisher": "Pushkin"},
            {"title": "The Bronze Horseman", "publisher": "Pushkin"},
            # Tolstoy
            {"title": "War and Peace", "publisher": "Tolstoy"},
            {"title": "Anna Karenina", "publisher": "Tolstoy"},
            {"title": "Resurrection", "publisher": "Tolstoy"},
            # Dostoevsky
            {"title": "Crime and Punishment", "publisher": "Dostoevsky"},
            {"title": "The Idiot", "publisher": "Dostoevsky"},
            {"title": "The Brothers Karamazov", "publisher": "Dostoevsky"},
            # Chekhov
            {"title": "The Cherry Orchard", "publisher": "Chekhov"},
            {"title": "Uncle Vanya", "publisher": "Chekhov"},
            {"title": "The Seagull", "publisher": "Chekhov"},
            # Gogol
            {"title": "Dead Souls", "publisher": "Gogol"},
            {"title": "The Government Inspector", "publisher": "Gogol"},
            {"title": "The Overcoat", "publisher": "Gogol"},
        ]
        
        # Dictionary to store created books
        book_objects = []
        
        for book_data in books_data:
            publisher = session.query(Publisher).filter_by(name=book_data["publisher"]).first()
            if publisher:
                book = Book(title=book_data["title"], id_publisher=publisher.id)
                session.add(book)
                book_objects.append(book)
        session.commit()
        
        # Create stock entries and sales
        base_date = datetime.now() - timedelta(days=60)
        
        for book in book_objects:
            # Each book is available in 2-4 shops
            num_shops = random.randint(2, 4)
            selected_shops = random.sample(shops, num_shops)
            
            for shop in selected_shops:
                # Create stock entry
                stock_count = random.randint(20, 100)
                stock = Stock(
                    id_book=book.id,
                    id_shop=shop.id,
                    count=stock_count
                )
                session.add(stock)
                session.flush()  # Get stock ID
                
                # Create 3-8 sales for this stock
                num_sales = random.randint(3, 8)
                for i in range(num_sales):
                    # Random date within the last 60 days
                    sale_date = base_date + timedelta(days=random.randint(0, 60))
                    sale_date = sale_date.replace(hour=random.randint(10, 20), minute=random.randint(0, 59))
                    
                    # Random quantity (1-5 books)
                    quantity = random.randint(1, 5)
                    
                    # Random price (500-1500 rubles)
                    price = round(random.uniform(300, 1500), 2)
                    
                    sale = Sale(
                        id_stock=stock.id,
                        count=quantity,
                        price=price,
                        date_sale=sale_date
                    )
                    session.add(sale)
                    
                    # Update stock count
                    stock.count -= quantity
        
        session.commit()
        print("Test data successfully added!")
        
        # Print summary
        print("\n" + "="*60)
        print("DATABASE SUMMARY")
        print("="*60)
        
        publishers_count = session.query(Publisher).count()
        books_count = session.query(Book).count()
        shops_count = session.query(Shop).count()
        stocks_count = session.query(Stock).count()
        sales_count = session.query(Sale).count()
        
        print(f"Publishers: {publishers_count}")
        print(f"Books: {books_count}")
        print(f"Shops: {shops_count}")
        print(f"Stock entries: {stocks_count}")
        print(f"Sales: {sales_count}")
        print("="*60)
        
    except Exception as e:
        session.rollback()
        print(f"Error seeding database: {e}")
        raise
    finally:
        session.close()


if __name__ == "__main__":
    seed_database()