feat(ai-studio): add image studio persistence layer
This commit is contained in:
chengma
@@ -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);
|
||||
"""
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
+6
-2
@@ -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`，全部通过。
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user