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)
|
||||
Reference in New Issue
Block a user