indigo
9c5abf6342
Idempotent script mirroring the existing note/schedule seed-data scripts: seeds Submission + Review rows (with real placeholder files on disk) for the same sample of Dragon Quest shots, so the submissions/reviews UI has realistic data to test against. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
262 lines
11 KiB
Python
262 lines
11 KiB
Python
#!/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/<task_id>/ (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!")
|