diff --git a/app/db.py b/app/db.py index 4a7652e..679f7cc 100644 --- a/app/db.py +++ b/app/db.py @@ -289,6 +289,89 @@ CREATE TABLE IF NOT EXISTS run_log_events ( CREATE INDEX IF NOT EXISTS idx_run_logs_started ON run_logs(started_at DESC, id DESC); CREATE INDEX IF NOT EXISTS idx_run_log_events_run ON run_log_events(run_id, id); + +CREATE TABLE IF NOT EXISTS image_studio_projects ( + id INTEGER PRIMARY KEY, + account_alias TEXT NOT NULL, + account_slug TEXT NOT NULL, + account_name TEXT, + item_id TEXT NOT NULL, + target_main_count INTEGER NOT NULL DEFAULT 9, + target_detail_count INTEGER NOT NULL DEFAULT 12, + draft_prompt TEXT, + status TEXT NOT NULL DEFAULT 'active', + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + deleted_at TEXT, + deleted_reason TEXT, + UNIQUE(account_alias, item_id) +); + +CREATE INDEX IF NOT EXISTS idx_image_studio_projects_status + ON image_studio_projects(status, updated_at DESC); + +CREATE TABLE IF NOT EXISTS image_studio_assets ( + id INTEGER PRIMARY KEY, + project_id INTEGER NOT NULL REFERENCES image_studio_projects(id) ON DELETE CASCADE, + kind TEXT NOT NULL, + remote_url TEXT, + local_path TEXT, + aspect_ratio TEXT, + parent_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL, + prompt TEXT, + status TEXT NOT NULL DEFAULT 'available', + source_order INTEGER NOT NULL DEFAULT 0, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL +); + +CREATE INDEX IF NOT EXISTS idx_image_studio_assets_project_kind + ON image_studio_assets(project_id, kind, source_order, id); +CREATE INDEX IF NOT EXISTS idx_image_studio_assets_parent + ON image_studio_assets(parent_asset_id); + +CREATE TABLE IF NOT EXISTS image_studio_jobs ( + id INTEGER PRIMARY KEY, + project_id INTEGER NOT NULL REFERENCES image_studio_projects(id) ON DELETE CASCADE, + source_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL, + output_asset_id INTEGER REFERENCES image_studio_assets(id) ON DELETE SET NULL, + generation_source TEXT NOT NULL DEFAULT 'cmhub', + provider TEXT NOT NULL DEFAULT 'cmhub', + job_type TEXT NOT NULL, + task_key TEXT NOT NULL UNIQUE, + task_id TEXT, + status TEXT NOT NULL DEFAULT 'pending', + prompt TEXT, + error TEXT, + attempts INTEGER NOT NULL DEFAULT 0, + points_cost INTEGER, + points_balance INTEGER, + call_id TEXT, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + submitted_at TEXT, + finished_at TEXT +); + +CREATE INDEX IF NOT EXISTS idx_image_studio_jobs_project_status + ON image_studio_jobs(project_id, status, id); +CREATE INDEX IF NOT EXISTS idx_image_studio_jobs_task_id + ON image_studio_jobs(task_id); + +CREATE TABLE IF NOT EXISTS image_studio_selections ( + id INTEGER PRIMARY KEY, + project_id INTEGER NOT NULL REFERENCES image_studio_projects(id) ON DELETE CASCADE, + selection_type TEXT NOT NULL, + position INTEGER NOT NULL, + asset_id INTEGER NOT NULL REFERENCES image_studio_assets(id) ON DELETE CASCADE, + created_at TEXT NOT NULL, + updated_at TEXT NOT NULL, + UNIQUE(project_id, selection_type, position), + UNIQUE(project_id, selection_type, asset_id) +); + +CREATE INDEX IF NOT EXISTS idx_image_studio_selections_project + ON image_studio_selections(project_id, selection_type, position); """ diff --git a/app/image_studio.py b/app/image_studio.py new file mode 100644 index 0000000..ce0c7d2 --- /dev/null +++ b/app/image_studio.py @@ -0,0 +1,565 @@ +"""Data service for the AI image studio.""" + +from __future__ import annotations + +import os +import uuid +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import datetime +from typing import Iterable, Optional + +from . import appconfig, db +from .config import make_slug + + +PROJECT_STATUS_ACTIVE = "active" +ASSET_STATUS_AVAILABLE = "available" +ASSET_STATUS_MISSING = "missing" +ASSET_STATUSES = {ASSET_STATUS_AVAILABLE, ASSET_STATUS_MISSING} +JOB_STATUSES = {"pending", "submitted", "running", "succeeded", "failed", "expired", "cancelled"} +JOB_RESUMABLE_STATUSES = {"submitted", "running"} +SELECTION_TYPES = {"main", "detail"} + + +@dataclass(frozen=True) +class ImageStudioProject: + id: int + account_alias: str + account_slug: str + account_name: Optional[str] + item_id: str + target_main_count: int + target_detail_count: int + draft_prompt: Optional[str] + status: str + created_at: str + updated_at: str + deleted_at: Optional[str] + deleted_reason: Optional[str] + + +@dataclass(frozen=True) +class ImageStudioAsset: + id: int + project_id: int + kind: str + remote_url: Optional[str] + local_path: Optional[str] + aspect_ratio: Optional[str] + parent_asset_id: Optional[int] + prompt: Optional[str] + status: str + source_order: int + created_at: str + updated_at: str + + +@dataclass(frozen=True) +class ImageStudioJob: + id: int + project_id: int + source_asset_id: Optional[int] + output_asset_id: Optional[int] + generation_source: str + provider: str + job_type: str + task_key: str + task_id: Optional[str] + status: str + prompt: Optional[str] + error: Optional[str] + attempts: int + points_cost: Optional[int] + points_balance: Optional[int] + call_id: Optional[str] + created_at: str + updated_at: str + submitted_at: Optional[str] + finished_at: Optional[str] + + +@dataclass(frozen=True) +class ImageStudioSelection: + id: int + project_id: int + selection_type: str + position: int + asset_id: int + created_at: str + updated_at: str + + +def _now() -> str: + return datetime.now().isoformat(timespec="seconds") + + +@contextmanager +def _connection(conn=None, path=None): + with db._connection(conn, path) as database: + yield database + + +def _row_to_dataclass(row, cls): + return None if row is None else cls(**dict(row)) + + +def _fetch_one(conn, sql, params, cls): + return _row_to_dataclass(conn.execute(sql, params).fetchone(), cls) + + +def _fetch_all(conn, sql, params, cls): + return [_row_to_dataclass(row, cls) for row in conn.execute(sql, params).fetchall()] + + +def _get(obj, name, default=None): + if obj is None: + return default + if isinstance(obj, dict): + return obj.get(name, default) + return getattr(obj, name, default) + + +def _safe_component(value, default): + text = str(value or "").strip() + if not text: + text = str(default) + safe = "".join(ch if ch.isalnum() or ch in "_-" else "_" for ch in text).strip("_") + return safe or str(default) + + +def _account_fields(account=None, account_alias=None, account_name=None, account_slug=None): + alias = str(account_alias or _get(account, "alias") or "").strip() + if not alias: + raise db.DbError("AI工场项目缺少账号别名") + name = str(account_name or _get(account, "account_name") or "").strip() or None + slug = str(account_slug or _get(account, "slug") or "").strip() or make_slug(alias) + return alias, name, _safe_component(slug, "account") + + +def _task_key(project_id): + return f"image-studio-{int(project_id)}-{uuid.uuid4().hex}" + + +def _ensure_assets_belong_to_project(database, project_id, asset_ids): + ids = [int(asset_id) for asset_id in asset_ids if asset_id is not None] + if not ids: + return + placeholders = ",".join("?" for _ in ids) + rows = database.execute( + f"SELECT id FROM image_studio_assets WHERE project_id = ? AND id IN ({placeholders})", + [int(project_id), *ids], + ).fetchall() + found = {int(row["id"]) for row in rows} + missing = sorted(set(ids) - found) + if missing: + raise db.DbError("AI工场资产不属于当前项目") + + +def get_project(project_id, path=None, conn=None, include_deleted=False): + clauses = ["id = ?"] + params = [int(project_id)] + if not include_deleted: + clauses.append("deleted_at IS NULL") + sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses) + with _connection(conn, path) as database: + return _fetch_one(database, sql, params, ImageStudioProject) + + +def get_project_by_account_item(account_alias, item_id, path=None, conn=None, include_deleted=False): + clauses = ["account_alias = ?", "item_id = ?"] + params = [str(account_alias).strip(), str(item_id).strip()] + if not include_deleted: + clauses.append("deleted_at IS NULL") + sql = "SELECT * FROM image_studio_projects WHERE " + " AND ".join(clauses) + with _connection(conn, path) as database: + return _fetch_one(database, sql, params, ImageStudioProject) + + +def create_or_get_project( + account=None, + item_id=None, + *, + account_alias=None, + account_name=None, + account_slug=None, + target_main_count=9, + target_detail_count=12, + draft_prompt="", + path=None, + conn=None, +): + alias, name, slug = _account_fields( + account, + account_alias=account_alias, + account_name=account_name, + account_slug=account_slug, + ) + item = str(item_id or _get(account, "item_id") or "").strip() + if not item: + raise db.DbError("AI工场项目缺少商品ID") + now = _now() + with _connection(conn, path) as database: + existing = get_project_by_account_item(alias, item, conn=database, include_deleted=True) + if existing is not None: + if existing.deleted_at is not None: + with database: + database.execute( + """ + UPDATE image_studio_projects + SET deleted_at = NULL, + deleted_reason = NULL, + status = ?, + updated_at = ? + WHERE id = ? + """, + (PROJECT_STATUS_ACTIVE, now, existing.id), + ) + return get_project(existing.id, conn=database) + return existing + with database: + cursor = database.execute( + """ + INSERT INTO image_studio_projects + (account_alias, account_slug, account_name, item_id, + target_main_count, target_detail_count, draft_prompt, + status, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + alias, + slug, + name, + item, + int(target_main_count), + int(target_detail_count), + str(draft_prompt or ""), + PROJECT_STATUS_ACTIVE, + now, + now, + ), + ) + project_id = cursor.lastrowid + return get_project(project_id, conn=database) + + +def list_projects(path=None, conn=None, include_deleted=False): + sql = "SELECT * FROM image_studio_projects" + if not include_deleted: + sql += " WHERE deleted_at IS NULL" + sql += " ORDER BY updated_at DESC, id DESC" + with _connection(conn, path) as database: + return _fetch_all(database, sql, (), ImageStudioProject) + + +def update_project_prompt(project_id, draft_prompt, path=None, conn=None): + with _connection(conn, path) as database: + with database: + database.execute( + """ + UPDATE image_studio_projects + SET draft_prompt = ?, updated_at = ? + WHERE id = ? AND deleted_at IS NULL + """, + (str(draft_prompt or ""), _now(), int(project_id)), + ) + return get_project(project_id, conn=database) + + +def soft_delete_project(project_id, reason="", path=None, conn=None): + now = _now() + with _connection(conn, path) as database: + with database: + database.execute( + """ + UPDATE image_studio_projects + SET deleted_at = ?, deleted_reason = ?, updated_at = ? + WHERE id = ? AND deleted_at IS NULL + """, + (now, str(reason or ""), now, int(project_id)), + ) + return get_project(project_id, conn=database, include_deleted=True) + + +def project_image_dirs(image_root, project): + root = os.path.abspath( + os.path.join( + str(image_root or "images"), + "pool", + _safe_component(_get(project, "account_slug"), "account"), + _safe_component(_get(project, "item_id"), "item"), + ) + ) + return { + "root": root, + "originals": os.path.join(root, "originals"), + "generated": os.path.join(root, "generated"), + "exports": os.path.join(root, "exports"), + } + + +def default_project_image_dirs(project, config=None): + return project_image_dirs(appconfig.image_dir(config), project) + + +def add_asset( + project_id, + kind, + *, + remote_url=None, + local_path=None, + aspect_ratio=None, + parent_asset_id=None, + prompt=None, + status=ASSET_STATUS_AVAILABLE, + source_order=0, + path=None, + conn=None, +): + now = _now() + abs_local_path = os.path.abspath(local_path) if local_path else None + with _connection(conn, path) as database: + with database: + cursor = database.execute( + """ + INSERT INTO image_studio_assets + (project_id, kind, remote_url, local_path, aspect_ratio, + parent_asset_id, prompt, status, source_order, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) + """, + ( + int(project_id), + str(kind), + remote_url, + abs_local_path, + aspect_ratio, + parent_asset_id, + prompt, + str(status or ASSET_STATUS_AVAILABLE), + int(source_order or 0), + now, + now, + ), + ) + asset_id = cursor.lastrowid + return get_asset(asset_id, conn=database) + + +def get_asset(asset_id, path=None, conn=None): + with _connection(conn, path) as database: + return _fetch_one( + database, + "SELECT * FROM image_studio_assets WHERE id = ?", + (int(asset_id),), + ImageStudioAsset, + ) + + +def list_assets(project_id, kind=None, include_missing=True, path=None, conn=None): + clauses = ["project_id = ?"] + params = [int(project_id)] + if kind is not None: + clauses.append("kind = ?") + params.append(str(kind)) + if not include_missing: + clauses.append("status != ?") + params.append(ASSET_STATUS_MISSING) + sql = "SELECT * FROM image_studio_assets WHERE " + " AND ".join(clauses) + sql += " ORDER BY source_order, id" + with _connection(conn, path) as database: + return _fetch_all(database, sql, params, ImageStudioAsset) + + +def mark_asset_status(asset_id, status, path=None, conn=None): + if str(status) not in ASSET_STATUSES: + raise db.DbError("AI工场资产状态必须是 available 或 missing") + with _connection(conn, path) as database: + with database: + database.execute( + """ + UPDATE image_studio_assets + SET status = ?, updated_at = ? + WHERE id = ? + """, + (str(status), _now(), int(asset_id)), + ) + return get_asset(asset_id, conn=database) + + +def create_job( + project_id, + *, + source_asset_id=None, + job_type="main", + prompt="", + task_key=None, + generation_source="cmhub", + provider="cmhub", + path=None, + conn=None, +): + now = _now() + task_key = str(task_key or _task_key(project_id)) + with _connection(conn, path) as database: + try: + with database: + _ensure_assets_belong_to_project(database, project_id, [source_asset_id]) + cursor = database.execute( + """ + INSERT INTO image_studio_jobs + (project_id, source_asset_id, generation_source, provider, + job_type, task_key, status, prompt, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?, 'pending', ?, ?, ?) + """, + ( + int(project_id), + source_asset_id, + str(generation_source or "cmhub"), + str(provider or "cmhub"), + str(job_type), + task_key, + str(prompt or ""), + now, + now, + ), + ) + job_id = cursor.lastrowid + except Exception as exc: + raise db.DbError(f"创建AI工场生图任务失败: {exc}") from exc + return get_job(job_id, conn=database) + + +def get_job(job_id, path=None, conn=None): + with _connection(conn, path) as database: + return _fetch_one( + database, + "SELECT * FROM image_studio_jobs WHERE id = ?", + (int(job_id),), + ImageStudioJob, + ) + + +def set_job_submitted(job_id, task_id, *, call_id=None, points_cost=None, points_balance=None, path=None, conn=None): + now = _now() + with _connection(conn, path) as database: + with database: + database.execute( + """ + UPDATE image_studio_jobs + SET task_id = ?, + status = 'submitted', + call_id = ?, + points_cost = ?, + points_balance = ?, + submitted_at = ?, + updated_at = ? + WHERE id = ? + """, + (str(task_id), call_id, points_cost, points_balance, now, now, int(job_id)), + ) + return get_job(job_id, conn=database) + + +def update_job_status( + job_id, + status, + *, + error=None, + output_asset_id=None, + points_balance=None, + increment_attempts=False, + path=None, + conn=None, +): + if str(status) not in JOB_STATUSES: + raise db.DbError("AI工场任务状态无效") + now = _now() + terminal = str(status) in {"succeeded", "failed", "expired", "cancelled"} + with _connection(conn, path) as database: + with database: + if output_asset_id is not None: + job = get_job(job_id, conn=database) + if job is None: + raise db.DbError("AI工场任务不存在") + _ensure_assets_belong_to_project(database, job.project_id, [output_asset_id]) + database.execute( + """ + UPDATE image_studio_jobs + SET status = ?, + error = ?, + output_asset_id = COALESCE(?, output_asset_id), + points_balance = COALESCE(?, points_balance), + attempts = attempts + ?, + finished_at = CASE WHEN ? THEN ? ELSE finished_at END, + updated_at = ? + WHERE id = ? + """, + ( + str(status), + error, + output_asset_id, + points_balance, + 1 if increment_attempts else 0, + 1 if terminal else 0, + now, + now, + int(job_id), + ), + ) + return get_job(job_id, conn=database) + + +def list_resumable_jobs(path=None, conn=None, project_id=None): + clauses = ["status IN (?, ?)", "task_id IS NOT NULL"] + params = ["submitted", "running"] + if project_id is not None: + clauses.append("project_id = ?") + params.append(int(project_id)) + sql = "SELECT * FROM image_studio_jobs WHERE " + " AND ".join(clauses) + sql += " ORDER BY updated_at, id" + with _connection(conn, path) as database: + return _fetch_all(database, sql, params, ImageStudioJob) + + +def replace_selections(project_id, selection_type, asset_ids: Iterable[int], path=None, conn=None): + selection = str(selection_type) + if selection not in SELECTION_TYPES: + raise db.DbError("终选类型必须是 main 或 detail") + ids = [int(asset_id) for asset_id in asset_ids] + if len(ids) != len(set(ids)): + raise db.DbError("同一资产不能重复加入同一类终选") + now = _now() + with _connection(conn, path) as database: + with database: + _ensure_assets_belong_to_project(database, project_id, ids) + database.execute( + """ + DELETE FROM image_studio_selections + WHERE project_id = ? AND selection_type = ? + """, + (int(project_id), selection), + ) + database.executemany( + """ + INSERT INTO image_studio_selections + (project_id, selection_type, position, asset_id, created_at, updated_at) + VALUES (?, ?, ?, ?, ?, ?) + """, + [ + (int(project_id), selection, index + 1, asset_id, now, now) + for index, asset_id in enumerate(ids) + ], + ) + return list_selections(project_id, selection, conn=database) + + +def list_selections(project_id, selection_type=None, path=None, conn=None): + clauses = ["project_id = ?"] + params = [int(project_id)] + if selection_type is not None: + clauses.append("selection_type = ?") + params.append(str(selection_type)) + sql = "SELECT * FROM image_studio_selections WHERE " + " AND ".join(clauses) + sql += " ORDER BY selection_type, position" + with _connection(conn, path) as database: + return _fetch_all(database, sql, params, ImageStudioSelection) diff --git a/docs/tasks/T-586.md b/docs/tasks/T-586.md index 5f23c58..5414238 100644 --- a/docs/tasks/T-586.md +++ b/docs/tasks/T-586.md @@ -3,7 +3,7 @@ id: T-586 title: AI工场数据地基：项目、图片资产、异步任务与终选顺序 SQLite 模型 phase: 7 deps: [T-538] -status: TODO +status: DONE created: 2026-07-11 --- @@ -40,4 +40,8 @@ created: 2026-07-11 ## 执行记录 -（完成后记录实现文件、迁移决策与验证结果。） +- 2026-07-11：完成 AI工场 SQLite 数据地基。 + - `app/db.py` 增加 `image_studio_projects`、`image_studio_assets`、`image_studio_jobs`、`image_studio_selections` 四张表与索引；沿用现有 `CREATE TABLE IF NOT EXISTS` ad-hoc 迁移风格，历史 DB 启动时自动补表，不影响 `batches/accounts/tasks`。 + - 新增 `app/image_studio.py` service 层，封装项目创建/软删除恢复、图片目录计算、资产 CRUD、cmhub job task key/task_id 生命周期、可续查任务列表、主图/详情图终选排序替换；校验 job/终选引用的资产必须属于同一项目。 + - 新增 `tests/test_image_studio.py` 覆盖 schema 初始化、项目唯一性、软删除恢复、路径隔离、资产父子关系、任务状态转换、续查查询、终选连续排序与错误路径。 + - 验证：在只套用 T-586 diff 的 clean worktree 中运行 `python -m ruff check app tests main.py`、`py -3.10 -m compileall app main.py`、`py -3.10 -m unittest discover -s tests`（343 tests）和 `git diff --check`，全部通过。 diff --git a/tests/test_image_studio.py b/tests/test_image_studio.py new file mode 100644 index 0000000..63a0aaf --- /dev/null +++ b/tests/test_image_studio.py @@ -0,0 +1,331 @@ +import os +import sys +import unittest +from types import SimpleNamespace + +sys.path.insert(0, os.path.dirname(__file__)) + +from _helpers import TempDirMixin + +from app import db, image_studio + + +class ImageStudioTests(TempDirMixin, unittest.TestCase): + def test_init_db_adds_image_studio_tables_without_breaking_existing_tables(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + + db.init_db(db_path) + db.init_db(db_path) + + conn = db.connect(db_path) + try: + tables = { + row["name"] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type = 'table'" + ).fetchall() + } + self.assertTrue({"batches", "accounts", "tasks"}.issubset(tables)) + self.assertTrue( + { + "image_studio_projects", + "image_studio_assets", + "image_studio_jobs", + "image_studio_selections", + }.issubset(tables) + ) + + projects_columns = { + row["name"] + for row in conn.execute("PRAGMA table_info(image_studio_projects)").fetchall() + } + self.assertTrue( + { + "account_alias", + "account_slug", + "item_id", + "target_main_count", + "target_detail_count", + "deleted_at", + }.issubset(projects_columns) + ) + finally: + conn.close() + + self.assert_removed(temp_dir) + + def test_project_crud_unique_per_account_and_image_dirs(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + + account = SimpleNamespace( + alias="alias A", + account_name="店铺A", + slug="alias_a_slug", + ) + project = image_studio.create_or_get_project( + account, + item_id="51100639510", + draft_prompt="初始提示词", + path=db_path, + ) + same_project = image_studio.create_or_get_project( + account_alias="alias A", + account_slug="ignored_slug", + item_id="51100639510", + path=db_path, + ) + other_account_project = image_studio.create_or_get_project( + account_alias="alias B", + account_slug="alias_b_slug", + item_id="51100639510", + path=db_path, + ) + + self.assertEqual(project.id, same_project.id) + self.assertNotEqual(project.id, other_account_project.id) + self.assertEqual("alias A", project.account_alias) + self.assertEqual("alias_a_slug", project.account_slug) + self.assertEqual("店铺A", project.account_name) + self.assertEqual("初始提示词", project.draft_prompt) + + updated = image_studio.update_project_prompt(project.id, "二次提示词", path=db_path) + self.assertEqual("二次提示词", updated.draft_prompt) + + dirs = image_studio.project_image_dirs(os.path.join(temp_dir, "images"), project) + self.assertEqual( + os.path.join( + temp_dir, + "images", + "pool", + "alias_a_slug", + "51100639510", + "originals", + ), + dirs["originals"], + ) + self.assertEqual( + os.path.join( + temp_dir, + "images", + "pool", + "alias_a_slug", + "51100639510", + "generated", + ), + dirs["generated"], + ) + self.assertEqual( + os.path.join( + temp_dir, + "images", + "pool", + "alias_a_slug", + "51100639510", + "exports", + ), + dirs["exports"], + ) + + deleted = image_studio.soft_delete_project(project.id, "测试删除", path=db_path) + self.assertEqual("测试删除", deleted.deleted_reason) + self.assertEqual(1, len(image_studio.list_projects(path=db_path))) + restored = image_studio.create_or_get_project( + account, + item_id="51100639510", + path=db_path, + ) + self.assertEqual(project.id, restored.id) + self.assertIsNone(restored.deleted_at) + + self.assert_removed(temp_dir) + + def test_asset_crud_parent_status_and_sorting(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + project = image_studio.create_or_get_project( + account_alias="alias", + account_slug="alias_slug", + item_id="51100639510", + path=db_path, + ) + + original_path = os.path.join(temp_dir, "old.png") + generated_path = os.path.join(temp_dir, "generated.png") + original = image_studio.add_asset( + project.id, + "original", + remote_url="https://example.test/old.png", + local_path=original_path, + aspect_ratio="1:1", + source_order=2, + path=db_path, + ) + generated = image_studio.add_asset( + project.id, + "generated_main", + local_path=generated_path, + parent_asset_id=original.id, + prompt="生成提示词", + source_order=1, + path=db_path, + ) + + self.assertEqual(os.path.abspath(original_path), original.local_path) + self.assertEqual(original.id, generated.parent_asset_id) + self.assertEqual("生成提示词", generated.prompt) + + assets = image_studio.list_assets(project.id, path=db_path) + self.assertEqual([generated.id, original.id], [asset.id for asset in assets]) + + missing = image_studio.mark_asset_status( + original.id, + image_studio.ASSET_STATUS_MISSING, + path=db_path, + ) + self.assertEqual(image_studio.ASSET_STATUS_MISSING, missing.status) + with self.assertRaises(db.DbError): + image_studio.mark_asset_status(original.id, "deleted", path=db_path) + available_assets = image_studio.list_assets( + project.id, + include_missing=False, + path=db_path, + ) + self.assertEqual([generated.id], [asset.id for asset in available_assets]) + + self.assert_removed(temp_dir) + + def test_job_lifecycle_and_resumable_query(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + project = image_studio.create_or_get_project( + account_alias="alias", + account_slug="alias_slug", + item_id="51100639510", + path=db_path, + ) + source_asset = image_studio.add_asset(project.id, "original", path=db_path) + + job = image_studio.create_job( + project.id, + source_asset_id=source_asset.id, + job_type="main", + task_key="stable-task-key", + prompt="生成主图", + path=db_path, + ) + self.assertEqual("stable-task-key", job.task_key) + self.assertEqual("pending", job.status) + self.assertEqual("cmhub", job.provider) + + submitted = image_studio.set_job_submitted( + job.id, + "cmhub-task-1", + call_id="call-1", + points_cost=2, + points_balance=98, + path=db_path, + ) + self.assertEqual("submitted", submitted.status) + self.assertEqual("cmhub-task-1", submitted.task_id) + self.assertEqual(2, submitted.points_cost) + self.assertEqual([job.id], [item.id for item in image_studio.list_resumable_jobs(path=db_path)]) + + running = image_studio.update_job_status( + job.id, + "running", + increment_attempts=True, + path=db_path, + ) + self.assertEqual("running", running.status) + self.assertEqual(1, running.attempts) + with self.assertRaises(db.DbError): + image_studio.update_job_status(job.id, "unknown", path=db_path) + + output_asset = image_studio.add_asset( + project.id, + "generated_main", + parent_asset_id=source_asset.id, + path=db_path, + ) + succeeded = image_studio.update_job_status( + job.id, + "succeeded", + output_asset_id=output_asset.id, + points_balance=96, + path=db_path, + ) + self.assertEqual("succeeded", succeeded.status) + self.assertEqual(output_asset.id, succeeded.output_asset_id) + self.assertEqual(96, succeeded.points_balance) + self.assertIsNotNone(succeeded.finished_at) + self.assertEqual([], image_studio.list_resumable_jobs(path=db_path)) + + with self.assertRaises(db.DbError): + image_studio.create_job( + project.id, + task_key="stable-task-key", + path=db_path, + ) + + self.assert_removed(temp_dir) + + def test_selections_are_consecutive_unique_and_replaceable(self): + with self.make_temp_dir() as temp_dir: + db_path = os.path.join(temp_dir, "cmshopee.db") + db.init_db(db_path) + project = image_studio.create_or_get_project( + account_alias="alias", + account_slug="alias_slug", + item_id="51100639510", + path=db_path, + ) + other_project = image_studio.create_or_get_project( + account_alias="other", + account_slug="other_slug", + item_id="51100639510", + path=db_path, + ) + first = image_studio.add_asset(project.id, "generated_main", source_order=1, path=db_path) + second = image_studio.add_asset(project.id, "generated_main", source_order=2, path=db_path) + third = image_studio.add_asset(project.id, "generated_detail", source_order=3, path=db_path) + other_asset = image_studio.add_asset(other_project.id, "generated_main", path=db_path) + + main = image_studio.replace_selections(project.id, "main", [second.id, first.id], path=db_path) + self.assertEqual([1, 2], [selection.position for selection in main]) + self.assertEqual([second.id, first.id], [selection.asset_id for selection in main]) + + detail = image_studio.replace_selections(project.id, "detail", [second.id, third.id], path=db_path) + self.assertEqual([second.id, third.id], [selection.asset_id for selection in detail]) + + replaced = image_studio.replace_selections(project.id, "main", [first.id], path=db_path) + self.assertEqual([1], [selection.position for selection in replaced]) + self.assertEqual([first.id], [selection.asset_id for selection in replaced]) + + all_selections = image_studio.list_selections(project.id, path=db_path) + self.assertEqual( + [("detail", 1, second.id), ("detail", 2, third.id), ("main", 1, first.id)], + [ + (selection.selection_type, selection.position, selection.asset_id) + for selection in all_selections + ], + ) + + with self.assertRaises(db.DbError): + image_studio.replace_selections(project.id, "main", [first.id, first.id], path=db_path) + with self.assertRaises(db.DbError): + image_studio.replace_selections(project.id, "invalid", [first.id], path=db_path) + with self.assertRaises(db.DbError): + image_studio.replace_selections(project.id, "main", [other_asset.id], path=db_path) + with self.assertRaises(db.DbError): + image_studio.create_job(project.id, source_asset_id=other_asset.id, path=db_path) + + self.assert_removed(temp_dir) + + +if __name__ == "__main__": + unittest.main()