yvjia/src/main/java/com/backendsys/modules/sdk/comfyui/utils/ComfyUtil.java

package com.backendsys.modules.sdk.comfyui.utils;

import com.backendsys.modules.common.config.security.utils.SecurityUtil;
import com.backendsys.modules.sdk.comfyui.dao.ComfyuiTaskDao;
import com.backendsys.modules.sdk.comfyui.entity.ComfyuiTask;
import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.util.concurrent.CompletableFuture;

@Component
public class ComfyUtil {

    @Autowired
    private ComfyuiTaskDao comfyuiTaskDao;

    /**
     * [新增] 生成任务的记录 (异步)
     */
    public void asyncInsertTaskRecord(String prompt_id, String prompt) {
        Long user_id = SecurityUtil.getUserId();
        CompletableFuture.runAsync(() -> {
            // [DB] 新增任务记录
            try {
                System.out.println("-- 新增任务记录 --");
                ComfyuiTask comfyuiTask = new ComfyuiTask();
                comfyuiTask.setUser_id(user_id);
                comfyuiTask.setPrompt_id(prompt_id);
                comfyuiTask.setGenerate_request(prompt);
                comfyuiTaskDao.insert(comfyuiTask);
            } catch (Exception e) {
                System.out.println(e.getMessage());
            }
        });
    }
    /**
     * [更新] 生成任务的记录 (异步)
     */
    public void asyncUpdateTaskRecord(String prompt_id, String response, Integer status) {
        CompletableFuture.runAsync(() -> {
            // [DB] 新增任务记录
            try {
                System.out.println("-- 更新任务记录 --");
                LambdaQueryWrapper<ComfyuiTask> wrapper = new LambdaQueryWrapper<>();
                wrapper.eq(ComfyuiTask::getPrompt_id, prompt_id);

                ComfyuiTask comfyuiTask = new ComfyuiTask();
                comfyuiTask.setStatus(status);
                comfyuiTask.setGenerate_response(response);
                comfyuiTaskDao.update(comfyuiTask, wrapper);
            } catch (Exception e) {
                System.out.println(e.getMessage());
            }
        });
    }

}