EasyFlow-Tech-Dev/EasyFlow
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 完善数据中枢联邦查询闭环

Browse Source 
- 重构数据源生命周期、元数据纳管与运行时切换

- 增加只读 SQL、查询审计、跨节点取消与工作流联动

- 完善管理端连接配置、元数据浏览与 SQL 工作台
This commit is contained in:
陈子默
2026-08-26 18:14:34 +08:00
parent 3e79e99925
commit e38821e48a
113 changed files with 17032 additions and 1151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/cancellation/WorkflowDatasetQueryCancellationListener.java Normal file
View File

@@ -0,0 +1,42 @@
package tech.easyflow.ai.easyagentsflow.cancellation;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.Event;
import com.easyagents.flow.core.chain.event.ChainStatusChangeEvent;
import com.easyagents.flow.core.chain.listener.ChainEventListener;
import org.springframework.stereotype.Component;
/**
* 将工作流取消终态桥接到活动数据集查询。
*/
@Component
public class WorkflowDatasetQueryCancellationListener
implements ChainEventListener {
private final WorkflowDatasetQueryCancellationRegistry registry;
/**
* 创建工作流查询取消监听器。
*
* @param registry 工作流查询取消登记表
*/
public WorkflowDatasetQueryCancellationListener(
WorkflowDatasetQueryCancellationRegistry registry) {
this.registry = registry;
}
/**
* 在工作流进入取消终态后取消该实例的活动查询。
*
* @param event 工作流事件
* @param chain 工作流实例
*/
@Override
public void onEvent(Event event, Chain chain) {
if (event instanceof ChainStatusChangeEvent statusChangeEvent
&& statusChangeEvent.getStatus() == ChainStatus.CANCELLED) {
registry.cancelExecution(chain.getStateInstanceId());
}
}
}

279
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/cancellation/WorkflowDatasetQueryCancellationRegistry.java Normal file
View File

