diff --git a/backend/generate_submission_example_data.py b/backend/generate_submission_example_data.py new file mode 100644 index 0000000..d5b3da8 --- /dev/null +++ b/backend/generate_submission_example_data.py @@ -0,0 +1,261 @@ +#!/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!")