#!/usr/bin/env python3 """ One-off script to generate example Submission (and Review) data for testing/demoing the work-submission and review flow. Reuses the same shot-sampling approach as generate_note_example_data.py and generate_schedule_example_data.py: picks the same bounded sample of Dragon Quest shots (first 15 candidates with the normal 4-task set) rather than touching all ~1400 shot tasks, so submissions line up with the tasks that already have start_date/deadline from the schedule seed data. Not every sampled task gets a submission (a "not started" task realistically has none) - about 70% do, with 1-3 versions each. Earlier versions are given a "retake" review, and the latest version is left pending, approved, or retake'd, with the task's status updated to match. Each submission gets a real placeholder file on disk under uploads/submissions// (a tiny real image for .jpg/.png so thumbnail generation works, plain bytes for other formats) so the review UI has something to open instead of a 404. Idempotent: skips any task that already has a submission, so it's safe to re-run without creating duplicates. """ from datetime import datetime, timedelta, timezone from pathlib import Path 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, Submission, Review, ReviewDecision from models.project import Project, ProjectMember from models.user import User, UserRole from utils.file_handler import file_handler from PIL import Image import io 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 # same sample window as generate_schedule_example_data.py RNG_SEED = 42 SUBMISSION_CHANCE = 0.7 # fraction of sampled tasks that get any submission # Plausible deliverable formats per task type (all within FileHandler's # SUPPORTED_FORMATS). Kept short - this is placeholder data, not a real # pipeline convention. TASK_TYPE_EXTENSIONS = { "layout": [".ma", ".mov"], "animation": [".mov", ".ma"], "simulation": [".mov", ".ma"], "lighting": [".exr", ".mov"], "compositing": [".mov", ".png", ".exr"], "previz": [".mov"], } DEFAULT_EXTENSIONS = [".mov"] SUBMISSION_NOTES = [ "First pass, ready for review.", "Addressed previous feedback, please take another look.", "Still WIP on some details but wanted to get eyes on it early.", "Final polish pass, should be good to approve.", None, # some submissions have no note at all ] RETAKE_FEEDBACK = [ "Good start, but needs another pass before this is ready - see notes on the task.", "Close, but a few issues need fixing before this can move forward.", "Not quite there yet, please revise and resubmit.", ] APPROVE_FEEDBACK = [ "Looks great, approved!", "Nice work, this is ready to move on.", "Approved - matches the brief.", ] IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png"} def make_placeholder_file(task: Task, extension: str, version: int, rng: random.Random) -> tuple[str, int]: """Write a small real file to disk and return (relative_path, size).""" task_dir = file_handler.create_directory_structure(task.id, "submission") base_name = f"{(task.shot_id and 'shot' or 'asset')}_{task.task_type}_v{version:03d}{extension}" filename = file_handler.generate_unique_filename(base_name, f"v{version:03d}") file_path = task_dir / filename if extension in IMAGE_EXTENSIONS: # A tiny but real image so thumbnail generation (PIL.Image.open) works. color = (rng.randint(30, 220), rng.randint(30, 220), rng.randint(30, 220)) img = Image.new("RGB", (64, 64), color) buf = io.BytesIO() img.save(buf, format="JPEG" if extension in (".jpg", ".jpeg") else "PNG") file_path.write_bytes(buf.getvalue()) else: file_path.write_bytes(f"Placeholder {extension} submission file for testing.\n".encode()) relative_path = file_handler.store_relative_path(str(file_path)) return relative_path, file_path.stat().st_size def generate_submission_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) now = datetime.now(timezone.utc) try: project = db.query(Project).filter(Project.name == PROJECT_NAME).first() if not project: logger.error(f"Project '{PROJECT_NAME}' not found") return member_ids = [m.user_id for m in db.query(ProjectMember).filter( ProjectMember.project_id == project.id ).all()] if not member_ids: logger.error(f"No members found on project '{PROJECT_NAME}'") return reviewer_ids = [u.id for u in db.query(User).join( ProjectMember, ProjectMember.user_id == User.id ).filter( ProjectMember.project_id == project.id, User.role.in_([UserRole.DIRECTOR, UserRole.COORDINATOR]) ).all()] or member_ids logger.info(f"Project '{PROJECT_NAME}' (id={project.id}), {len(member_ids)} members, {len(reviewer_ids)} potential reviewers") # 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_seen = 0 tasks_skipped = 0 tasks_submitted = 0 submissions_created = 0 reviews_created = 0 for shot in sample_shots: tasks = db.query(Task).filter( Task.shot_id == shot.id, Task.deleted_at.is_(None) ).all() for task in tasks: tasks_seen += 1 existing = db.query(Submission).filter(Submission.task_id == task.id).first() if existing: tasks_skipped += 1 continue if rng.random() >= SUBMISSION_CHANCE: continue extensions = TASK_TYPE_EXTENSIONS.get(task.task_type, DEFAULT_EXTENSIONS) version_count = rng.choice([1, 1, 2, 3]) submitter_id = task.assigned_user_id or rng.choice(member_ids) # Spread submissions over the days leading up to the deadline # (or today, if the task has none) so submitted_at reads as a # real history rather than everything happening "now". anchor = ( datetime.combine(task.deadline, datetime.min.time(), tzinfo=timezone.utc) if task.deadline else now ) submitted_at = anchor - timedelta(days=version_count * 2) last_submission = None last_decision = None for version in range(1, version_count + 1): extension = rng.choice(extensions) relative_path, file_size = make_placeholder_file(task, extension, version, rng) thumbnail_path = None if extension in IMAGE_EXTENSIONS: thumbnail_path = file_handler.create_thumbnail(relative_path) submission = Submission( task_id=task.id, user_id=submitter_id, file_path=relative_path, file_name=Path(relative_path).name, version_number=version, notes=rng.choice(SUBMISSION_NOTES), submitted_at=submitted_at ) db.add(submission) db.flush() # need submission.id for a Review FK submissions_created += 1 submitted_at += timedelta(days=2) is_last_version = version == version_count if not is_last_version: # Earlier versions were superseded because they needed work. review = Review( submission_id=submission.id, reviewer_id=rng.choice(reviewer_ids), decision=ReviewDecision.RETAKE, feedback=rng.choice(RETAKE_FEEDBACK), reviewed_at=submission.submitted_at + timedelta(hours=rng.randint(2, 20)) ) db.add(review) reviews_created += 1 else: # Latest version: leave some pending review, resolve the rest. outcome = rng.choice(["pending", "approved", "approved", "retake"]) if outcome != "pending": decision = ReviewDecision.APPROVED if outcome == "approved" else ReviewDecision.RETAKE feedback = rng.choice(APPROVE_FEEDBACK if outcome == "approved" else RETAKE_FEEDBACK) review = Review( submission_id=submission.id, reviewer_id=rng.choice(reviewer_ids), decision=decision, feedback=feedback, reviewed_at=submission.submitted_at + timedelta(hours=rng.randint(2, 20)) ) db.add(review) reviews_created += 1 last_decision = outcome last_submission = submission if last_submission: task.status = { "approved": "approved", "retake": "retake", "pending": "submitted" }[last_decision] tasks_submitted += 1 db.commit() logger.info(f"Tasks seen: {tasks_seen}, skipped (already had a submission): {tasks_skipped}") logger.info(f"Tasks given submissions: {tasks_submitted}") logger.info(f"Submissions created: {submissions_created}, reviews created: {reviews_created}") except Exception as e: logger.error(f"Generation failed: {e}") db.rollback() raise finally: db.close() if __name__ == "__main__": logger.info("Generating submission example data...") generate_submission_example_data() logger.info("Done!")