@@ -0,0 +1,279 @@
package tech.easyflow.ai.easyagentsflow.cancellation;
import java.math.BigInteger;
import java.time.Duration;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Objects;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.datacenter.federation.DatacenterFederationQueryCancellationService;
/**
* 维护工作流实例到活动数据集查询的跨节点取消映射。
*/
@Component
public class WorkflowDatasetQueryCancellationRegistry {
private static final String ACTIVE_KEY_PREFIX =
"easyflow:workflow:dataset-query:active:";
private static final String CANCELLED_KEY_PREFIX =
"easyflow:workflow:dataset-query:cancelled:";
private static final Duration STATE_TTL = Duration.ofHours(24);
private static final Logger log = LoggerFactory.getLogger(
WorkflowDatasetQueryCancellationRegistry.class);
private final DatacenterFederationQueryCancellationService cancellationService;
private final ObjectProvider<StringRedisTemplate> redisTemplateProvider;
private final ConcurrentHashMap<String, ConcurrentHashMap<String, BigInteger>>
localActive = new ConcurrentHashMap<>();
private final ConcurrentHashMap<String, Long> localCancelledUntil =
new ConcurrentHashMap<>();
/**
* 创建工作流查询取消登记表。
*
* @param cancellationService 数据中枢查询取消服务
* @param redisTemplateProvider 可选 Redis 模板
*/
public WorkflowDatasetQueryCancellationRegistry(
DatacenterFederationQueryCancellationService cancellationService,
ObjectProvider<StringRedisTemplate> redisTemplateProvider) {
this.cancellationService = cancellationService;
this.redisTemplateProvider = redisTemplateProvider;
}
/**
* 在数据库查询开始前登记工作流与 QueryId。
*
* @param stateInstanceId 工作流实例 ID
* @param account 执行账号
* @param queryId 查询 UUID
* @return 必须关闭的登记句柄
* @throws BusinessException 参数缺失时抛出
*/
public Registration register(
String stateInstanceId,
LoginAccount account,
String queryId) {
String instanceId = requireText(stateInstanceId, "工作流实例 ID 不能为空");
String normalizedQueryId = requireText(queryId, "queryId 不能为空");
BigInteger tenantId = requireTenantId(account);
localActive.computeIfAbsent(
instanceId,
ignored -> new ConcurrentHashMap<>())
.put(normalizedQueryId, tenantId);
persistActive(instanceId, normalizedQueryId, tenantId);
Registration registration = new Registration(
instanceId, normalizedQueryId, tenantId);
if (isExecutionCancelled(instanceId)) {
cancelQuery(normalizedQueryId, tenantId);
}
return registration;
}
/**
* 取消工作流当前登记的全部数据集查询。
*
* @param stateInstanceId 工作流实例 ID
* @return 是否发现至少一个活动查询
*/
public boolean cancelExecution(String stateInstanceId) {
String instanceId = requireText(stateInstanceId, "工作流实例 ID 不能为空");
long expiresAt = System.currentTimeMillis() + STATE_TTL.toMillis();
localCancelledUntil.put(instanceId, expiresAt);
persistCancellationMarker(instanceId);
Map<String, BigInteger> active = new LinkedHashMap<>();
ConcurrentHashMap<String, BigInteger> local = localActive.get(instanceId);
if (local != null) {
active.putAll(local);
}
loadPersistedActive(instanceId).forEach(active::putIfAbsent);
active.forEach((queryId, tenantId) -> {
try {
cancelQuery(queryId, tenantId);
} catch (RuntimeException exception) {
// 单个驱动取消失败不能阻断同一工作流的其他活动查询。
log.error("Failed to cancel workflow dataset query {} for {}",
queryId, instanceId, exception);
}
});
return !active.isEmpty();
}
/**
* 清理本机已过期的取消墓碑，避免不可达工作流持续占用内存。
*/
@Scheduled(fixedDelayString =
"${easyflow.workflow.dataset-query-cancel-cleanup-ms:60000}")
public void cleanupExpiredLocalMarkers() {
long now = System.currentTimeMillis();
localCancelledUntil.entrySet().removeIf(
entry -> entry.getValue() <= now);
}
private void cancelQuery(String queryId, BigInteger tenantId) {
LoginAccount account = new LoginAccount();
account.setTenantId(tenantId);
cancellationService.cancel(queryId, account);
}
private boolean isExecutionCancelled(String stateInstanceId) {
Long localDeadline = localCancelledUntil.get(stateInstanceId);
if (localDeadline != null && localDeadline > System.currentTimeMillis()) {
return true;
}
StringRedisTemplate redisTemplate = redisTemplateProvider.getIfAvailable();
if (redisTemplate == null) {
return false;
}
try {
return Boolean.TRUE.equals(redisTemplate.hasKey(
cancelledKey(stateInstanceId)));
} catch (RuntimeException exception) {
log.warn("Failed to read workflow dataset cancellation marker for {}",
stateInstanceId, exception);
return false;
}
}
private void persistActive(
String stateInstanceId,
String queryId,
BigInteger tenantId) {
StringRedisTemplate redisTemplate = redisTemplateProvider.getIfAvailable();
if (redisTemplate == null) {
return;
}
try {
String key = activeKey(stateInstanceId);
redisTemplate.opsForHash().put(key, queryId, tenantId.toString());
redisTemplate.expire(key, STATE_TTL);
} catch (RuntimeException exception) {
log.warn("Failed to persist workflow dataset query mapping for {}",
stateInstanceId, exception);
}
}
private Map<String, BigInteger> loadPersistedActive(String stateInstanceId) {
StringRedisTemplate redisTemplate = redisTemplateProvider.getIfAvailable();
if (redisTemplate == null) {
return Map.of();
}
try {
Map<Object, Object> entries = redisTemplate.opsForHash()
.entries(activeKey(stateInstanceId));
Map<String, BigInteger> active = new LinkedHashMap<>();
entries.forEach((queryId, tenantId) -> {
try {
active.put(String.valueOf(queryId),
new BigInteger(String.valueOf(tenantId)));
} catch (RuntimeException exception) {
log.warn("Ignored invalid workflow dataset query mapping for {}",
stateInstanceId, exception);
}
});
return active;
} catch (RuntimeException exception) {
log.warn("Failed to load workflow dataset query mappings for {}",
stateInstanceId, exception);
return Map.of();
}
}
private void persistCancellationMarker(String stateInstanceId) {
StringRedisTemplate redisTemplate = redisTemplateProvider.getIfAvailable();
if (redisTemplate == null) {
return;
}
try {
redisTemplate.opsForValue().set(
cancelledKey(stateInstanceId), "1", STATE_TTL);
} catch (RuntimeException exception) {
log.warn("Failed to persist workflow dataset cancellation marker for {}",
stateInstanceId, exception);
}
}
private void unregister(
String stateInstanceId,
String queryId,
BigInteger tenantId) {
localActive.computeIfPresent(stateInstanceId, (ignored, queries) -> {
queries.remove(queryId, tenantId);
return queries.isEmpty() ? null : queries;
});
StringRedisTemplate redisTemplate = redisTemplateProvider.getIfAvailable();
if (redisTemplate == null) {
return;
}
try {
redisTemplate.opsForHash().delete(
activeKey(stateInstanceId), queryId);
} catch (RuntimeException exception) {
log.warn("Failed to remove workflow dataset query mapping for {}",
stateInstanceId, exception);
}
}
private String activeKey(String stateInstanceId) {
return ACTIVE_KEY_PREFIX + stateInstanceId;
}
private String cancelledKey(String stateInstanceId) {
return CANCELLED_KEY_PREFIX + stateInstanceId;
}
private String requireText(String value, String message) {
if (value == null || value.isBlank()) {
throw new BusinessException(message);
}
return value.trim();
}
private BigInteger requireTenantId(LoginAccount account) {
if (account == null || account.getTenantId() == null) {
throw new BusinessException("工作流数据集查询缺少执行租户");
}
return account.getTenantId();
}
/**
* 单次工作流查询登记句柄。
*/
public final class Registration implements AutoCloseable {
private final String stateInstanceId;
private final String queryId;
private final BigInteger tenantId;
private final AtomicBoolean closed = new AtomicBoolean();
private Registration(
String stateInstanceId,
String queryId,
BigInteger tenantId) {
this.stateInstanceId = stateInstanceId;
this.queryId = queryId;
this.tenantId = tenantId;
}
/**
* 幂等注销当前工作流查询。
*/
@Override
public void close() {
if (closed.compareAndSet(false, true)) {
unregister(stateInstanceId, queryId, tenantId);
}
}
}
}

