load_fixtures.py - master - Askoline/Homework03

load_fixtures.py
294 строки9 KB
Askoline
first_commit
05 авг 2026, 14:29
05 авг 2026, 14:29a97374b
# -*- coding: utf-8 -*-
# load_fixtures.py
# Script to load data from JSON fixtures into database

import json
import os
from datetime import datetime
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from pathlib import Path

from models import Publisher, Book, Shop, Stock, Sale, Base
from config import get_database_url, FIXTURES_DIR


class FixtureLoader:
    """Class for loading fixtures into database"""
    
    def __init__(self, engine):
        self.engine = engine
        self.Session = sessionmaker(bind=engine)
        self.session = self.Session()
        
        # Cache for quick lookups
        self.publisher_cache = {}  # name -> publisher object
        self.shop_cache = {}       # name -> shop object
        self.book_cache = {}       # (title, publisher_name) -> book object
        self.stock_cache = {}      # (book_title, shop_name) -> stock object
    
    def load_publishers(self, filepath):
        """Load publishers from JSON file"""
        print(f"Loading publishers from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        for item in data:
            publisher = Publisher(name=item['name'])
            self.session.add(publisher)
            self.publisher_cache[item['name']] = publisher
        
        self.session.commit()
        print(f"  Loaded {len(data)} publishers")
    
    def load_shops(self, filepath):
        """Load shops from JSON file"""
        print(f"Loading shops from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        for item in data:
            shop = Shop(name=item['name'])
            self.session.add(shop)
            self.shop_cache[item['name']] = shop
        
        self.session.commit()
        print(f"  Loaded {len(data)} shops")
    
    def load_books(self, filepath):
        """Load books from JSON file"""
        print(f"Loading books from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Refresh caches to get IDs
        self._refresh_caches()
        
        for item in data:
            publisher_name = item['publisher_name']
            publisher = self.publisher_cache.get(publisher_name)
            
            if not publisher:
                print(f"  Warning: Publisher '{publisher_name}' not found, skipping book '{item['title']}'")
                continue
            
            book = Book(
                title=item['title'],
                id_publisher=publisher.id
            )
            self.session.add(book)
            self.book_cache[(item['title'], publisher_name)] = book
        
        self.session.commit()
        print(f"  Loaded {len(data)} books")
    
    def load_stock(self, filepath):
        """Load stock entries from JSON file"""
        print(f"Loading stock from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Refresh caches to get IDs
        self._refresh_caches()
        
        for item in data:
            book_title = item['book_title']
            shop_name = item['shop_name']
            count = item['count']
            
            # Find book and shop
            book = None
            for (title, publisher_name), book_obj in self.book_cache.items():
                if title == book_title:
                    book = book_obj
                    break
            
            shop = self.shop_cache.get(shop_name)
            
            if not book:
                print(f"  Warning: Book '{book_title}' not found, skipping stock entry")
                continue
            
            if not shop:
                print(f"  Warning: Shop '{shop_name}' not found, skipping stock entry")
                continue
            
            stock = Stock(
                id_book=book.id,
                id_shop=shop.id,
                count=count
            )
            self.session.add(stock)
            self.stock_cache[(book_title, shop_name)] = stock
        
        self.session.commit()
        print(f"  Loaded {len(data)} stock entries")
    
    def load_sales(self, filepath):
        """Load sales from JSON file"""
        print(f"Loading sales from {filepath}...")
        
        with open(filepath, 'r', encoding='utf-8') as f:
            data = json.load(f)
        
        # Refresh caches to get IDs
        self._refresh_caches()
        
        for item in data:
            book_title = item['book_title']
            shop_name = item['shop_name']
            
            # Find stock entry
            stock = self.stock_cache.get((book_title, shop_name))
            
            if not stock:
                print(f"  Warning: Stock for '{book_title}' in '{shop_name}' not found, skipping sale")
                continue
            
            # Parse date
            try:
                sale_date = datetime.strptime(item['date'], "%Y-%m-%d %H:%M:%S")
            except ValueError:
                sale_date = datetime.now()
            
            sale = Sale(
                id_stock=stock.id,
                count=item['quantity'],
                price=item['price'],
                date_sale=sale_date
            )
            self.session.add(sale)
            
            # Update stock count
            stock.count -= item['quantity']
        
        self.session.commit()
        print(f"  Loaded {len(data)} sales")
    
    def _refresh_caches(self):
        """Refresh caches with data from database"""
        # Refresh publishers
        publishers = self.session.query(Publisher).all()
        self.publisher_cache = {p.name: p for p in publishers}
        
        # Refresh shops
        shops = self.session.query(Shop).all()
        self.shop_cache = {s.name: s for s in shops}
        
        # Refresh books
        books = self.session.query(Book).all()
        self.book_cache = {}
        for book in books:
            publisher = self.session.query(Publisher).filter_by(id=book.id_publisher).first()
            if publisher:
                self.book_cache[(book.title, publisher.name)] = book
        
        # Refresh stock
        stocks = self.session.query(Stock).all()
        self.stock_cache = {}
        for stock in stocks:
            book = self.session.query(Book).filter_by(id=stock.id_book).first()
            shop = self.session.query(Shop).filter_by(id=stock.id_shop).first()
            if book and shop:
                self.stock_cache[(book.title, shop.name)] = stock
    
    def load_all(self):
        """Load all fixtures in correct order"""
        try:
            # Load in correct order (respecting foreign keys)
            fixtures = [
                ('publishers.json', self.load_publishers),
                ('shops.json', self.load_shops),
                ('books.json', self.load_books),
                ('stock.json', self.load_stock),
                ('sales.json', self.load_sales),
            ]
            
            for filename, loader_func in fixtures:
                filepath = FIXTURES_DIR / filename
                if filepath.exists():
                    loader_func(filepath)
                else:
                    print(f"Warning: {filepath} not found, skipping")
            
            print("\n" + "="*60)
            print("FIXTURES LOADING COMPLETED SUCCESSFULLY!")
            print("="*60)
            
            # Print summary
            self.print_summary()
            
        except Exception as e:
            self.session.rollback()
            print(f"\nError loading fixtures: {e}")
            import traceback
            traceback.print_exc()
            raise
        finally:
            self.session.close()
    
    def print_summary(self):
        """Print database summary"""
        session = self.Session()
        
        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"\nDatabase Summary:")
        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}")
        
        session.close()
    
    def clear_database(self):
        """Clear all data from database"""
        print("Clearing database...")
        self.session.query(Sale).delete()
        self.session.query(Stock).delete()
        self.session.query(Book).delete()
        self.session.query(Shop).delete()
        self.session.query(Publisher).delete()
        self.session.commit()
        print("Database cleared")


def main():
    """Main function"""
    print("="*60)
    print("FIXTURE LOADER")
    print("="*60)
    
    # Get database URL
    db_url = get_database_url()
    print(f"\nUsing database: {db_url}")
    
    # Create engine
    engine = create_engine(db_url, echo=False)
    
    # Create tables if they don't exist
    Base.metadata.create_all(engine)
    
    # Create loader
    loader = FixtureLoader(engine)
    
    # Ask if user wants to clear existing data
    response = input("\nClear existing data before loading? (y/n): ").strip().lower()
    if response == 'y':
        loader.clear_database()
    
    # Load fixtures
    loader.load_all()


if __name__ == "__main__":
    main()