#!/usr/bin/env python3 """ One-off script to generate example start_date/deadline data on shot tasks, for testing/demoing the Project Schedule (Gantt) page. Reuses the same shot-sampling approach as generate_note_example_data.py: the dev database is heavily populated with disposable load-test fixtures, so this picks a small, bounded sample of shots (from Dragon Quest) that have the normal 4-task set, rather than scheduling all ~1400 shot tasks. For each sampled shot, tasks are scheduled in pipeline order (layout -> animation -> simulation -> lighting -> compositing, any other task types appended after) with realistic overlapping durations, staggered around today's date so the chart's "Today" marker falls inside the visible range. Idempotent: skips any task that already has a start_date set, so it's safe to re-run without clobbering manually-edited dates. """ from datetime import date, timedelta from sqlalchemy import create_engine from sqlalchemy.orm import sessionmaker from database import DATABASE_URL, Base from models.shot import Shot from models.task import Task from models.project import Project import logging import random import models # noqa: F401 - ensures every model is registered on Base.metadata logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) PROJECT_NAME = "Dragon Quest" SAMPLE_SIZE = 15 RNG_SEED = 42 # Canonical pipeline order used to sequence a shot's tasks; anything not # listed here keeps its natural (query) order, appended at the end. PIPELINE_ORDER = ["layout", "animation", "simulation", "lighting", "compositing", "previz"] def pipeline_index(task_type: str) -> int: return PIPELINE_ORDER.index(task_type) if task_type in PIPELINE_ORDER else len(PIPELINE_ORDER) def generate_schedule_example_data(): engine = create_engine(DATABASE_URL) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base.metadata.create_all(bind=engine) db = SessionLocal() rng = random.Random(RNG_SEED) today = date.today() try: project = db.query(Project).filter(Project.name == PROJECT_NAME).first() if not project: logger.error(f"Project '{PROJECT_NAME}' not found") return logger.info(f"Project '{PROJECT_NAME}' (id={project.id})") # Shots with the normal 4-task set (excludes shots with zero tasks # and the handful of outliers with extra bulk-test tasks piled on). shots = db.query(Shot).filter( Shot.project_id == project.id, Shot.deleted_at.is_(None) ).order_by(Shot.id).all() candidates = [ s for s in shots if db.query(Task).filter(Task.shot_id == s.id, Task.deleted_at.is_(None)).count() == 4 ] sample_shots = candidates[:SAMPLE_SIZE] logger.info(f"{len(candidates)} candidate shots with a 4-task set, sampling {len(sample_shots)}") tasks_scheduled = 0 tasks_skipped = 0 for shot in sample_shots: tasks = db.query(Task).filter( Task.shot_id == shot.id, Task.deleted_at.is_(None) ).all() tasks.sort(key=lambda t: pipeline_index(t.task_type)) # Stagger each shot's pipeline start across a window that straddles # today, so the demo chart shows a mix of past/current/future work. shot_start = today + timedelta(days=rng.randint(-20, 20)) cursor = shot_start for task in tasks: if task.start_date is not None: tasks_skipped += 1 # Still advance the cursor so later tasks in this shot # don't bunch up if an earlier one was already scheduled. cursor = task.deadline or cursor continue duration = rng.randint(4, 12) task.start_date = cursor task.deadline = cursor + timedelta(days=duration) tasks_scheduled += 1 # Next task starts a little before this one ends (realistic # pipeline overlap) rather than strictly back-to-back. cursor = task.deadline - timedelta(days=rng.randint(0, 4)) db.commit() logger.info(f"Tasks scheduled: {tasks_scheduled}, skipped (already scheduled): {tasks_skipped}") except Exception as e: logger.error(f"Generation failed: {e}") db.rollback() raise finally: db.close() if __name__ == "__main__": logger.info("Generating schedule example data...") generate_schedule_example_data() logger.info("Done!")