7
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/config/ChainExecutorConfig.java
View File

@@ -14,6 +14,7 @@ import org.springframework.context.annotation.Configuration;
import tech.easyflow.ai.easyagentsflow.listener.ChainErrorListenerForSave;
import tech.easyflow.ai.easyagentsflow.listener.ChainEventListenerForSave;
import tech.easyflow.ai.easyagentsflow.listener.NodeErrorListenerForSave;
import tech.easyflow.ai.easyagentsflow.cancellation.WorkflowDatasetQueryCancellationListener;
import tech.easyflow.ai.easyagentsflow.upload.WorkflowApiUploadCleanupListener;
import javax.annotation.Resource;
@@ -40,6 +41,9 @@ public class ChainExecutorConfig {
@Resource
private WorkflowApiUploadCleanupListener workflowApiUploadCleanupListener;
@Resource
private WorkflowDatasetQueryCancellationListener
workflowDatasetQueryCancellationListener;
@Resource
private WorkflowExecutionBudgetProperties workflowExecutionBudgetProperties;
@Resource
private WorkflowRuntimeProperties workflowRuntimeProperties;
@@ -91,6 +95,9 @@ public class ChainExecutorConfig {
chainExecutor.addEventListener(
ChainStatusChangeEvent.class,
workflowApiUploadCleanupListener);
chainExecutor.addEventListener(
ChainStatusChangeEvent.class,
workflowDatasetQueryCancellationListener);
chainExecutor.addErrorListener(new ChainErrorListenerForSave());
chainExecutor.addNodeErrorListener(new NodeErrorListenerForSave());
}

117
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/easyagentsflow/service/WorkflowDatacenterContentService.java
View File

@@ -14,12 +14,13 @@ import tech.easyflow.datacenter.execution.model.DatasetRef;
import tech.easyflow.datacenter.execution.service.DatacenterDatasetQueryService;
import tech.easyflow.datacenter.meta.entity.DatacenterCatalog;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.datacenter.meta.enums.DatacenterMetadataStatus;
import tech.easyflow.datacenter.meta.enums.DatacenterSensitivityLevel;
import tech.easyflow.datacenter.meta.service.DatacenterDatasetRegistryService;
import javax.annotation.Resource;
import java.math.BigInteger;
import java.util.ArrayList;
import java.util.Comparator;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
@@ -34,6 +35,7 @@ public class WorkflowDatacenterContentService {
public static final String LLM_NODE_TYPE = "llmNode";
public static final String QUERY_DATA_CONTEXT = "queryDataContext";
public static final String SEARCH_SOURCE_MISSING_MESSAGE = "查询数据节点未选择连接服务";
public static final String SEARCH_TABLE_MISSING_MESSAGE = "查询数据节点未选择已接入表";
public static final String SEARCH_SQL_MISSING_MESSAGE = "查询数据节点未设置 SQL";
public static final String SAVE_EXPIRED_MESSAGE = "写入数据节点配置已过期，请重新选择已接入表";
public static final String INVALID_QUERY_CONTEXT_MESSAGE = "查询上下文配置无效，请重新选择查询数据节点";
@@ -135,11 +137,16 @@ public class WorkflowDatacenterContentService {
if (datasetRef == null || datasetRef.getSourceId() == null) {
throw new BusinessException(SEARCH_SOURCE_MISSING_MESSAGE);
}
if (datasetRef.getTableId() == null) {
throw new BusinessException(SEARCH_TABLE_MISSING_MESSAGE);
}
String querySql = data == null ? null : trimToNull(data.getString("querySql"));
if (!StringUtils.hasText(querySql)) {
throw new BusinessException(SEARCH_SQL_MISSING_MESSAGE);
}
return datasetRef;
DatasetRef boundRef = bindAuthoritativeTenant(datasetRef);
data.put("datasetRef", boundRef);
return boundRef;
}
public DatasetRef requireSaveDatasetRef(JSONObject data) {
@@ -150,6 +157,31 @@ public class WorkflowDatacenterContentService {
if (datasetRef == null || datasetRef.getTableId() == null) {
throw new BusinessException(SAVE_EXPIRED_MESSAGE);
}
DatasetRef boundRef = bindAuthoritativeTenant(datasetRef);
data.put("datasetRef", boundRef);
return boundRef;
}
/**
* 依据当前租户可见的权威 Source/Table 覆盖调用方声明的租户字段。
*
* @param datasetRef 工作流数据集引用
* @return 已绑定权威租户的数据集引用
*/
private DatasetRef bindAuthoritativeTenant(DatasetRef datasetRef) {
DatacenterTable table = datasetRef.getTableId() == null
? null : registryService.getTableWithFields(datasetRef.getTableId());
BigInteger sourceId = table == null ? datasetRef.getSourceId() : table.getSourceId();
if (sourceId == null) {
throw new BusinessException(SEARCH_SOURCE_MISSING_MESSAGE);
}
DatacenterSource source = registryService.getSourceRequired(sourceId);
if (table != null && (!sourceId.equals(table.getSourceId())
|| !java.util.Objects.equals(source.getTenantId(), table.getTenantId()))) {
throw new BusinessException("数据集引用与当前租户不一致");
}
datasetRef.setTenantId(source.getTenantId());
datasetRef.setSourceId(sourceId);
return datasetRef;
}
@@ -172,41 +204,68 @@ public class WorkflowDatacenterContentService {
if (datasetRef == null || datasetRef.getSourceId() == null) {
throw new BusinessException(SEARCH_SOURCE_MISSING_MESSAGE);
}
if (datasetRef.getTableId() == null) {
throw new BusinessException(SEARCH_TABLE_MISSING_MESSAGE);
}
DatacenterSource source = registryService.getSourceRequired(datasetRef.getSourceId());
List<DatacenterTable> managedTables = registryService.listManagedTables(datasetRef.getSourceId(), datasetRef.getCatalogId());
managedTables.sort(Comparator.comparing(table -> table.getTableName() == null ? "" : table.getTableName()));
DatacenterTable fullTable = registryService.getTableWithFields(
datasetRef.getTableId());
if (fullTable == null
|| !datasetRef.getSourceId().equals(fullTable.getSourceId())
|| !Integer.valueOf(1).equals(fullTable.getQueryable())
|| !DatacenterMetadataStatus.ACTIVE.name().equals(
fullTable.getMetadataStatus())) {
throw new BusinessException("查询数据节点绑定的表不可用");
}
DatacenterCatalog catalog = registryService.getCatalogById(
fullTable.getCatalogId());
if (StringUtils.hasText(datasetRef.getCatalogName())
&& (catalog == null || !datasetRef.getCatalogName().equals(
catalog.getCatalogName()))) {
throw new BusinessException("查询数据节点绑定的命名空间已变化");
}
JSONObject sourceSummary = new JSONObject();
sourceSummary.put("sourceName", source.getSourceName());
sourceSummary.put("sourceType", source.getSourceType());
JSONArray tables = new JSONArray();
for (DatacenterTable table : managedTables) {
DatacenterTable fullTable = registryService.getTableWithFields(table.getId());
DatacenterCatalog catalog = registryService.getCatalogById(fullTable.getCatalogId());
if (StringUtils.hasText(datasetRef.getCatalogName())
&& (catalog == null || !datasetRef.getCatalogName().equals(catalog.getCatalogName()))) {
continue;
}
JSONObject tableSummary = new JSONObject();
tableSummary.put("catalogName", catalog == null ? null : catalog.getCatalogName());
tableSummary.put("tableName", fullTable.getTableName());
tableSummary.put("tableDesc", fullTable.getTableDesc());
JSONArray fields = new JSONArray();
if (fullTable.getFields() != null) {
for (DatacenterTableField field : fullTable.getFields()) {
JSONObject fieldSummary = new JSONObject();
fieldSummary.put("fieldName", field.getFieldName());
fieldSummary.put("fieldDesc", field.getFieldDesc());
fieldSummary.put("fieldType", resolveFieldType(field));
fields.add(fieldSummary);
JSONObject tableSummary = new JSONObject();
tableSummary.put("catalogName", catalog == null ? null : catalog.getCatalogName());
tableSummary.put("tableName", fullTable.getTableName());
tableSummary.put("tableDesc", fullTable.getTableDesc());
JSONArray fields = new JSONArray();
if (fullTable.getFields() != null) {
for (DatacenterTableField field : fullTable.getFields()) {
if (!isQueryablePublicField(field)) {
continue;
}
JSONObject fieldSummary = new JSONObject();
fieldSummary.put("fieldName", field.getFieldName());
fieldSummary.put("fieldDesc", field.getFieldDesc());
fieldSummary.put("fieldType", resolveFieldType(field));
fields.add(fieldSummary);
}
tableSummary.put("fields", fields);
tables.add(tableSummary);
}
tableSummary.put("fields", fields);
tables.add(tableSummary);
sourceSummary.put("tables", tables);
return sourceSummary;
}
/**
* 判断字段是否允许暴露给 SQL 生成上下文。
*
* @param field 字段元数据
* @return 是否为当前可查询的公开字段
*/
private boolean isQueryablePublicField(DatacenterTableField field) {
return field != null
&& Integer.valueOf(1).equals(field.getQueryable())
&& DatacenterMetadataStatus.ACTIVE.name().equals(
field.getMetadataStatus())
&& DatacenterSensitivityLevel.PUBLIC.name().equals(
field.getSensitivityLevel());
}
private void injectQueryDataContext(JSONObject data, Map<String, JSONObject> nodeMap) {
if (data == null) {
return;
@@ -217,7 +276,7 @@ public class WorkflowDatacenterContentService {
removeQueryDataContextPlaceholder(data);
return;
}
Map<BigInteger, JSONObject> sourceSummaries = new LinkedHashMap<>();
Map<String, JSONObject> sourceSummaries = new LinkedHashMap<>();
Set<String> visitedNodeIds = new LinkedHashSet<>();
for (int i = 0; i < nodeIds.size(); i++) {
String nodeId = trimToNull(nodeIds.getString(i));
@@ -229,7 +288,10 @@ public class WorkflowDatacenterContentService {
throw new BusinessException(INVALID_QUERY_CONTEXT_MESSAGE);
}
DatasetRef datasetRef = requireSearchDatasetRef(targetNode.getJSONObject("data"));
sourceSummaries.putIfAbsent(datasetRef.getSourceId(), buildSourceSummary(datasetRef));
String summaryKey = datasetRef.getSourceId() + ":" + datasetRef.getTableId();
if (!sourceSummaries.containsKey(summaryKey)) {
sourceSummaries.put(summaryKey, buildSourceSummary(datasetRef));
}
}
String contextValue = QUERY_CONTEXT_PROMPT + "\n" + JSON.toJSONString(new ArrayList<>(sourceSummaries.values()));
upsertQueryDataContextParameter(data, contextValue);
@@ -343,6 +405,7 @@ public class WorkflowDatacenterContentService {
private DatasetRef copyDatasetRef(DatasetRef datasetRef) {
DatasetRef copy = new DatasetRef();
copy.setTenantId(datasetRef.getTenantId());
copy.setSourceId(datasetRef.getSourceId());
copy.setCatalogId(datasetRef.getCatalogId());
copy.setCatalogName(datasetRef.getCatalogName());

4
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/node/SaveDatasetNode.java
View File

@@ -6,7 +6,6 @@ import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.runtime.RetryableTriggerException;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.util.IoBulkhead;
import com.mybatisflex.core.tenant.TenantManager;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import tech.easyflow.ai.utils.WorkFlowUtil;
@@ -55,7 +54,6 @@ public class SaveDatasetNode extends BaseNode {
rows.add(item instanceof JSONObject json ? json : JSONObject.from(item));
}
try {
TenantManager.ignoreTenantCondition();
try (IoBulkhead.Permit ignored = IoBulkhead.dataset().acquire(resolveIoTarget())) {
writeService.saveRowsIdempotently(
datasetRef,
@@ -77,8 +75,6 @@ public class SaveDatasetNode extends BaseNode {
} catch (Exception ex) {
log.error("工作流保存数据到统一数据集失败，datasetRef={}", datasetRef, ex);
throw ex;
} finally {
TenantManager.restoreTenantCondition();
}
}

67
easyflow-modules/easyflow-module-ai/src/main/java/tech/easyflow/ai/node/SearchDatasetNode.java
View File

@@ -7,12 +7,13 @@ import com.easyagents.flow.core.chain.repository.LoopInputReference;
import com.easyagents.flow.core.node.BaseNode;
import com.easyagents.flow.core.util.IoBulkhead;
import com.mybatisflex.core.row.Row;
import com.mybatisflex.core.tenant.TenantManager;
import tech.easyflow.common.util.SpringContextUtil;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.common.web.exceptions.BusinessException;
import tech.easyflow.datacenter.execution.model.DatacenterSqlQueryRequest;
import tech.easyflow.datacenter.execution.model.DatasetRef;
import tech.easyflow.datacenter.execution.service.DatacenterDatasetQueryService;
import tech.easyflow.ai.easyagentsflow.cancellation.WorkflowDatasetQueryCancellationRegistry;
import java.util.HashMap;
import java.util.List;
@@ -55,32 +56,38 @@ public class SearchDatasetNode extends BaseNode {
Map<String, Object> params =
chain.getExecutionState().resolveParameters(this);
DatacenterDatasetQueryService queryService = SpringContextUtil.getBean(DatacenterDatasetQueryService.class);
WorkflowDatasetQueryCancellationRegistry cancellationRegistry =
SpringContextUtil.getBean(
WorkflowDatasetQueryCancellationRegistry.class);
LoginAccount account = tech.easyflow.ai.utils.WorkFlowUtil.getOperator(chain);
DatacenterSqlQueryRequest request = buildRuntimeRequest(params);
Map<String, Object> result = new HashMap<>();
try {
TenantManager.ignoreTenantCondition();
try (IoBulkhead.Permit ignored = IoBulkhead.dataset().acquire(resolveIoTarget())) {
String resultId = chain.getStateInstanceId()
+ ":dataset:"
+ UUID.randomUUID();
int rowCount =
chain.storeProducedLoopInputOutsideLock(
resultId,
sink -> queryService.consumeBySql(
request,
QUERY_PAGE_SIZE,
sink::accept),
0L,
chain.currentFencingClaimId(),
chain.currentClaimGeneration());
result.put(
resolveOutputKey("data"),
new LoopInputReference(
resultId, rowCount));
return result;
}
} finally {
TenantManager.restoreTenantCondition();
String queryId = UUID.randomUUID().toString();
try (WorkflowDatasetQueryCancellationRegistry.Registration registration =
cancellationRegistry.register(
chain.getStateInstanceId(), account, queryId);
IoBulkhead.Permit ignored = IoBulkhead.dataset()
.acquire(resolveIoTarget())) {
String resultId = chain.getStateInstanceId()
+ ":dataset:"
+ queryId;
int rowCount =
chain.storeProducedLoopInputOutsideLock(
resultId,
sink -> queryService.consumeBySql(
request,
QUERY_PAGE_SIZE,
account,
queryId,
sink::accept),
0L,
chain.currentFencingClaimId(),
chain.currentClaimGeneration());
result.put(
resolveOutputKey("data"),
new LoopInputReference(
resultId, rowCount));
return result;
}
}
@@ -97,6 +104,9 @@ public class SearchDatasetNode extends BaseNode {
}
private DatacenterSqlQueryRequest buildRuntimeRequest(Map<String, Object> params) {
if (datasetRef == null || datasetRef.getSourceId() == null) {
throw new BusinessException("数据集绑定缺少连接信息，请重新选择数据集");
}
DatacenterSqlQueryRequest request = new DatacenterSqlQueryRequest();
request.setDatasetRef(copyDatasetRef());
request.setSql(resolveQuerySql(params));
@@ -128,12 +138,13 @@ public class SearchDatasetNode extends BaseNode {
private DatasetRef copyDatasetRef() {
DatasetRef copy = new DatasetRef();
copy.setTenantId(datasetRef == null ? null : datasetRef.getTenantId());
copy.setSourceId(datasetRef == null ? null : datasetRef.getSourceId());
copy.setCatalogId(datasetRef == null ? null : datasetRef.getCatalogId());
copy.setCatalogName(datasetRef == null ? null : datasetRef.getCatalogName());
copy.setTableId(null);
copy.setTableName(null);
copy.setVersionId(null);
copy.setTableId(datasetRef == null ? null : datasetRef.getTableId());
copy.setTableName(datasetRef == null ? null : datasetRef.getTableName());
copy.setVersionId(datasetRef == null ? null : datasetRef.getVersionId());
return copy;
}

131
easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/easyagentsflow/cancellation/WorkflowDatasetQueryCancellationRegistryTest.java Normal file
View File

@@ -0,0 +1,131 @@
package tech.easyflow.ai.easyagentsflow.cancellation;
import com.easyagents.flow.core.chain.Chain;
import com.easyagents.flow.core.chain.ChainStatus;
import com.easyagents.flow.core.chain.event.ChainStatusChangeEvent;
import java.math.BigInteger;
import java.util.UUID;
import org.junit.Assert;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.mockito.Mockito;
import org.springframework.beans.factory.ObjectProvider;
import org.springframework.data.redis.core.StringRedisTemplate;
import tech.easyflow.common.entity.LoginAccount;
import tech.easyflow.datacenter.federation.DatacenterFederationQueryCancellationService;
/**
* {@link WorkflowDatasetQueryCancellationRegistry} 取消竞态回归测试。
*/
public class WorkflowDatasetQueryCancellationRegistryTest {
/**
* 验证工作流取消会使用登记租户取消活动 QueryId。
*/
@Test
public void shouldCancelRegisteredQuery() {
Fixture fixture = fixture();
String queryId = UUID.randomUUID().toString();
LoginAccount account = account(BigInteger.valueOf(7L));
try (WorkflowDatasetQueryCancellationRegistry.Registration ignored =
fixture.registry.register("instance-a", account, queryId)) {
Assert.assertTrue(fixture.registry.cancelExecution("instance-a"));
}
ArgumentCaptor<LoginAccount> accountCaptor =
ArgumentCaptor.forClass(LoginAccount.class);
Mockito.verify(fixture.cancellationService).cancel(
Mockito.eq(queryId), accountCaptor.capture());
Assert.assertEquals(BigInteger.valueOf(7L),
accountCaptor.getValue().getTenantId());
}
/**
* 验证单个查询取消失败不会阻断同一工作流的其他活动查询。
*/
@Test
public void shouldIsolateCancellationFailuresBetweenQueries() {
Fixture fixture = fixture();
String failedQueryId = UUID.randomUUID().toString();
String succeedingQueryId = UUID.randomUUID().toString();
LoginAccount account = account(BigInteger.valueOf(7L));
Mockito.doThrow(new IllegalStateException("simulated cancellation failure"))
.when(fixture.cancellationService)
.cancel(Mockito.eq(failedQueryId), Mockito.any(LoginAccount.class));
try (WorkflowDatasetQueryCancellationRegistry.Registration ignoredA =
fixture.registry.register("instance-failure", account, failedQueryId);
WorkflowDatasetQueryCancellationRegistry.Registration ignoredB =
fixture.registry.register("instance-failure", account, succeedingQueryId)) {
Assert.assertTrue(fixture.registry.cancelExecution("instance-failure"));
}
Mockito.verify(fixture.cancellationService).cancel(
Mockito.eq(failedQueryId), Mockito.any(LoginAccount.class));
Mockito.verify(fixture.cancellationService).cancel(
Mockito.eq(succeedingQueryId), Mockito.any(LoginAccount.class));
}
/**
* 验证取消先于节点登记发生时，后续查询仍会立即收到取消。
*/
@Test
public void shouldCancelQueryRegisteredAfterWorkflowCancellation() {
Fixture fixture = fixture();
String queryId = UUID.randomUUID().toString();
Assert.assertFalse(fixture.registry.cancelExecution("instance-b"));
try (WorkflowDatasetQueryCancellationRegistry.Registration ignored =
fixture.registry.register(
"instance-b", account(BigInteger.ONE), queryId)) {
Mockito.verify(fixture.cancellationService).cancel(
Mockito.eq(queryId), Mockito.any(LoginAccount.class));
}
}
/**
* 验证监听器只桥接取消终态。
*/
@Test
public void listenerShouldBridgeOnlyCancelledStatus() {
WorkflowDatasetQueryCancellationRegistry registry = Mockito.mock(
WorkflowDatasetQueryCancellationRegistry.class);
WorkflowDatasetQueryCancellationListener listener =
new WorkflowDatasetQueryCancellationListener(registry);
Chain chain = Mockito.mock(Chain.class);
Mockito.when(chain.getStateInstanceId()).thenReturn("instance-c");
listener.onEvent(new ChainStatusChangeEvent(
chain, ChainStatus.SUCCEEDED, ChainStatus.RUNNING), chain);
listener.onEvent(new ChainStatusChangeEvent(
chain, ChainStatus.CANCELLED, ChainStatus.RUNNING), chain);
Mockito.verify(registry).cancelExecution("instance-c");
}
@SuppressWarnings("unchecked")
private Fixture fixture() {
DatacenterFederationQueryCancellationService cancellationService =
Mockito.mock(
DatacenterFederationQueryCancellationService.class);
ObjectProvider<StringRedisTemplate> redisProvider =
Mockito.mock(ObjectProvider.class);
Mockito.when(redisProvider.getIfAvailable()).thenReturn(null);
return new Fixture(
new WorkflowDatasetQueryCancellationRegistry(
cancellationService, redisProvider),
cancellationService);
}
private LoginAccount account(BigInteger tenantId) {
LoginAccount account = new LoginAccount();
account.setTenantId(tenantId);
return account;
}
private record Fixture(
WorkflowDatasetQueryCancellationRegistry registry,
DatacenterFederationQueryCancellationService cancellationService) {
}
}

20
easyflow-modules/easyflow-module-ai/src/test/java/tech/easyflow/ai/easyagentsflow/service/WorkflowDatacenterContentServiceTest.java
View File

@@ -9,6 +9,8 @@ import org.mockito.Mockito;
import tech.easyflow.datacenter.entity.DatacenterTable;
import tech.easyflow.datacenter.entity.DatacenterTableField;
import tech.easyflow.datacenter.meta.entity.DatacenterSource;
import tech.easyflow.datacenter.meta.enums.DatacenterMetadataStatus;
import tech.easyflow.datacenter.meta.enums.DatacenterSensitivityLevel;
import tech.easyflow.datacenter.meta.service.DatacenterDatasetRegistryService;
import java.lang.reflect.Field;
@@ -41,13 +43,22 @@ public class WorkflowDatacenterContentServiceTest {
DatacenterSource source = Mockito.mock(DatacenterSource.class);
Mockito.when(source.getSourceName()).thenReturn("ama 实验基线模型预算");
Mockito.when(source.getSourceType()).thenReturn("EXCEL");
Mockito.when(source.getTenantId()).thenReturn(BigInteger.ONE);
DatacenterTableField modelId = mockField("col_id", "模型ID", "VARCHAR");
DatacenterTableField inputPrice = mockField("token", "输入价格", "DECIMAL");
DatacenterTableField hidden = mockField("hidden_col", "受限字段", "VARCHAR");
Mockito.when(hidden.getQueryable()).thenReturn(0);
DatacenterTable table = Mockito.mock(DatacenterTable.class);
Mockito.when(table.getId()).thenReturn(TABLE_ID);
Mockito.when(table.getSourceId()).thenReturn(SOURCE_ID);
Mockito.when(table.getTenantId()).thenReturn(BigInteger.ONE);
Mockito.when(table.getTableName()).thenReturn("Sheet1");
Mockito.when(table.getFields()).thenReturn(List.of(modelId, inputPrice));
Mockito.when(table.getQueryable()).thenReturn(1);
Mockito.when(table.getMetadataStatus()).thenReturn(
DatacenterMetadataStatus.ACTIVE.name());
Mockito.when(table.getFields()).thenReturn(List.of(
modelId, inputPrice, hidden));
Mockito.when(registryService.getSourceRequired(SOURCE_ID)).thenReturn(source);
Mockito.when(registryService.listManagedTables(SOURCE_ID, null))
@@ -77,6 +88,7 @@ public class WorkflowDatacenterContentServiceTest {
Assert.assertTrue(contextValue.contains("col_id"));
Assert.assertTrue(contextValue.contains("模型ID"));
Assert.assertTrue(contextValue.contains("token AS input_price"));
Assert.assertFalse(contextValue.contains("hidden_col"));
}
/**
@@ -103,6 +115,7 @@ public class WorkflowDatacenterContentServiceTest {
private JSONObject buildWorkflowRoot() {
JSONObject datasetRef = new JSONObject();
datasetRef.put("sourceId", SOURCE_ID);
datasetRef.put("tableId", TABLE_ID);
JSONObject queryData = new JSONObject();
queryData.put("datasetRef", datasetRef);
@@ -160,6 +173,11 @@ public class WorkflowDatacenterContentServiceTest {
Mockito.when(field.getFieldName()).thenReturn(fieldName);
Mockito.when(field.getFieldDesc()).thenReturn(fieldDesc);
Mockito.when(field.getJdbcType()).thenReturn(jdbcType);
Mockito.when(field.getQueryable()).thenReturn(1);
Mockito.when(field.getMetadataStatus()).thenReturn(
DatacenterMetadataStatus.ACTIVE.name());
Mockito.when(field.getSensitivityLevel()).thenReturn(
DatacenterSensitivityLevel.PUBLIC.name());
return field;
}