feat(product-suite): make image pulls cancellable

app/image_studio.py

@@ -116,6 +116,15 @@ class ImageStudioProjectConflictError(ImageStudioError):
     """Raised when a draft cannot be bound because the formal project already exists."""
 
 
+class ImageStudioPullCancelled(ImageStudioError):
+    """Raised when a read-only Shopee image pull stops at a safe boundary."""
+
+    def __init__(self, project=None, assets=None):
+        super().__init__("用户停止拉取蝦皮主图")
+        self.project = project
+        self.assets = list(assets or [])
+
+
 def _now() -> str:
     return datetime.now().isoformat(timespec="seconds")
 
@@ -825,6 +834,93 @@ def remove_original_assets_if_unused(project_id, asset_ids, path=None, conn=None
         return [_row_to_dataclass(by_id[asset_id], ImageStudioAsset) for asset_id in ordered_ids]
 
 
+def restore_original_asset_snapshot(project_id, snapshots, path=None, conn=None):
+    """Restore status/order for original assets that existed before a pull."""
+
+    project_id = int(project_id)
+    normalized = []
+    seen = set()
+    for value in snapshots or []:
+        asset_id = int(_get(value, "id"))
+        if asset_id in seen:
+            continue
+        status = str(_get(value, "status") or ASSET_STATUS_AVAILABLE)
+        if status not in ASSET_STATUSES:
+            raise db.DbError("商品原图快照状态无效")
+        normalized.append(
+            {
+                "id": asset_id,
+                "status": status,
+                "source_order": max(0, int(_get(value, "source_order") or 0)),
+            }
+        )
+        seen.add(asset_id)
+    if not normalized:
+        return []
+    ids = [item["id"] for item in normalized]
+    placeholders = ",".join("?" for _ in ids)
+    with _connection(conn, path) as database:
+        with database:
+            rows = database.execute(
+                f"""
+                SELECT id FROM image_studio_assets
+                WHERE project_id = ? AND kind = ? AND id IN ({placeholders})
+                """,
+                [project_id, ASSET_KIND_ORIGINAL, *ids],
+            ).fetchall()
+            if {int(row["id"]) for row in rows} != set(ids):
+                raise db.DbError("拉取前商品原图快照已失效")
+            now = _now()
+            for item in normalized:
+                database.execute(
+                    """
+                    UPDATE image_studio_assets
+                    SET status = ?, source_order = ?, updated_at = ?
+                    WHERE id = ? AND project_id = ? AND kind = ?
+                    """,
+                    (
+                        item["status"],
+                        item["source_order"],
+                        now,
+                        item["id"],
+                        project_id,
+                        ASSET_KIND_ORIGINAL,
+                    ),
+                )
+            remaining_rows = database.execute(
+                f"""
+                SELECT id FROM image_studio_assets
+                WHERE project_id = ? AND kind = ?
+                  AND id NOT IN ({placeholders})
+                ORDER BY source_order, id
+                """,
+                [project_id, ASSET_KIND_ORIGINAL, *ids],
+            ).fetchall()
+            next_order = max(
+                [int(item["source_order"]) for item in normalized] + [0]
+            ) + 1
+            for row in remaining_rows:
+                database.execute(
+                    """
+                    UPDATE image_studio_assets
+                    SET source_order = ?, updated_at = ?
+                    WHERE id = ? AND project_id = ? AND kind = ?
+                    """,
+                    (
+                        next_order,
+                        now,
+                        int(row["id"]),
+                        project_id,
+                        ASSET_KIND_ORIGINAL,
+                    ),
+                )
+                next_order += 1
+        return [
+            get_asset(item["id"], conn=database)
+            for item in normalized
+        ]
+
+
 def sync_original_asset_urls(project_id, image_urls, path=None, conn=None, max_assets=16):
     """Store the read-only Shopee main image URL snapshot as remote-only assets."""
 
@@ -1183,10 +1279,14 @@ def pull_remote_main_image_urls(
     config=None,
     login_timeout=8,
     on_step=None,
+    should_stop=None,
 ):
     """Create/open an AI studio project and read Shopee main image URLs without editing."""
 
     cfg = appconfig.load_config() if config is None else config
+    should_stop = should_stop or (lambda: False)
+    if should_stop():
+        raise ImageStudioPullCancelled()
     database_path = _db_path(path, cfg)
     db.init_db(database_path)
     item = _normalize_item_id(item_id)
@@ -1196,6 +1296,8 @@ def pull_remote_main_image_urls(
         raise ImageStudioError(f"AI工场账号不可用：{exc}") from exc
 
     project = create_or_get_project(account, item_id=item, path=database_path)
+    if should_stop():
+        raise ImageStudioPullCancelled(project=project)
     readiness = _ensure_account_ready_for_read(
         account,
         path=database_path,
@@ -1203,16 +1305,26 @@ def pull_remote_main_image_urls(
         login_timeout=login_timeout,
         on_step=on_step,
     )
+    if should_stop():
+        raise ImageStudioPullCancelled(project=project)
     cdp = None
     try:
+        if should_stop():
+            raise ImageStudioPullCancelled(project=project)
         _notify_step(on_step, "open_product", "start", f"商品 {item}")
         cdp = editor.open_product(account, item, on_step=on_step, bring_to_front=False)
+        if should_stop():
+            raise ImageStudioPullCancelled(project=project)
         _notify_step(on_step, "open_product", "success", f"商品 {item}")
         _notify_step(on_step, "read_main_images", "start", f"商品 {item}")
         images = editor.read_product_image_urls(cdp)
+        if should_stop():
+            raise ImageStudioPullCancelled(project=project)
         if not images:
             raise ImageStudioError("未读取到蝦皮商品主图 URL")
         assets = sync_original_asset_urls(project.id, images, path=database_path)
+        if should_stop():
+            raise ImageStudioPullCancelled(project=project, assets=assets)
         _notify_step(on_step, "read_main_images", "success", f"读取 {len(images)} 张主图 URL")
         project = get_project(project.id, path=database_path)
         return {