From a144b315c6904764c51899843bf957a33a85eb4d Mon Sep 17 00:00:00 2001 From: fenyf Date: Mon, 21 Jul 2025 18:01:51 +0800 Subject: [PATCH 01/10] feat(ai_recognition): implement initial AI-based sensitive column rule feat(ai_recognition): Implement rule priority configuration using the strategy pattern feat(ai_recognition): Implement the process of sending sensitive columns in batches to the AI feat(ai_recognition): Implement the process of sending sensitive columns in batches to the AI feat(ai_recognition): Implement the process of sending sensitive columns in batches to the AI(completely restructured version) --- client | 2 +- pom.xml | 11 + .../V_4_3_4_20__alter_sensitive_rule.sql | 10 + server/odc-service/pom.xml | 8 + .../datasecurity/SensitiveRuleEntity.java | 9 + .../SensitiveColumnRecognizer.java | 72 ------ .../datasecurity/SensitiveColumnScanner.java | 79 +++++++ .../SensitiveColumnScanningTask.java | 101 ++++++-- .../SensitiveColumnScanningTaskManager.java | 6 +- .../datasecurity/SensitiveColumnService.java | 2 +- .../odc/service/datasecurity/ai/AIConfig.java | 65 ++++++ .../datasecurity/ai/AIInferenceService.java | 56 +++++ .../datasecurity/ai/PromptTemplateLoader.java | 151 ++++++++++++ .../ColumnRecognizerFactory.java | 15 +- .../factory/ScanningStrategyFactory.java | 93 ++++++++ .../datasecurity/model/RecognitionResult.java | 35 +++ .../datasecurity/model/ScanResult.java | 70 ++++++ .../datasecurity/model/ScanningModeType.java | 31 +++ .../model/SensitiveColumnScanningReq.java | 2 + .../datasecurity/model/SensitiveRule.java | 6 + .../datasecurity/model/SensitiveRuleType.java | 7 +- .../recognizer/AIColumnRecognizer.java | 171 ++++++++++++++ .../recognizer/ColumnRecognizer.java | 44 +++- .../recognizer/GroovyColumnRecognizer.java | 30 ++- .../recognizer/PathColumnRecognizer.java | 30 ++- .../recognizer/RegexColumnRecognizer.java | 39 +++- .../strategy/AbstractScanningStrategy.java | 109 +++++++++ .../strategy/JointRecognitionStrategy.java | 89 +++++++ .../strategy/RulesAndAiStrategy.java | 72 ++++++ .../strategy/RulesOnlyStrategy.java | 59 +++++ .../strategy/ScanningStrategy.java | 55 +++++ ...itive_column_recognize_prompt_templete.txt | 46 ++++ .../SensitiveColumnScannerTest.java | 202 ++++++++++++++++ .../recognizer/AIColumnRecognizerTest.java | 218 ++++++++++++++++++ .../GroovyColumnRecognizerTest.java | 112 ++++++--- .../recognizer/PathColumnRecognizerTest.java | 58 ++++- .../recognizer/RegexColumnRecognizerTest.java | 55 ++++- 37 files changed, 2036 insertions(+), 184 deletions(-) create mode 100644 server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql delete mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnRecognizer.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java rename server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/{ => factory}/ColumnRecognizerFactory.java (77%) create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java create mode 100644 server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java diff --git a/client b/client index 273fa3c4cc..5b4c9fbf18 160000 --- a/client +++ b/client @@ -1 +1 @@ -Subproject commit 273fa3c4cc7d87f942ade231ce9220b98fa595e2 +Subproject commit 5b4c9fbf18cde5f7c2b3455eef434d96c7ddf771 diff --git a/pom.xml b/pom.xml index 42262c105b..54505aea16 100644 --- a/pom.xml +++ b/pom.xml @@ -196,6 +196,17 @@ ${springboot.version} + + + com.openai + openai-java + 2.16.0 + + + org.jetbrains.kotlin + kotlin-stdlib-jdk7 + 1.9.24 + org.codehaus.groovy diff --git a/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql new file mode 100644 index 0000000000..5a1fd21059 --- /dev/null +++ b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql @@ -0,0 +1,10 @@ +-- Add AI-related columns to table `data_security_sensitive_rule` +alter table `data_security_sensitive_rule` + add column `ai_sensitive_types` text default null comment 'A list of sensitive data types for AI rules, stored as a JSON array string.'; + +alter table `data_security_sensitive_rule` + add column `ai_confidence_threshold` integer default 80 comment 'Confidence threshold for AI rules, with a value range of 0-100.'; + +alter table `data_security_sensitive_rule` + add column `ai_custom_prompt` text default null comment 'User-defined custom prompt for AI rules.'; + diff --git a/server/odc-service/pom.xml b/server/odc-service/pom.xml index 5d28d6077c..e811f237ef 100644 --- a/server/odc-service/pom.xml +++ b/server/odc-service/pom.xml @@ -75,6 +75,14 @@ commons-beanutils commons-beanutils + + com.openai + openai-java + + + org.jetbrains.kotlin + kotlin-stdlib-jdk7 + org.apache.commons commons-compress diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java b/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java index 405452ab90..8dedbf4e12 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java @@ -113,4 +113,13 @@ public class SensitiveRuleEntity { @Column(name = "update_time", nullable = false, insertable = false, updatable = false) private Date updateTime; + @Convert(converter = JsonListConverter.class) + @Column(name = "ai_sensitive_types") + private List aiSensitiveTypes; + + @Column(name = "ai_confidence_threshold") + private Integer aiConfidenceThreshold; + + @Column(name = "ai_custom_prompt") + private String aiCustomPrompt; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnRecognizer.java deleted file mode 100644 index 6e7c4a979b..0000000000 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnRecognizer.java +++ /dev/null @@ -1,72 +0,0 @@ -/* - * Copyright (c) 2023 OceanBase. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.oceanbase.odc.service.datasecurity; - -import java.util.ArrayList; -import java.util.List; - -import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; -import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; -import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - -/** - * @author gaoda.xy - * @date 2023/5/30 10:30 - */ -public class SensitiveColumnRecognizer implements ColumnRecognizer { - - private Long sensitiveRuleId; - private Long maskingAlgorithmId; - private SensitiveLevel sensitiveLevel; - private final List sensitiveRules; - private final List recognizers; - - public SensitiveColumnRecognizer(List rules) { - this.sensitiveRules = rules; - this.recognizers = new ArrayList<>(); - for (SensitiveRule rule : this.sensitiveRules) { - this.recognizers.add(ColumnRecognizerFactory.create(rule)); - } - } - - public Long sensitiveRuleId() { - return this.sensitiveRuleId; - } - - public Long maskingAlgorithmId() { - return this.maskingAlgorithmId; - } - - public SensitiveLevel sensitiveLevel() { - return this.sensitiveLevel; - } - - @Override - public boolean recognize(DBTableColumn column) { - for (int i = 0; i < recognizers.size(); i++) { - if (recognizers.get(i).recognize(column)) { - SensitiveRule rule = sensitiveRules.get(i); - this.sensitiveRuleId = rule.getId(); - this.maskingAlgorithmId = rule.getMaskingAlgorithmId(); - this.sensitiveLevel = rule.getLevel(); - return true; - } - } - return false; - } - -} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java new file mode 100644 index 0000000000..8cc51db195 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java @@ -0,0 +1,79 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity; + +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +import com.oceanbase.odc.service.datasecurity.factory.ColumnRecognizerFactory; +import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.odc.service.datasecurity.strategy.ScanningStrategy; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 敏感列识别的编排器，负责根据不同的扫描模式执行识别策略。 + * 使用策略模式重构，消除重复代码。 + */ +public class SensitiveColumnScanner { + + private final List basicRecognizers; + private final List aiRecognizers; + private final ScanningStrategyFactory strategyFactory; + + public SensitiveColumnScanner(List rules, ScanningStrategyFactory strategyFactory) { + // 在构造时，就将规则分好类，并创建对应的识别器 + this.basicRecognizers = rules.stream() + .filter(r -> r.getType() != SensitiveRuleType.AI) + .map(ColumnRecognizerFactory::create) + .collect(Collectors.toList()); + this.aiRecognizers = rules.stream() + .filter(r -> r.getType() == SensitiveRuleType.AI) + .map(ColumnRecognizerFactory::create) + .collect(Collectors.toList()); + this.strategyFactory = strategyFactory; + } + + /** + * 核心扫描方法 + * + * @param column 待扫描的列 + * @param mode 用户选择的扫描模式 + * @return 包含一个或两个结果的最终扫描报告 + */ + public ScanResult scan(DBTableColumn column, ScanningModeType mode) { + ScanningStrategy strategy = strategyFactory.getStrategy(mode); + return strategy.scan(column, basicRecognizers, aiRecognizers); + } + + /** + * 批量扫描方法 + * + * @param columns 待扫描的列列表 + * @param mode 用户选择的扫描模式 + * @return 扫描结果映射，key为列名，value为扫描结果 + */ + public Map scanBatch(List columns, ScanningModeType mode) { + ScanningStrategy strategy = strategyFactory.getStrategy(mode); + return strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java index 00d3ccc390..0220d35d7f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java @@ -20,11 +20,18 @@ import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; +import java.util.function.Function; +import java.util.stream.Collectors; import com.oceanbase.odc.core.shared.constant.ErrorCodes; import com.oceanbase.odc.service.connection.database.model.Database; +import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumn; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnMeta; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo; @@ -40,56 +47,88 @@ public class SensitiveColumnScanningTask implements Callable { private final Database database; - private final SensitiveColumnRecognizer recognizer; + private final SensitiveColumnScanner scanner; + private final ScanningModeType scanningMode; private final SensitiveColumnScanningTaskInfo taskInfo; private final Map> table2Columns; private final Map> view2Columns; private final Set existsSensitiveColumns; + private final Map ruleMap; - public SensitiveColumnScanningTask(Database database, List rules, + public SensitiveColumnScanningTask(Database database, List rules, ScanningModeType scanningMode, SensitiveColumnScanningTaskInfo taskInfo, List existsSensitiveColumns, Map> table2Columns, Map> view2Columns) { this.database = database; - this.recognizer = new SensitiveColumnRecognizer(rules); + // 【修改】接收扫描模式，并创建新的扫描器 + this.scanningMode = scanningMode; + ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); + this.scanner = new SensitiveColumnScanner(rules, strategyFactory); this.table2Columns = table2Columns; this.view2Columns = view2Columns; this.taskInfo = taskInfo; this.existsSensitiveColumns = new HashSet<>(existsSensitiveColumns); + // 【修改】将规则列表转换为 Map，方便通过 ID 快速查找 + this.ruleMap = rules.stream().collect(Collectors.toMap(SensitiveRule::getId, Function.identity())); + } + + /** + * 生成列的唯一标识符 + */ + private String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } @Override - public Void call() throws Exception { + public Void call() { try { taskInfo.setStatus(ScanningTaskStatus.RUNNING); + // 调用重构后的 scanColumns 方法 scanColumns(table2Columns, SensitiveColumnType.TABLE_COLUMN); scanColumns(view2Columns, SensitiveColumnType.VIEW_COLUMN); + taskInfo.setStatus(ScanningTaskStatus.SUCCESS); } catch (Exception e) { - taskInfo.setCompleteTime(new Date()); taskInfo.setStatus(ScanningTaskStatus.FAILED); taskInfo.setErrorCode(ErrorCodes.Unexpected); - taskInfo.setErrorMsg(String.format("Some errors happen when scanning sensitive column, database=%s", - database.getName())); + taskInfo.setErrorMsg(String.format("Error during sensitive column scanning on database=%s, reason=%s", + database.getName(), e.getMessage())); + } finally { + taskInfo.setCompleteTime(new Date()); } return null; } + // 【修改】scanColumns 的核心逻辑改为批量扫描 private void scanColumns(Map> object2Columns, SensitiveColumnType columnType) { - for (String objectName : object2Columns.keySet()) { + for (Map.Entry> entry : object2Columns.entrySet()) { + String objectName = entry.getKey(); + List columns = entry.getValue(); + + // 【改为批量扫描】一次性扫描整个表的所有列 + Map scanResults = this.scanner.scanBatch(columns, this.scanningMode); + List sensitiveColumns = new ArrayList<>(); - for (DBTableColumn dbTableColumn : object2Columns.get(objectName)) { - if (recognizer.recognize(dbTableColumn) && !existsSensitiveColumns - .contains(new SensitiveColumnMeta(database.getId(), objectName, dbTableColumn.getName()))) { - SensitiveColumn column = new SensitiveColumn(); - column.setType(columnType); - column.setDatabase(database); - column.setTableName(objectName); - column.setColumnName(dbTableColumn.getName()); - column.setMaskingAlgorithmId(recognizer.maskingAlgorithmId()); - column.setSensitiveRuleId(recognizer.sensitiveRuleId()); - column.setLevel(recognizer.sensitiveLevel()); - sensitiveColumns.add(column); - existsSensitiveColumns - .add(new SensitiveColumnMeta(database.getId(), objectName, dbTableColumn.getName())); + for (DBTableColumn dbTableColumn : columns) { + String columnKey = getColumnKey(dbTableColumn); + ScanResult scanResult = scanResults.get(columnKey); + + if (scanResult != null) { + // 根据扫描模式获取最终的识别结果 + Optional finalResultOpt = scanResult.getFinalResult(this.scanningMode); + + // 如果最终有识别结果，则处理 + finalResultOpt.ifPresent(finalResult -> { + SensitiveColumnMeta meta = new SensitiveColumnMeta(database.getId(), objectName, + dbTableColumn.getName()); + if (!existsSensitiveColumns.contains(meta)) { + SensitiveColumn column = createSensitiveColumn(columnType, objectName, dbTableColumn, + finalResult); + sensitiveColumns.add(column); + existsSensitiveColumns.add(meta); + } + }); } } taskInfo.addSensitiveColumns(sensitiveColumns); @@ -97,4 +136,22 @@ private void scanColumns(Map> object2Columns, Sensit } } + // 【新增】辅助方法，用于创建 SensitiveColumn 对象，使代码更清晰 + private SensitiveColumn createSensitiveColumn(SensitiveColumnType columnType, String objectName, + DBTableColumn dbTableColumn, RecognitionResult result) { + SensitiveColumn column = new SensitiveColumn(); + column.setType(columnType); + column.setDatabase(database); + column.setTableName(objectName); + column.setColumnName(dbTableColumn.getName()); + // 从 RecognitionResult 获取 ruleId 和 level + column.setSensitiveRuleId(result.getMatchedRuleId()); + column.setLevel(result.getLevel()); + // 通过 ruleId 从我们保存的 ruleMap 中找到对应的规则，再获取脱敏算法ID + SensitiveRule matchedRule = this.ruleMap.get(result.getMatchedRuleId()); + if (matchedRule != null) { + column.setMaskingAlgorithmId(matchedRule.getMaskingAlgorithmId()); + } + return column; + } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java index 72a93c1676..1dc3543692 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java @@ -37,6 +37,7 @@ import com.oceanbase.odc.core.shared.constant.ResourceType; import com.oceanbase.odc.service.connection.database.model.Database; import com.oceanbase.odc.service.connection.model.ConnectionConfig; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnMeta; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo.ScanningTaskStatus; @@ -65,6 +66,7 @@ public class SensitiveColumnScanningTaskManager { private StatefulUuidStateIdGenerator statefulUuidStateIdGenerator; public SensitiveColumnScanningTaskInfo start(List databases, List rules, + ScanningModeType scanningMode, // 新增参数 ConnectionConfig connectionConfig, Map> databaseId2SensitiveColumns) { ConnectionSession session = new DefaultConnectSessionFactory(connectionConfig).generateSession(); try { @@ -101,8 +103,8 @@ public SensitiveColumnScanningTaskInfo start(List databases, List()), + SensitiveColumnScanningTask subTask = new SensitiveColumnScanningTask(database, rules, scanningMode, + taskInfo, sensitiveColumns, database2Table2ColumnsList.getOrDefault(database, new HashMap<>()), database2View2ColumnsList.getOrDefault(database, new HashMap<>())); try { executor.submit(subTask); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java index 6569b068d9..f7129a8fd7 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java @@ -399,7 +399,7 @@ public SensitiveColumnScanningTaskInfo startScanning(@NotNull Long projectId, PreConditions.notEmpty(rules, "sensitiveRules"); ConnectionConfig connectionConfig = databaseService.findDataSourceForConnectById(databases.get(0).getId()); Map> databaseId2SensitiveColumns = listExistSensitiveColumns(databaseIds); - return scanningTaskManager.start(databases, rules, connectionConfig, databaseId2SensitiveColumns); + return scanningTaskManager.start(databases, rules, req.getScanningMode(), connectionConfig, databaseId2SensitiveColumns); } @Transactional(rollbackFor = Exception.class) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java new file mode 100644 index 0000000000..28898c754c --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java @@ -0,0 +1,65 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import org.springframework.context.annotation.Bean; +import org.springframework.stereotype.Component; + +import com.openai.client.OpenAIClient; +import com.openai.client.okhttp.OpenAIOkHttpClient; + +import lombok.Data; + +/** + * 后期需要弄成动态配置 + */ +@Data +@Component +// @ConfigurationProperties(prefix = "datasecurity.ai") +public class AIConfig { + /** + * 是否启用AI服务 + */ + private boolean enabled = true;; + + /** + * API 密钥 + */ + //private String apiKey = "sk-c6bbbbde1b7e420b897d0662301c6d7c"; + private String apiKey = "token-abc123"; + + /** + * API 的基础URL + */ + //private String baseUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1"; + private String baseUrl = "http://172.25.17.78:8000/v1"; + + /** + * 默认使用的模型名称 + */ + //private String model = "qwen2.5-3b-instruct"; + private String model = "nlora"; + + @Bean + public OpenAIClient openAIClient() { + return OpenAIOkHttpClient.builder() + .apiKey(this.apiKey) + .baseUrl(this.baseUrl) + .build(); + } + + +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java new file mode 100644 index 0000000000..98342d582f --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import java.util.HashMap; +import java.util.Map; + +import org.springframework.stereotype.Service; + +import com.openai.client.OpenAIClient; +import com.openai.core.JsonBoolean; +import com.openai.core.JsonValue; +import com.openai.models.chat.completions.ChatCompletion; +import com.openai.models.chat.completions.ChatCompletionCreateParams; + +@Service +public class AIInferenceService { + + private final AIConfig aiConfig; + private final OpenAIClient openAIClient; + + public AIInferenceService(AIConfig aiConfig) { + this.aiConfig = aiConfig; + // 根据AIConfig创建OpenAIClient + this.openAIClient = aiConfig.openAIClient(); + } + + public ChatCompletion chat(String prompt) { + //Map bodyParams = new HashMap<>(); + //bodyParams.put("enable_thinking", JsonBoolean.from(false)); + String model = aiConfig.getModel(); + try { + ChatCompletionCreateParams params = ChatCompletionCreateParams.builder() + .addUserMessage(prompt) + .model(model) + //.additionalBodyProperties(bodyParams) + .build(); + return openAIClient.chat().completions().create(params); + } catch (Exception e) { + throw new RuntimeException("调用阿里云AI服务失败: " + e.getMessage(), e); + } + } +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java new file mode 100644 index 0000000000..d904de9220 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java @@ -0,0 +1,151 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import java.util.List; +import java.util.Objects; +import java.util.stream.Collectors; + +import javax.annotation.PostConstruct; + +import org.springframework.stereotype.Component; + +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +import lombok.var; + +/** + * AI 提示词 (Prompt) 模板加载器和构建器（重构版）。 + *

+ * 该类负责加载结构化的 AI 提示词模板，并根据列元数据、指定的敏感类型和用户自定义提示来构建最终的提示词。 + *

+ */ +@Component +public class PromptTemplateLoader { + + private static final String TEMPLATE_PATH = "/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt"; + + // 定义新的三个占位符 + private static final String COLUMN_PLACEHOLDER = "{DBTableColumn}"; + private static final String TYPES_PLACEHOLDER = "{sensitiveTypes}"; + private static final String PROMPT_PLACEHOLDER = "{customPrompt}"; + + private String template; + + @PostConstruct + public void init() { + try (var inputStream = PromptTemplateLoader.class.getResourceAsStream(TEMPLATE_PATH)) { + if (Objects.isNull(inputStream)) { + throw new IllegalStateException("AI prompt template file not found: " + TEMPLATE_PATH); + } + try (var reader = new java.io.BufferedReader(new java.io.InputStreamReader(inputStream))) { + this.template = reader.lines().collect(Collectors.joining(System.lineSeparator())); + } + } catch (Exception e) { + // 在实际项目中，这里应该使用日志系统 + e.printStackTrace(); + throw new IllegalStateException("Failed to load AI prompt template", e); + } + } + + /** + * 【新】根据列元数据、敏感类型列表和用户提示，构建最终的 AI 提示词。 + * + * @param column 数据库表列的元数据对象 + * @param sensitiveTypes 用户指定的敏感类型列表 (例如 ["联系方式", "身份信息"]) + * @param customPrompt 用户为该规则自定义的补充说明提示 + * @return 填充了所有信息的完整提示词字符串 + */ + public String buildPrompt(DBTableColumn column, List sensitiveTypes, String customPrompt) { + if (this.template == null || this.template.isEmpty()) { + throw new IllegalStateException("Prompt template is not available. Check loading status."); + } + if (column == null) { + throw new IllegalArgumentException("Input column cannot be null."); + } + + // 1. 格式化列元数据 + String formattedColumn = formatColumnMetadata(column); + + // 2. 格式化敏感类型列表 + String formattedTypes = (sensitiveTypes == null || sensitiveTypes.isEmpty()) + ? "None specified" + : String.join(", ", sensitiveTypes); + + // 3. 格式化用户自定义提示 + String formattedPrompt = (customPrompt == null || customPrompt.trim().isEmpty()) + ? "None" + : customPrompt; + + // 4. 依次替换模板中的三个占位符 + return this.template + .replace(COLUMN_PLACEHOLDER, formattedColumn) + .replace(TYPES_PLACEHOLDER, formattedTypes) + .replace(PROMPT_PLACEHOLDER, formattedPrompt); + } + + /** + * 根据列元数据列表、敏感类型列表和用户提示，构建批量处理的 AI 提示词 + * + * @param columnsJson 数据库表列的元数据对象列表的JSON字符串 + * @param sensitiveTypes 用户指定的敏感类型列表 (例如 ["联系方式", "身份信息"]) + * @param customPrompt 用户为该规则自定义的补充说明提示 + * @return 填充了所有信息的完整提示词字符串 + */ + public String buildPrompt(String columnsJson, List sensitiveTypes, String customPrompt) { + if (this.template == null || this.template.isEmpty()) { + throw new IllegalStateException("Prompt template is not available. Check loading status."); + } + if (columnsJson == null) { + throw new IllegalArgumentException("Input columnsJson cannot be null."); + } + + // 1. 格式化敏感类型列表 + String formattedTypes = (sensitiveTypes == null || sensitiveTypes.isEmpty()) + ? "None specified" + : String.join(", ", sensitiveTypes); + + // 2. 格式化用户自定义提示 + String formattedPrompt = (customPrompt == null || customPrompt.trim().isEmpty()) + ? "None" + : customPrompt; + + // 3. 替换模板中的占位符 + return this.template + .replace(COLUMN_PLACEHOLDER, columnsJson) + .replace(TYPES_PLACEHOLDER, formattedTypes) + .replace(PROMPT_PLACEHOLDER, formattedPrompt); + } + + /** + * 将列的元数据格式化为对 AI 模型友好的字符串。 + * (此方法逻辑不变) + */ + private String formatColumnMetadata(DBTableColumn column) { + StringBuilder contextBuilder = new StringBuilder(); + contextBuilder.append("Schema Name: ").append(formatValue(column.getSchemaName())).append("\n"); + contextBuilder.append("Table Name: ").append(formatValue(column.getTableName())).append("\n"); + contextBuilder.append("Column Name: ").append(formatValue(column.getName())).append("\n"); + contextBuilder.append("Data Type: ").append(formatValue(column.getTypeName())).append("\n"); + contextBuilder.append("Column Comment: ").append(formatValue(column.getComment())); + return contextBuilder.toString(); + } + + private String formatValue(String value) { + return (value == null || value.trim().isEmpty()) ? "N/A" : value; + } + +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ColumnRecognizerFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java similarity index 77% rename from server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ColumnRecognizerFactory.java rename to server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java index 50cfc35c7a..8f67a4a4d3 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ColumnRecognizerFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023 OceanBase. + * Copyright (c) 2025 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,11 +13,12 @@ * See the License for the specific language governing permissions and * limitations under the License. */ -package com.oceanbase.odc.service.datasecurity; +package com.oceanbase.odc.service.datasecurity.factory; import com.oceanbase.odc.core.shared.constant.ErrorCodes; import com.oceanbase.odc.core.shared.exception.UnsupportedException; import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.recognizer.AIColumnRecognizer; import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; import com.oceanbase.odc.service.datasecurity.recognizer.GroovyColumnRecognizer; import com.oceanbase.odc.service.datasecurity.recognizer.PathColumnRecognizer; @@ -34,16 +35,16 @@ public class ColumnRecognizerFactory { public static ColumnRecognizer create(@NonNull SensitiveRule rule) { switch (rule.getType()) { case REGEX: - return new RegexColumnRecognizer(rule.getDatabaseRegexExpression(), rule.getTableRegexExpression(), - rule.getColumnRegexExpression(), rule.getColumnCommentRegexExpression()); + return new RegexColumnRecognizer(rule); case PATH: - return new PathColumnRecognizer(rule.getPathIncludes(), rule.getPathExcludes()); + return new PathColumnRecognizer(rule); case GROOVY: - return new GroovyColumnRecognizer(rule.getGroovyScript()); + return new GroovyColumnRecognizer(rule); + case AI: + return new AIColumnRecognizer(rule); default: String errorMsg = String.format("Unsupported sensitive rule type: %s", rule.getType().name()); throw new UnsupportedException(ErrorCodes.BadArgument, new Object[] {errorMsg}, errorMsg); } } - } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java new file mode 100644 index 0000000000..5bf9aed38d --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java @@ -0,0 +1,93 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.factory; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import org.springframework.stereotype.Component; + +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; +import com.oceanbase.odc.service.datasecurity.strategy.JointRecognitionStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.RulesAndAiStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.RulesOnlyStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.ScanningStrategy; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 扫描策略工厂类 + * 根据扫描模式类型返回对应的策略实例 + * + * @author Assistant + * @date 2025/1/27 + */ +@Component +public class ScanningStrategyFactory { + + private final Map strategies = new HashMap<>(); + + public ScanningStrategyFactory() { + // 预创建所有策略实例 + strategies.put(ScanningModeType.RULES_ONLY, new RulesOnlyStrategy()); + strategies.put(ScanningModeType.JOINT_RECOGNITION, new JointRecognitionStrategy()); + strategies.put(ScanningModeType.RULES_AND_AI, new RulesAndAiStrategy()); + } + + /** + * 根据扫描模式获取对应的策略 + * + * @param mode 扫描模式 + * @return 对应的策略实例 + */ + public ScanningStrategy getStrategy(ScanningModeType mode) { + ScanningStrategy strategy = strategies.get(mode); + if (strategy == null) { + // 返回默认的无操作策略 + return new NoOpStrategy(); + } + return strategy; + } + + /** + * 无操作策略实现，用于处理未知或不支持的扫描模式 + */ + private static class NoOpStrategy implements ScanningStrategy { + @Override + public ScanResult scan(DBTableColumn column, + java.util.List basicRecognizers, + java.util.List aiRecognizers) { + return new ScanResult(Optional.empty(), Optional.empty()); + } + + @Override + public Map scanBatch(java.util.List columns, + java.util.List basicRecognizers, + java.util.List aiRecognizers) { + Map results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + results.put(columnKey, new ScanResult(Optional.empty(), Optional.empty())); + } + return results; + } + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java new file mode 100644 index 0000000000..17c1dbdff8 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.model; + +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import lombok.Builder; +import lombok.Data; + +@Data +@Builder +public class RecognitionResult { + // 基础信息 + private boolean matched; + private Long matchedRuleId; + private SensitiveLevel level; + private SensitiveRuleType sourceRuleType; + + // AI 规则 + private String sensitiveType; // AI 判断出的具体敏感类型 + private Integer confidence; // AI 的置信度 +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java new file mode 100644 index 0000000000..44078f34fb --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java @@ -0,0 +1,70 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.model; // 建议放在 model 包下 + +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import lombok.AllArgsConstructor; +import lombok.Getter; + +@Getter +@AllArgsConstructor +public class ScanResult { + // 基础规则的识别结果 (如果有) + private final Optional basicRuleResult; + // AI 规则的识别结果 (如果有) + private final Optional aiRuleResult; + + /** + * 根据扫描模式获取最终的识别结果 + * + * @param scanningMode 扫描模式 + * @return 最终的识别结果 + */ + public Optional getFinalResult(ScanningModeType scanningMode) { + switch (scanningMode) { + case RULES_ONLY: + return basicRuleResult; + case JOINT_RECOGNITION: + // 对于联合识别，Scanner已经做过决策，直接返回存在的那个结果 + return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; + case RULES_AND_AI: + // 对于差异化展示模式，可以根据业务需求调整优先级策略 + return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; + default: + return Optional.empty(); + } + } + + /** + * 判断是否有任何识别结果 + */ + public boolean hasAnyResult() { + return basicRuleResult.isPresent() || aiRuleResult.isPresent(); + } + + /** + * 获取所有可用的结果（用于差异化展示场景） + */ + public List getAllResults() { + List results = new ArrayList<>(); + basicRuleResult.ifPresent(results::add); + aiRuleResult.ifPresent(results::add); + return results; + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java new file mode 100644 index 0000000000..3ce04a4be6 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java @@ -0,0 +1,31 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.model; + +/** + * @author fenyf + * @date 2025/7/18 17:52 + */ +public enum ScanningModeType { + + RULES_ONLY, + + RULES_AND_AI, + + JOINT_RECOGNITION, + + +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningReq.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningReq.java index fd77c7f642..ba36f1c087 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningReq.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningReq.java @@ -35,4 +35,6 @@ public class SensitiveColumnScanningReq { @NotNull private Boolean allSensitiveRules; private List sensitiveRuleIds; + @NotNull + private ScanningModeType scanningMode = ScanningModeType.JOINT_RECOGNITION; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java index 770038c56a..dec6e18bfa 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java @@ -70,6 +70,12 @@ public class SensitiveRule implements SecurityResource, SingleOrganizationResour private List pathExcludes = new ArrayList<>(); + private List aiSensitiveTypes = new ArrayList<>(); + + private Integer aiConfidenceThreshold; + + private String aiCustomPrompt; + @NotNull private Long maskingAlgorithmId; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRuleType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRuleType.java index 0673cbbf95..54ba6a4be0 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRuleType.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRuleType.java @@ -33,5 +33,10 @@ public enum SensitiveRuleType { /** * Path expression fuzzy match */ - PATH + PATH, + + /** + * AI-based sensitive data detection + */ + AI } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java new file mode 100644 index 0000000000..8e58c838cb --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -0,0 +1,171 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.recognizer; + +import java.io.IOException; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import java.util.stream.Collectors; + +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.collect.Lists; +import com.oceanbase.odc.service.common.util.SpringContextUtil; +import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; +import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; +import com.openai.models.chat.completions.ChatCompletion; +import lombok.Data; + +/** + * AI 列识别器（最终版） + */ +public class AIColumnRecognizer implements ColumnRecognizer { + + private final SensitiveRule aiRule; // 直接保存整个规则对象 + // @Value() + private static final int BATCH_SIZE = 50; // 这个批次大小需要能全局设置 + private static final ObjectMapper objectMapper = new ObjectMapper(); // 用于解析JSON + // private static final Pattern JSON_PATTERN = + // Pattern.compile("(?s)```json\\s*(\\{.*\\})\\s*```|(\\{.*\\})"); + //// 匹配 {...} 或 [...] + private static final Pattern JSON_PATTERN = Pattern + .compile("(?s)```json\\s*([\\{\\[].*[\\}\\]])\\s*```|([\\{\\[].*[\\}\\]])"); + + public AIColumnRecognizer(SensitiveRule rule) { + this.aiRule = rule; + } + + @Override + public Optional recognize(DBTableColumn column) { + // 通过调用批量识别方法来实现单个识别 + Map> batchResult = recognizeBatch(Collections.singletonList(column)); + String columnKey = getColumnKey(column); + return batchResult.getOrDefault(columnKey, Optional.empty()); + } + + @Override + public Map> recognizeBatch(List columns) { + if (columns == null || columns.isEmpty()) { + return Collections.emptyMap(); + } + // 1. 获取依赖的服务 + PromptTemplateLoader promptTemplateLoader = SpringContextUtil.getBean(PromptTemplateLoader.class); + AIInferenceService aiService = SpringContextUtil.getBean(AIInferenceService.class); + + // 2. 将所有待处理的列切分成多个小批次 + List> batches = Lists.partition(columns, BATCH_SIZE); + Map> finalAiResults = new HashMap<>(); + + try { + // 3. 遍历每一个小批次，分别调用 AI + for (List batch : batches) { + // a. 为这个小批次构建 Prompt + String prompt = buildBatchPrompt(promptTemplateLoader, batch); + // b. 调用 AI + ChatCompletion completion = aiService.chat(prompt); + String rawContent = completion.choices().get(0).message().content().orElse("[]"); + + // c. 使用正则表达式从AI的返回结果中安全地提取JSON数组字符串 + Matcher matcher = JSON_PATTERN.matcher(rawContent); + String jsonArrayResponse = "[]"; // 提供一个安全的默认值，以防匹配失败 + if (matcher.find()) { + // group(1) 对应被 ```json [...] ``` 包裹的内容, group(2) 对应裸露的 [...] + // 使用 Optional 来优雅地处理可能为null的捕获组 + jsonArrayResponse = Optional.ofNullable(matcher.group(1)).orElse(matcher.group(2)); + } + + // d. 解析提取出的、更纯净的 JSON 数组 + List batchResults = objectMapper.readValue(jsonArrayResponse, + new TypeReference>() { + }); + + // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 + int maxIndex = Math.min(batch.size(), batchResults.size()); + for (int i = 0; i < maxIndex; i++) { + DBTableColumn column = batch.get(i); + String columnKey = getColumnKey(column); + AiResponseDto dto = batchResults.get(i); + + if (dto.isSensitive()) { + RecognitionResult result = RecognitionResult.builder() + .matched(true) + .matchedRuleId(this.aiRule.getId()) + .level(dto.getRiskLevel()) + .sourceRuleType(SensitiveRuleType.AI) + .sensitiveType(dto.getSensitiveType()) + .confidence(dto.getConfidence()) + .build(); + finalAiResults.put(columnKey, Optional.of(result)); + } else { + finalAiResults.put(columnKey, Optional.empty()); + } + } + + // 如果AI返回的结果数量与输入不匹配，记录警告信息 + if (batchResults.size() != batch.size()) { + System.err.println("警告: AI返回结果数量(" + batchResults.size() + + ")与输入列数量(" + batch.size() + ")不匹配"); + } + } + } catch (Exception e) { + // 在实际项目中，应使用日志系统记录详细错误 + e.printStackTrace(); + // 出现异常时返回当前已成功识别的结果，或返回空map + return finalAiResults; + } + // 4. 返回包含所有 AI 识别结果的完整 Map + return finalAiResults; + } + + // 构建批量 Prompt 的新逻辑 + private String buildBatchPrompt(PromptTemplateLoader promptTemplateLoader, List batch) + throws IOException { + // 将一批列的元数据转换为 JSON 数组字符串 + List> columnMetadataList = batch.stream().map(c -> { + Map meta = new HashMap<>(); + meta.put("schemaName", c.getSchemaName()); + meta.put("tableName", c.getTableName()); + meta.put("columnName", c.getName()); + meta.put("comment", c.getComment()); + meta.put("dataType", c.getTypeName()); + return meta; + }).collect(Collectors.toList()); + String columnsJson = objectMapper.writeValueAsString(columnMetadataList); + // 调用 PromptTemplateLoader 来填充 + return promptTemplateLoader.buildPrompt(columnsJson, aiRule.getAiSensitiveTypes(), aiRule.getAiCustomPrompt()); + } + + // 移除私有方法，直接使用接口的默认实现 + + // 用于承载 AI 返回的 JSON 数据的内部类 + @Data + private static class AiResponseDto { + private boolean sensitive; + private SensitiveLevel riskLevel; + private Integer confidence; + private String sensitiveType; + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java index 9a301990bf..381d08d119 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java @@ -15,6 +15,13 @@ */ package com.oceanbase.odc.service.datasecurity.recognizer; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** @@ -29,6 +36,39 @@ public interface ColumnRecognizer { * @param column column {@link DBTableColumn} * @return recognizing result */ - boolean recognize(DBTableColumn column); + Optional recognize(DBTableColumn column); + + /** + * Batch recognizing the columns in database + * + * @param columns list of columns {@link DBTableColumn} + * @return map of recognizing results, key is column identifier, value is + * recognizing result + */ + default Map> recognizeBatch(List columns) { + if (columns == null || columns.isEmpty()) { + return Collections.emptyMap(); + } -} + Map> results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional result = recognize(column); + results.put(columnKey, result); + } + return results; + } + + /** + * Generate a unique key for the column + * + * @param column column {@link DBTableColumn} + * @return unique column key + */ + default String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java index 9cf782d9a7..e2b203abec 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java @@ -15,8 +15,13 @@ */ package com.oceanbase.odc.service.datasecurity.recognizer; +import java.util.Optional; + import org.codehaus.groovy.control.CompilerConfiguration; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.odc.service.datasecurity.util.SecureAstCustomizerUtil; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; @@ -32,26 +37,41 @@ */ public class GroovyColumnRecognizer implements ColumnRecognizer { + private final SensitiveRule rule; private final Script script; private static final String COLUMN_KEYWORD = "column"; - public GroovyColumnRecognizer(String groovyScript) { + // 【修改】构造函数接收 SensitiveRule 对象 + public GroovyColumnRecognizer(SensitiveRule rule) { + this.rule = rule; CompilerConfiguration config = new CompilerConfiguration(); config.addCompilationCustomizers(SecureAstCustomizerUtil.buildSecureASTCustomizer()); GroovyShell shell = new GroovyShell(config); - this.script = shell.parse(groovyScript); + this.script = shell.parse(rule.getGroovyScript()); } @Override - public boolean recognize(DBTableColumn column) { + public Optional recognize(DBTableColumn column) { try { GroovyColumnMeta groovyColumnMeta = new GroovyColumnMeta(column); Binding binding = new Binding(); binding.setVariable(COLUMN_KEYWORD, groovyColumnMeta); script.setBinding(binding); - return (boolean) script.run(); + + // 【修改】获取脚本执行结果，并根据结果构建返回 + boolean matched = (boolean) script.run(); + if (matched) { + RecognitionResult result = RecognitionResult.builder() + .matched(true) + .matchedRuleId(this.rule.getId()) + .level(this.rule.getLevel()) + .sourceRuleType(SensitiveRuleType.GROOVY) + .build(); + return Optional.of(result); + } + return Optional.empty(); } catch (Exception e) { - return false; + return Optional.empty(); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java index 3669113ead..6fd489d518 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java @@ -17,12 +17,16 @@ import java.util.List; import java.util.Objects; +import java.util.Optional; import java.util.stream.Collectors; import org.apache.commons.io.FilenameUtils; import org.apache.commons.io.IOCase; import com.oceanbase.odc.core.shared.PreConditions; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** @@ -31,33 +35,43 @@ */ public class PathColumnRecognizer implements ColumnRecognizer { + private final SensitiveRule rule; private final List pathIncludeMatchers; private final List pathExcludeMatchers; - public PathColumnRecognizer(List pathIncludes, List pathExcludes) { - pathIncludeMatchers = pathIncludes.stream().map(FieldPathMatcher::new).collect(Collectors.toList()); - pathExcludeMatchers = pathExcludes.stream().map(FieldPathMatcher::new).collect(Collectors.toList()); + // 【修改】构造函数接收 SensitiveRule 对象 + public PathColumnRecognizer(SensitiveRule rule) { + this.rule = rule; + pathIncludeMatchers = rule.getPathIncludes().stream().map(FieldPathMatcher::new).collect(Collectors.toList()); + pathExcludeMatchers = rule.getPathExcludes().stream().map(FieldPathMatcher::new).collect(Collectors.toList()); } @Override - public boolean recognize(DBTableColumn column) { + public Optional recognize(DBTableColumn column) { try { String schemaName = column.getSchemaName(); String tableName = column.getTableName(); String columnName = column.getName(); for (FieldPathMatcher matcher : pathExcludeMatchers) { if (matcher.match(schemaName, tableName, columnName)) { - return false; + return Optional.empty(); } } for (FieldPathMatcher matcher : pathIncludeMatchers) { if (matcher.match(schemaName, tableName, columnName)) { - return true; + // 【修改】匹配成功，构建并返回 RecognitionResult + RecognitionResult result = RecognitionResult.builder() + .matched(true) + .matchedRuleId(this.rule.getId()) + .level(this.rule.getLevel()) + .sourceRuleType(SensitiveRuleType.PATH) + .build(); + return Optional.of(result); } } - return false; + return Optional.empty(); } catch (Exception e) { - return false; + return Optional.empty(); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java index 02b15d1152..f3017e4cd0 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java @@ -15,9 +15,13 @@ */ package com.oceanbase.odc.service.datasecurity.recognizer; +import java.util.Optional; import java.util.regex.Pattern; import com.oceanbase.odc.common.util.StringUtils; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import lombok.NonNull; @@ -28,6 +32,8 @@ */ public class RegexColumnRecognizer implements ColumnRecognizer { + // 【修改】直接保存整个规则对象，以便获取 ID 和 Level + private final SensitiveRule rule; private final Pattern databasePattern; private final Pattern tablePattern; private final Pattern columnPattern; @@ -35,35 +41,44 @@ public class RegexColumnRecognizer implements ColumnRecognizer { private static final long MATCH_TIMEOUT_MILLIS = 100L; - public RegexColumnRecognizer(String databaseRegex, String tableRegex, String columnRegex, String commentRegex) { - databasePattern = StringUtils.isNotBlank(databaseRegex) ? Pattern.compile(databaseRegex) : null; - tablePattern = StringUtils.isNotBlank(tableRegex) ? Pattern.compile(tableRegex) : null; - columnPattern = StringUtils.isNotBlank(columnRegex) ? Pattern.compile(columnRegex) : null; - columnCommentPattern = StringUtils.isNotBlank(commentRegex) ? Pattern.compile(commentRegex) : null; + // 【修改】构造函数接收 SensitiveRule 对象 + public RegexColumnRecognizer(SensitiveRule rule) { + this.rule = rule; + databasePattern = StringUtils.isNotBlank(rule.getDatabaseRegexExpression()) ? Pattern.compile(rule.getDatabaseRegexExpression()) : null; + tablePattern = StringUtils.isNotBlank(rule.getTableRegexExpression()) ? Pattern.compile(rule.getTableRegexExpression()) : null; + columnPattern = StringUtils.isNotBlank(rule.getColumnRegexExpression()) ? Pattern.compile(rule.getColumnRegexExpression()) : null; + columnCommentPattern = StringUtils.isNotBlank(rule.getColumnCommentRegexExpression()) ? Pattern.compile(rule.getColumnCommentRegexExpression()) : null; } @Override - public boolean recognize(DBTableColumn column) { + public Optional recognize(DBTableColumn column) { try { if (databasePattern != null && !databasePattern .matcher(new TimeoutCharSequence(column.getSchemaName(), getTimeoutMillis())).matches()) { - return false; + return Optional.empty(); } if (tablePattern != null && !tablePattern .matcher(new TimeoutCharSequence(column.getTableName(), getTimeoutMillis())).matches()) { - return false; + return Optional.empty(); } if (columnPattern != null && !columnPattern .matcher(new TimeoutCharSequence(column.getName(), getTimeoutMillis())).matches()) { - return false; + return Optional.empty(); } if (columnCommentPattern != null && !columnCommentPattern .matcher(new TimeoutCharSequence(column.getComment(), getTimeoutMillis())).matches()) { - return false; + return Optional.empty(); } - return true; + // 【修改】如果所有条件都通过，说明匹配成功，构建并返回 RecognitionResult + RecognitionResult result = RecognitionResult.builder() + .matched(true) + .matchedRuleId(this.rule.getId()) + .level(this.rule.getLevel()) + .sourceRuleType(SensitiveRuleType.REGEX) + .build(); + return Optional.of(result); } catch (Exception e) { - return false; + return Optional.empty(); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java new file mode 100644 index 0000000000..617cf34f08 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java @@ -0,0 +1,109 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 抽象扫描策略基类，提供公共的工具方法 + * + * @author Assistant + * @date 2025/1/27 + */ +public abstract class AbstractScanningStrategy implements ScanningStrategy { + + /** + * 从识别器列表中找到第一个匹配的结果 + * + * @param recognizers 识别器列表 + * @param column 待识别的列 + * @return 第一个匹配的识别结果 + */ + protected Optional findFirstMatch(List recognizers, DBTableColumn column) { + for (ColumnRecognizer recognizer : recognizers) { + Optional result = recognizer.recognize(column); + if (result.isPresent()) { + return result; + } + } + return Optional.empty(); + } + + /** + * 批量查找所有列的第一个匹配结果 + * + * @param recognizers 识别器列表 + * @param columns 待识别的列列表 + * @return 列标识符到识别结果的映射 + */ + protected Map> findAllFirstMatches(List recognizers, + List columns) { + if (recognizers.isEmpty() || columns.isEmpty()) { + return createEmptyResultMap(columns); + } + + // 尝试使用批量识别（优先用于AI识别器） + if (recognizers.size() == 1) { + ColumnRecognizer recognizer = recognizers.get(0); + return recognizer.recognizeBatch(columns); + } + + // 多个识别器时，逐个处理以保证优先级 + Map> results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional result = findFirstMatch(recognizers, column); + results.put(columnKey, result); + } + return results; + } + + /** + * 为列列表创建空结果映射 + * + * @param columns 列列表 + * @return 空结果映射 + */ + protected Map> createEmptyResultMap(List columns) { + Map> results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + results.put(columnKey, Optional.empty()); + } + return results; + } + + /** + * 生成列的唯一标识符 + * + * @param column 列信息 + * @return 列标识符 + */ + protected String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java new file mode 100644 index 0000000000..6595104b3e --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java @@ -0,0 +1,89 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 联合识别策略实现 + * 优先使用基础规则识别，如果基础规则没有匹配到，则使用AI识别器作为补充 + * 不一致时信任规则结果 + * + * @author Assistant + * @date 2025/1/27 + */ +public class JointRecognitionStrategy extends AbstractScanningStrategy { + + @Override + public ScanResult scan(DBTableColumn column, List basicRecognizers, + List aiRecognizers) { + Optional basicResult = findFirstMatch(basicRecognizers, column); + + // 如果基础规则已经识别出来，就以此为准，不再调用AI + if (basicResult.isPresent()) { + return new ScanResult(basicResult, Optional.empty()); + } + + // 否则，调用AI作为补充 + Optional aiResult = findFirstMatch(aiRecognizers, column); + return new ScanResult(Optional.empty(), aiResult); + } + + @Override + public Map scanBatch(List columns, List basicRecognizers, + List aiRecognizers) { + Map> basicResults = findAllFirstMatches(basicRecognizers, columns); + + // 收集没有被基础规则识别的列 + List remainingColumns = new ArrayList<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional basicResult = basicResults.getOrDefault(columnKey, Optional.empty()); + if (!basicResult.isPresent()) { + remainingColumns.add(column); + } + } + + // 对剩余的列进行AI识别 + Map> aiResults = findAllFirstMatches(aiRecognizers, remainingColumns); + + // 合并结果 + Map results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional basicResult = basicResults.getOrDefault(columnKey, Optional.empty()); + + if (basicResult.isPresent()) { + results.put(columnKey, new ScanResult(basicResult, Optional.empty())); + } else { + Optional aiResult = aiResults.getOrDefault(columnKey, Optional.empty()); + results.put(columnKey, new ScanResult(Optional.empty(), aiResult)); + } + } + + return results; + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java new file mode 100644 index 0000000000..ac441b2b76 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java @@ -0,0 +1,72 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 规则+AI策略实现 + * 同时执行基础规则和AI识别，用于差异化展示两种结果 + * + * @author Assistant + * @date 2025/1/27 + */ +public class RulesAndAiStrategy extends AbstractScanningStrategy { + + @Override + public ScanResult scan(DBTableColumn column, List basicRecognizers, + List aiRecognizers) { + Optional basicResult = findFirstMatch(basicRecognizers, column); + Optional aiResult = findFirstMatch(aiRecognizers, column); + return new ScanResult(basicResult, aiResult); + } + + @Override + public Map scanBatch(List columns, List basicRecognizers, + List aiRecognizers) { + // 并行执行基础规则和AI识别 + CompletableFuture>> basicFuture = CompletableFuture + .supplyAsync(() -> findAllFirstMatches(basicRecognizers, columns)); + CompletableFuture>> aiFuture = CompletableFuture + .supplyAsync(() -> findAllFirstMatches(aiRecognizers, columns)); + + // 等待两个任务完成并合并结果 + CompletableFuture.allOf(basicFuture, aiFuture).join(); + + Map> basicResults = basicFuture.join(); + Map> aiResults = aiFuture.join(); + + Map results = new HashMap<>(); + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional basicResult = basicResults.getOrDefault(columnKey, Optional.empty()); + Optional aiResult = aiResults.getOrDefault(columnKey, Optional.empty()); + results.put(columnKey, new ScanResult(basicResult, aiResult)); + } + + return results; + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java new file mode 100644 index 0000000000..f494c39047 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 仅规则扫描策略实现 + * 只使用基础规则进行识别，忽略AI识别器 + * + * @author Assistant + * @date 2025/1/27 + */ +public class RulesOnlyStrategy extends AbstractScanningStrategy { + + @Override + public ScanResult scan(DBTableColumn column, List basicRecognizers, + List aiRecognizers) { + Optional basicResult = findFirstMatch(basicRecognizers, column); + return new ScanResult(basicResult, Optional.empty()); + } + + @Override + public Map scanBatch(List columns, List basicRecognizers, + List aiRecognizers) { + Map> basicResults = findAllFirstMatches(basicRecognizers, columns); + Map results = new HashMap<>(); + + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional basicResult = basicResults.getOrDefault(columnKey, Optional.empty()); + results.put(columnKey, new ScanResult(basicResult, Optional.empty())); + } + + return results; + } +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java new file mode 100644 index 0000000000..2cdfdbc30f --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java @@ -0,0 +1,55 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.List; +import java.util.Map; + +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 敏感列扫描策略接口 + * + * @author Assistant + * @date 2025/1/27 + */ +public interface ScanningStrategy { + + /** + * 执行单个列的扫描 + * + * @param column 待扫描的列 + * @param basicRecognizers 基础规则识别器列表 + * @param aiRecognizers AI识别器列表 + * @return 扫描结果 + */ + ScanResult scan(DBTableColumn column, List basicRecognizers, + List aiRecognizers); + + /** + * 执行批量列的扫描 + * + * @param columns 待扫描的列列表 + * @param basicRecognizers 基础规则识别器列表 + * @param aiRecognizers AI识别器列表 + * @return 扫描结果映射，key为列标识符，value为扫描结果 + */ + Map scanBatch(List columns, List basicRecognizers, + List aiRecognizers); +} \ No newline at end of file diff --git a/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt b/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt new file mode 100644 index 0000000000..b38b562000 --- /dev/null +++ b/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt @@ -0,0 +1,46 @@ +You are a data classification expert specializing in data privacy and security. Your task is to determine if a database column's content likely belongs to one of the user-specified sensitive data categories. + +Analyze the JSON array of database columns provided below. For EACH column in the array, determine if it belongs to any of the specified sensitive data categories. + +1. **Sensitive Categories to Check:** +{sensitiveTypes} + +2. **JSON Array of Columns to Analyze:** +{DBTableColumn} + +3. **Additional User-Provided Hint:** +{customPrompt} + +Based on all the information above, return a JSON array where each object corresponds to a column from the input array in the SAME order. +For each column, respond with a JSON object in the following format ONLY. Do not add any other text or explanations. +The final output should be a single, valid JSON array. Please do not output Markdown code blocks, such as ```json. + +Example of expected output format for 3 columns: +[ + { + "sensitive": true, + "riskLevel": "HIGH", + "confidence": 95, + "sensitiveType": "Financial Information" + }, + { + "sensitive": true, + "riskLevel": "MEDIUM", + "confidence": 80, + "sensitiveType": "Personal Identification" + }, + { + "sensitive": false, + "riskLevel": "LOW", + "confidence": 30, + "sensitiveType": null + } +] + +Format for each object in the output array: +{ + "sensitive": boolean, + "riskLevel": "HIGH" | "MEDIUM" | "LOW", + "confidence": number (an int value between 0 and 100), + "sensitiveType": string (the specific category it belongs to, or null) +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java new file mode 100644 index 0000000000..8bd122b809 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java @@ -0,0 +1,202 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; + +import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 敏感列扫描器单元测试类 + */ +public class SensitiveColumnScannerTest { + + private List rules; + private SensitiveColumnScanner scanner; + + @Before + public void setUp() { + // 创建测试规则 + SensitiveRule regexRule = createTestRegexRule(1L); + SensitiveRule aiRule = createTestAIRule(2L); + rules = Arrays.asList(regexRule, aiRule); + ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); + scanner = new SensitiveColumnScanner(rules, strategyFactory); + } + + @Test + public void scan_rulesOnlyMode_returnBasicResult() { + // 准备阶段：创建能被正则规则匹配的测试列 + DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); + + // 执行阶段：使用仅规则模式扫描 + ScanResult result = scanner.scan(column, ScanningModeType.RULES_ONLY); + + // 验证阶段：确认返回基础规则结果 + Assert.assertNotNull("结果不应为null", result); + Assert.assertTrue("基础规则结果应存在", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("AI规则结果应不存在", result.getAiRuleResult().isPresent()); + + RecognitionResult basicResult = result.getBasicRuleResult().get(); + Assert.assertTrue("应匹配成功", basicResult.isMatched()); + Assert.assertEquals("匹配的规则ID应为1", Long.valueOf(1), basicResult.getMatchedRuleId()); + } + + @Test + public void scan_jointRecognitionMode_ruleMatched_returnBasicResultOnly() { + // 准备阶段：创建能被正则规则匹配的测试列 + DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); + + // 执行阶段：使用联合识别模式扫描 + ScanResult scanResult = scanner.scan(column, ScanningModeType.JOINT_RECOGNITION); + + // 验证阶段：确认只返回基础规则结果 + Assert.assertNotNull("结果不应为null", scanResult); + Assert.assertTrue("基础规则结果应存在", scanResult.getBasicRuleResult().isPresent()); + Assert.assertFalse("AI规则结果应不存在", scanResult.getAiRuleResult().isPresent()); + } + + @Test + public void scan_jointRecognitionMode_ruleNotMatched_returnAiResult() { + // 准备阶段：创建不能被正则规则匹配的测试列 + DBTableColumn column = createTestColumn("test_db", "products", "name", "product name"); + + // 执行阶段：使用联合识别模式扫描 + ScanResult result = scanner.scan(column, ScanningModeType.JOINT_RECOGNITION); + + // 验证阶段：确认基础规则结果不存在（因为没有mock AI识别器，所以AI结果也不存在） + Assert.assertNotNull("结果不应为null", result); + Assert.assertFalse("基础规则结果应不存在", result.getBasicRuleResult().isPresent()); + } + + @Test + public void scan_rulesAndAiMode_returnBothResults() { + // 准备阶段：创建能被正则规则匹配的测试列 + DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); + + // 执行阶段：使用规则+AI模式扫描 + ScanResult result = scanner.scan(column, ScanningModeType.RULES_AND_AI); + + // 验证阶段：确认返回基础规则结果（AI结果不存在因为没有mock） + Assert.assertNotNull("结果不应为null", result); + Assert.assertTrue("基础规则结果应存在", result.getBasicRuleResult().isPresent()); + } + + @Test + public void scanBatch_emptyList_returnEmptyMap() { + // 执行阶段：使用仅规则模式批量扫描空列表 + Map result = scanner.scanBatch(Collections.emptyList(), ScanningModeType.RULES_ONLY); + + // 验证阶段：确认返回空Map + Assert.assertNotNull("结果不应为null", result); + Assert.assertTrue("空列表应返回空Map", result.isEmpty()); + } + + @Test + public void scanBatch_rulesOnlyMode_returnBasicResults() { + // 准备阶段：创建多个测试列，其中一列能被正则规则匹配 + DBTableColumn column1 = createTestColumn("test_db", "users", "email", "user email address"); // 能匹配 + DBTableColumn column2 = createTestColumn("test_db", "products", "name", "product name"); // 不能匹配 + List columns = Arrays.asList(column1, column2); + + // 执行阶段：使用仅规则模式批量扫描 + Map results = scanner.scanBatch(columns, ScanningModeType.RULES_ONLY); + + // 验证阶段：确认返回正确的批量扫描结果 + Assert.assertNotNull("结果不应为null", results); + Assert.assertEquals("应返回两个扫描结果", 2, results.size()); + Assert.assertTrue("应包含users.email的结果", results.containsKey("users.email")); + Assert.assertTrue("应包含products.name的结果", results.containsKey("products.name")); + + // 验证email列的扫描结果（应该匹配） + ScanResult emailResult = results.get("users.email"); + Assert.assertTrue("email的基础规则结果应存在", emailResult.getBasicRuleResult().isPresent()); + Assert.assertEquals("email匹配的规则ID应为1", Long.valueOf(1), emailResult.getBasicRuleResult().get().getMatchedRuleId()); + + // 验证name列的扫描结果（不应该匹配） + ScanResult nameResult = results.get("products.name"); + Assert.assertFalse("name的基础规则结果应不存在", nameResult.getBasicRuleResult().isPresent()); + } + + // --- 辅助方法 --- + + /** + * 创建测试用的正则规则 + * @param id 规则ID + * @return 正则规则对象 + */ + private SensitiveRule createTestRegexRule(Long id) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.REGEX); + rule.setEnabled(true); + rule.setLevel(SensitiveLevel.HIGH); + rule.setDatabaseRegexExpression("^\\S+$"); + rule.setTableRegexExpression("^\\S+$"); + rule.setColumnRegexExpression("^\\S*email\\S*$"); + rule.setColumnCommentRegexExpression("^[\\S\\s]*email[\\S\\s]*$"); + return rule; + } + + /** + * 创建测试用的AI规则 + * @param id 规则ID + * @return AI规则对象 + */ + private SensitiveRule createTestAIRule(Long id) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.AI); + rule.setEnabled(true); + rule.setLevel(SensitiveLevel.HIGH); + rule.setAiSensitiveTypes(Arrays.asList("联系方式", "财务信息", "身份信息")); + rule.setAiConfidenceThreshold(80); + rule.setAiCustomPrompt("请识别敏感数据列"); + return rule; + } + + /** + * 创建测试用的数据库列 + * @param schemaName 数据库名 + * @param tableName 表名 + * @param columnName 列名 + * @param comment 列注释 + * @return 数据库列对象 + */ + private DBTableColumn createTestColumn(String schemaName, String tableName, String columnName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setSchemaName(schemaName); + column.setTableName(tableName); + column.setName(columnName); + column.setComment(comment); + column.setTypeName("VARCHAR"); + return column; + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java new file mode 100644 index 0000000000..ad0398bcef --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java @@ -0,0 +1,218 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.recognizer; + +import static org.mockito.Mockito.*; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; + +import com.oceanbase.odc.service.common.util.SpringContextUtil; +import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; +import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; +import com.openai.models.chat.completions.ChatCompletion; + +/** + * AI列识别器单元测试类 + */ +public class AIColumnRecognizerTest { + + @Mock + private PromptTemplateLoader promptTemplateLoader; + + @Mock + private AIInferenceService aiInferenceService; + + @Mock + private SpringContextUtil springContextUtil; + + private SensitiveRule aiRule; + private AIColumnRecognizer aiColumnRecognizer; + + @Before + public void setUp() { + MockitoAnnotations.initMocks(this); + + // Mock SpringContextUtil + mockStatic(SpringContextUtil.class); + when(SpringContextUtil.getBean(PromptTemplateLoader.class)).thenReturn(promptTemplateLoader); + when(SpringContextUtil.getBean(AIInferenceService.class)).thenReturn(aiInferenceService); + + // 创建测试用的AI规则 + aiRule = createTestAIRule(1L); + aiColumnRecognizer = new AIColumnRecognizer(aiRule); + } + + @Test + public void recognize_singleColumn_returnEmpty() { + // 准备阶段：创建测试列 + DBTableColumn column = createTestColumn("test_db", "users", "email", "用户邮箱"); + + // 执行阶段：调用单列识别方法 + Optional result = aiColumnRecognizer.recognize(column); + + // 验证阶段：确认单列识别返回空结果 + Assert.assertFalse("单列识别应返回空结果", result.isPresent()); + } + + @Test + public void recognizeBatch_emptyList_returnEmptyMap() { + // 执行阶段：调用批量识别方法，传入空列表 + Map result = aiColumnRecognizer.recognizeBatch(Collections.emptyList()); + + // 验证阶段：确认返回空Map + Assert.assertNotNull("结果不应为null", result); + Assert.assertTrue("空列表应返回空Map", result.isEmpty()); + } + + @Test + public void recognizeBatch_nullList_returnEmptyMap() { + // 执行阶段：调用批量识别方法，传入null + Map result = aiColumnRecognizer.recognizeBatch(null); + + // 验证阶段：确认返回空Map + Assert.assertNotNull("结果不应为null", result); + Assert.assertTrue("null列表应返回空Map", result.isEmpty()); + } + + @Test + public void recognizeBatch_singleColumn_returnResult() throws Exception { + // 准备阶段：创建测试数据和模拟AI服务响应 + DBTableColumn column = createTestColumn("test_db", "users", "email", "用户邮箱"); + List columns = Arrays.asList(column); + + String prompt = "test prompt"; + when(promptTemplateLoader.buildPrompt(anyString(), anyList(), anyString())).thenReturn(prompt); + + String aiResponse = "[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 95, \"sensitiveType\": \"联系方式\"}]"; + ChatCompletion chatCompletion = mock(ChatCompletion.class); + when(chatCompletion.choices()).thenReturn(Arrays.asList(mock(com.openai.models.chat.completions.ChatCompletion.Choice.class))); + when(chatCompletion.choices().get(0).message()).thenReturn(mock(com.openai.models.chat.completions.ChatCompletionMessage.class)); + when(chatCompletion.choices().get(0).message().content()).thenReturn(Optional.of(aiResponse)); + when(aiInferenceService.chat(prompt)).thenReturn(chatCompletion); + + // 执行阶段：调用批量识别方法 + Map result = aiColumnRecognizer.recognizeBatch(columns); + + // 验证阶段：确认返回正确的识别结果 + Assert.assertNotNull("结果不应为null", result); + Assert.assertEquals("应返回一个识别结果", 1, result.size()); + Assert.assertTrue("应包含指定列的结果", result.containsKey("users.email")); + + RecognitionResult recognitionResult = result.get("users.email"); + Assert.assertTrue("应识别为敏感列", recognitionResult.isMatched()); + Assert.assertEquals("规则ID应匹配", aiRule.getId(), recognitionResult.getMatchedRuleId()); + Assert.assertEquals("规则类型应为AI", SensitiveRuleType.AI, recognitionResult.getSourceRuleType()); + Assert.assertEquals("风险等级应为HIGH", SensitiveLevel.HIGH, recognitionResult.getLevel()); + Assert.assertEquals("置信度应为95", Double.valueOf(95), recognitionResult.getConfidence()); + Assert.assertEquals("敏感类型应为联系方式", "联系方式", recognitionResult.getSensitiveType()); + } + + @Test + public void recognizeBatch_multipleColumns_returnResults() throws Exception { + // 准备阶段：创建多个测试列和模拟AI服务响应 + DBTableColumn column1 = createTestColumn("test_db", "users", "email", "用户邮箱"); + DBTableColumn column2 = createTestColumn("test_db", "employees", "salary", "员工薪资"); + DBTableColumn column3 = createTestColumn("test_db", "products", "name", "产品名称"); + List columns = Arrays.asList(column1, column2, column3); + + String prompt = "test prompt"; + when(promptTemplateLoader.buildPrompt(anyString(), anyList(), anyString())).thenReturn(prompt); + + // 模拟AI响应：前两列为敏感列，第三列（name）为非敏感列 + String aiResponse = "[" + + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 95, \"sensitiveType\": \"联系方式\"}," + + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 90, \"sensitiveType\": \"财务信息\"}," + + "{\"sensitive\": false, \"riskLevel\": \"LOW\", \"confidence\": 20, \"sensitiveType\": null}" + + "]"; + ChatCompletion chatCompletion = mock(ChatCompletion.class); + when(chatCompletion.choices()).thenReturn(Arrays.asList(mock(com.openai.models.chat.completions.ChatCompletion.Choice.class))); + when(chatCompletion.choices().get(0).message()).thenReturn(mock(com.openai.models.chat.completions.ChatCompletionMessage.class)); + when(chatCompletion.choices().get(0).message().content()).thenReturn(Optional.of(aiResponse)); + when(aiInferenceService.chat(prompt)).thenReturn(chatCompletion); + + // 执行阶段：调用批量识别方法 + Map result = aiColumnRecognizer.recognizeBatch(columns); + + // 验证阶段：确认返回正确的识别结果数量（只有敏感列会被返回） + Assert.assertNotNull("结果不应为null", result); + Assert.assertEquals("应返回两个识别结果（只有敏感列会被返回）", 2, result.size()); + Assert.assertTrue("应包含users.email的结果", result.containsKey("users.email")); + Assert.assertTrue("应包含employees.salary的结果", result.containsKey("employees.salary")); + Assert.assertFalse("不应包含products.name的结果（非敏感列）", result.containsKey("products.name")); + + // 验证email列的识别结果 + RecognitionResult emailResult = result.get("users.email"); + Assert.assertTrue("email应识别为敏感列", emailResult.isMatched()); + Assert.assertEquals("email敏感类型应为联系方式", "联系方式", emailResult.getSensitiveType()); + + // 验证salary列的识别结果 + RecognitionResult salaryResult = result.get("employees.salary"); + Assert.assertTrue("salary应识别为敏感列", salaryResult.isMatched()); + Assert.assertEquals("salary敏感类型应为财务信息", "财务信息", salaryResult.getSensitiveType()); + } + + // --- 辅助方法 --- + + /** + * 创建测试用的AI规则 + * @param id 规则ID + * @return AI规则对象 + */ + private SensitiveRule createTestAIRule(Long id) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.AI); + rule.setEnabled(true); + rule.setLevel(SensitiveLevel.HIGH); + rule.setAiSensitiveTypes(Arrays.asList("联系方式", "财务信息", "身份信息")); + rule.setAiConfidenceThreshold(80); + rule.setAiCustomPrompt("请识别敏感数据列"); + return rule; + } + + /** + * 创建测试用的数据库列 + * @param schemaName 数据库名 + * @param tableName 表名 + * @param columnName 列名 + * @param comment 列注释 + * @return 数据库列对象 + */ + private DBTableColumn createTestColumn(String schemaName, String tableName, String columnName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setSchemaName(schemaName); + column.setTableName(tableName); + column.setName(columnName); + column.setComment(comment); + column.setTypeName("VARCHAR"); + return column; + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java index a936364325..f13c0c4b57 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java @@ -15,12 +15,18 @@ */ package com.oceanbase.odc.service.datasecurity.recognizer; +import java.util.Optional; + import org.codehaus.groovy.control.MultipleCompilationErrorsException; import org.junit.Assert; import org.junit.Rule; import org.junit.Test; import org.junit.rules.ExpectedException; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** @@ -29,38 +35,62 @@ */ public class GroovyColumnRecognizerTest { + // 原作者使用的 Junit 4 异常测试方式，我们予以保留 @Rule public ExpectedException thrown = ExpectedException.none(); @Test public void test_recognize_true() { - ColumnRecognizer recognizer = new GroovyColumnRecognizer(buildGroovyScript()); - DBTableColumn dbTableColumn = createDBTableColumn(); - Assert.assertTrue(recognizer.recognize(dbTableColumn)); + // 【修改】通过辅助方法创建规则和识别器 + SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); + ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); + DBTableColumn dbTableColumn = createTestColumn(); + + // 【修改】调用新的 recognize 方法并检查 Optional 返回值 + Optional resultOpt = recognizer.recognize(dbTableColumn); + + // 【修改】断言结果存在，并验证内容 + Assert.assertTrue("脚本匹配成功，应返回有值的 Optional", resultOpt.isPresent()); + RecognitionResult result = resultOpt.get(); + Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), result.getMatchedRuleId()); + Assert.assertEquals("规则类型应为 GROOVY", SensitiveRuleType.GROOVY, result.getSourceRuleType()); } @Test public void test_recognize_false() { - ColumnRecognizer recognizer = new GroovyColumnRecognizer(buildGroovyScript()); - DBTableColumn dbTableColumn = createDBTableColumn(); - dbTableColumn.setTableName("unmatched"); - Assert.assertFalse(recognizer.recognize(dbTableColumn)); + SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); + ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); + DBTableColumn dbTableColumn = createTestColumn(); + dbTableColumn.setTableName("unmatched_table"); // 使脚本匹配失败 + + Optional resultOpt = recognizer.recognize(dbTableColumn); + + // 【修改】断言结果为空 Optional + Assert.assertFalse("脚本匹配失败，应返回空的 Optional", resultOpt.isPresent()); } @Test public void test_recognize_nullColumnName() { - ColumnRecognizer recognizer = new GroovyColumnRecognizer(buildGroovyScript()); - DBTableColumn dbTableColumn = createDBTableColumn(); - dbTableColumn.setTableName(null); - Assert.assertFalse(recognizer.recognize(dbTableColumn)); + SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); + ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); + DBTableColumn dbTableColumn = createTestColumn(); + dbTableColumn.setName(null); // 脚本内部会因为 null.equals(...) 抛异常 + + Optional resultOpt = recognizer.recognize(dbTableColumn); + + // 【修改】断言结果为空 Optional，因为脚本执行异常被捕获 + Assert.assertFalse("脚本执行异常，应返回空的 Optional", resultOpt.isPresent()); } + // --- 以下所有安全校验测试，保留原作者的意图和断言方式 --- + @Test public void test_securityInterceptor_systemExit() { thrown.expect(Exception.class); thrown.expectMessage("Method call is not security"); String script = "System.exit(-1);"; - new GroovyColumnRecognizer(script); + // 【修改】调用新的构造函数 + new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @Test @@ -68,9 +98,9 @@ public void test_securityInterceptor_forLoop() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("ForStatements are not allowed"); String script = "for (int i = 0; i < 1; i++) {\n" - + " i = 0;\n" - + "}"; - new GroovyColumnRecognizer(script); + + " i = 0;\n" + + "}"; + new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @Test @@ -78,43 +108,45 @@ public void test_securityInterceptor_whileLoop() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("WhileStatements are not allowed"); String script = "while(true) {\n" - + " int i = 0;\n" - + "}"; - new GroovyColumnRecognizer(script); + + " int i = 0;\n" + + "}"; + new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @Test - public void test_securityInterceptor_threadSleep() throws InterruptedException { + public void test_securityInterceptor_threadSleep() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("java.lang.Thread"); String script = "Thread.sleep(1000);"; - new GroovyColumnRecognizer(script); + new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @Test - public void test_securityInterceptor_importPackage() throws InterruptedException { + public void test_securityInterceptor_importPackage() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("java.lang.System"); String script = "import java.lang.System;"; - new GroovyColumnRecognizer(script); + new GroovyColumnRecognizer(createGroovyRule(1L, script)); } - private String buildGroovyScript() { + // --- 辅助方法 --- + + private String buildDefaultGroovyScript() { return "if (column.name.equals(\"column\")) {\n" - + " if (column.table.equalsIgnoreCase(\"IAM_USER\")) {\n" - + " if (column.schema.length() > 0) {\n" - + " if (column.comment.indexOf(\"user\") > 0) {\n" - + " if (column.type.toLowerCase().equals(\"varchar\")) {\n" - + " return true;\n" - + " }\n" - + " }\n" - + " }\n" - + " }\n" - + "}\n" - + "return false;"; + + " if (column.table.equalsIgnoreCase(\"iam_user\")) {\n" + + " if (column.schema.length() > 0) {\n" + + " if (column.comment.indexOf(\"user\") > 0) {\n" + + " if (column.type.toLowerCase().equals(\"varchar\")) {\n" + + " return true;\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + "}\n" + + "return false;"; } - private DBTableColumn createDBTableColumn() { + private DBTableColumn createTestColumn() { DBTableColumn dbTableColumn = new DBTableColumn(); dbTableColumn.setSchemaName("odc_meta"); dbTableColumn.setTableName("iam_user"); @@ -124,4 +156,14 @@ private DBTableColumn createDBTableColumn() { return dbTableColumn; } + // 【新增】辅助方法，用于快速创建测试用的 Groovy 规则 + private SensitiveRule createGroovyRule(Long id, String script) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.GROOVY); + rule.setGroovyScript(script); + rule.setLevel(SensitiveLevel.HIGH); + rule.setEnabled(true); + return rule; + } } diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java index 8c3f19aa2a..21dec644e1 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java @@ -16,10 +16,16 @@ package com.oceanbase.odc.service.datasecurity.recognizer; import java.util.Arrays; +import java.util.List; +import java.util.Optional; import org.junit.Assert; import org.junit.Test; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** @@ -30,20 +36,44 @@ public class PathColumnRecognizerTest { @Test public void test_recognize_true() { - ColumnRecognizer recognizer = new PathColumnRecognizer(Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); - Assert.assertTrue(recognizer.recognize(createDBTableColumn("a", "b12", "c"))); - Assert.assertTrue(recognizer.recognize(createDBTableColumn("a12", "34b56", "c"))); - Assert.assertTrue(recognizer.recognize(createDBTableColumn("a12", "34b", "c"))); + // 【修改】通过辅助方法创建规则，并用规则创建识别器 + SensitiveRule rule = createPathRule(1L, Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); + ColumnRecognizer recognizer = new PathColumnRecognizer(rule); + + // 【修改】断言方式改为检查 Optional.isPresent() + // 原作者的每个测试用例都予以保留 + Optional result1 = recognizer.recognize(createDBTableColumn("a", "b12", "c")); + Assert.assertTrue("路径 'a.b12.c' 应匹配成功", result1.isPresent()); + // 增加对返回内容的校验，使测试更严谨 + Assert.assertEquals(rule.getId(), result1.get().getMatchedRuleId()); + + Optional result2 = recognizer.recognize(createDBTableColumn("a12", "34b56", "c")); + Assert.assertTrue("路径 'a12.34b56.c' 应匹配成功", result2.isPresent()); + + Optional result3 = recognizer.recognize(createDBTableColumn("a12", "34b", "c")); + Assert.assertTrue("路径 'a12.34b.c' 应匹配成功", result3.isPresent()); } @Test public void test_recognize_false() { - ColumnRecognizer recognizer = new PathColumnRecognizer(Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); - Assert.assertFalse(recognizer.recognize(createDBTableColumn("a", "b", "c"))); - Assert.assertFalse(recognizer.recognize(createDBTableColumn("a12", "b34", "c56"))); - Assert.assertFalse(recognizer.recognize(createDBTableColumn("a12", "b34", null))); + // 【修改】通过辅助方法创建规则 + SensitiveRule rule = createPathRule(1L, Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); + ColumnRecognizer recognizer = new PathColumnRecognizer(rule); + + // 【修改】断言方式改为检查 !Optional.isPresent() + // 原作者的每个测试用例都予以保留 + Assert.assertFalse("路径 'a.b.c' 应被排除，匹配失败", + recognizer.recognize(createDBTableColumn("a", "b", "c")).isPresent()); + + Assert.assertFalse("路径 'a12.b34.c56' 不应匹配，匹配失败", + recognizer.recognize(createDBTableColumn("a12", "b34", "c56")).isPresent()); + + Assert.assertFalse("路径 'a12.b34.null' 不应匹配，匹配失败", + recognizer.recognize(createDBTableColumn("a12", "b34", null)).isPresent()); } + // --- 辅助方法 --- + private DBTableColumn createDBTableColumn(String schemaName, String tableName, String columnName) { DBTableColumn column = new DBTableColumn(); column.setSchemaName(schemaName); @@ -51,4 +81,16 @@ private DBTableColumn createDBTableColumn(String schemaName, String tableName, S column.setName(columnName); return column; } + + // 【新增】辅助方法，用于快速创建测试用的 Path 规则 + private SensitiveRule createPathRule(Long id, List includes, List excludes) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.PATH); + rule.setPathIncludes(includes); + rule.setPathExcludes(excludes); + rule.setLevel(SensitiveLevel.HIGH); + rule.setEnabled(true); + return rule; + } } diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java index 524b2acd91..2e19ca422f 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java @@ -15,9 +15,15 @@ */ package com.oceanbase.odc.service.datasecurity.recognizer; +import java.util.Optional; + import org.junit.Assert; import org.junit.Test; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** @@ -28,20 +34,53 @@ public class RegexColumnRecognizerTest { @Test public void recognize_returnTrue() { - ColumnRecognizer recognizer = createRegexColumnRecognizer(); - Assert.assertTrue(recognizer.recognize(createDBTableColumn("xxx", "xxx", "user_email", "email of user"))); + // 【修改】通过辅助方法创建规则，并用规则创建识别器 + SensitiveRule rule = createTestRegexRule(1L); + ColumnRecognizer recognizer = new RegexColumnRecognizer(rule); + DBTableColumn column = createDBTableColumn("xxx", "xxx", "user_email", "email of user"); + + // 【修改】调用新的 recognize 方法并检查 Optional 返回值 + Optional resultOpt = recognizer.recognize(column); + + // 【修改】断言结果存在，并验证内容 + Assert.assertTrue("正则表达式应匹配成功", resultOpt.isPresent()); + Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), resultOpt.get().getMatchedRuleId()); } @Test public void recognize_returnFalse() { - ColumnRecognizer recognizer = createRegexColumnRecognizer(); - Assert.assertFalse(recognizer.recognize(createDBTableColumn("xxx", "xxx", "user_email", null))); - Assert.assertFalse(recognizer.recognize(createDBTableColumn(" ", "xxx", "user_email", null))); - Assert.assertFalse(recognizer.recognize(createDBTableColumn("xxx", "xxx", "user", "email"))); + // 【修改】通过辅助方法创建规则 + SensitiveRule rule = createTestRegexRule(1L); + ColumnRecognizer recognizer = new RegexColumnRecognizer(rule); + + // 【修改】断言方式改为检查 !Optional.isPresent() + // 原作者的每个测试用例都予以保留 + Assert.assertFalse("Comment 为 null 时不应匹配", + recognizer.recognize(createDBTableColumn("xxx", "xxx", "user_email", null)).isPresent()); + + Assert.assertFalse("SchemaName 不匹配时应失败", + recognizer.recognize(createDBTableColumn(" ", "xxx", "user_email", "email of user")).isPresent()); + + Assert.assertFalse("ColumnName 和 Comment 都不匹配时应失败", + recognizer.recognize(createDBTableColumn("xxx", "xxx", "user", "some info")).isPresent()); } - private RegexColumnRecognizer createRegexColumnRecognizer() { - return new RegexColumnRecognizer("^\\S+$", "^\\S+$", "^\\S*email\\S*$", "^[\\S\\s]*email[\\S\\s]*$"); + + // --- 辅助方法 --- + + // 【修改】原 createRegexColumnRecognizer 替换为 createTestRegexRule + private SensitiveRule createTestRegexRule(Long id) { + SensitiveRule rule = new SensitiveRule(); + rule.setId(id); + rule.setType(SensitiveRuleType.REGEX); + rule.setEnabled(true); + rule.setLevel(SensitiveLevel.HIGH); + // 保留原作者的正则表达式 + rule.setDatabaseRegexExpression("^\\S+$"); + rule.setTableRegexExpression("^\\S+$"); + rule.setColumnRegexExpression("^\\S*email\\S*$"); + rule.setColumnCommentRegexExpression("^[\\S\\s]*email[\\S\\s]*$"); + return rule; } private DBTableColumn createDBTableColumn(String schemaName, String tableName, String columnName, String comment) { From ddfc830c319666fc324da98f6a303e1211f2a24a Mon Sep 17 00:00:00 2001 From: fenyf Date: Thu, 31 Jul 2025 15:10:30 +0800 Subject: [PATCH 02/10] feature(ai_recognition): Function optimization and code refactoring MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(ai_recognition): 解决查询接口缺少AI相关字段 fix(ai_recognition): 去除AI识别器冗余逻辑 feature(ai_recognition): 为扫描任务增加多线程和异步 feature(ai_recognition): 去除置信度阈值功能 feature(ai_recognition): 完善为AI识别结果指定默认算法功能 feature(ai_recognition): 中断扫描功能 feature(ai_recognition): 优化提示词拼接逻辑 feature(ai_recognition): 优化AI调用代码 feature(ai_recognition): 删除冗余识别模式 feature(ai_recognition): 优化提示词 --- .../V_4_3_4_20__alter_sensitive_rule.sql | 3 - .../v2/SensitiveColumnController.java | 7 + .../datasecurity/SensitiveRuleEntity.java | 3 - .../datasecurity/MaskingAlgorithmService.java | 15 ++ .../SensitiveColumnScanningTask.java | 197 ++++++++++++---- .../SensitiveColumnScanningTaskManager.java | 13 ++ .../datasecurity/SensitiveColumnService.java | 15 ++ .../datasecurity/SensitiveRuleService.java | 2 + .../odc/service/datasecurity/ai/AIConfig.java | 48 ++-- .../datasecurity/ai/AIInferenceService.java | 38 +-- .../odc/service/datasecurity/ai/AIParam.java | 30 +++ .../datasecurity/ai/PromptTemplateLoader.java | 99 ++------ .../factory/ScanningStrategyFactory.java | 2 - .../model/DefaultSensitiveType.java | 177 ++++++++++++++ .../datasecurity/model/RecognitionResult.java | 1 - .../datasecurity/model/ScanResult.java | 3 - .../datasecurity/model/ScanningModeType.java | 4 - .../SensitiveColumnScanningTaskInfo.java | 16 +- .../datasecurity/model/SensitiveRule.java | 3 - .../recognizer/AIColumnRecognizer.java | 156 +++++++------ .../strategy/RulesAndAiStrategy.java | 72 ------ ...nsitive_column_recognize_system_prompt.txt | 81 +++++++ ...itive_column_recognize_prompt_templete.txt | 46 ---- .../SensitiveColumnScannerTest.java | 202 ---------------- .../recognizer/AIColumnRecognizerTest.java | 218 ------------------ 25 files changed, 659 insertions(+), 792 deletions(-) create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java delete mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java create mode 100644 server/odc-service/src/main/resources/ai-prompt-template/sensitive_column_recognize_system_prompt.txt delete mode 100644 server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt delete mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java delete mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java diff --git a/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql index 5a1fd21059..163ce9a466 100644 --- a/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql +++ b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql @@ -2,9 +2,6 @@ alter table `data_security_sensitive_rule` add column `ai_sensitive_types` text default null comment 'A list of sensitive data types for AI rules, stored as a JSON array string.'; -alter table `data_security_sensitive_rule` - add column `ai_confidence_threshold` integer default 80 comment 'Confidence threshold for AI rules, with a value range of 0-100.'; - alter table `data_security_sensitive_rule` add column `ai_custom_prompt` text default null comment 'User-defined custom prompt for AI rules.'; diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java index 9fd98510b3..381760f830 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java @@ -143,4 +143,11 @@ public SuccessResponse getScanningResults(@Path return Responses.success(service.getScanningResults(projectId, taskId)); } + @ApiOperation(value = "stopScanning", notes = "Stop a sensitive column scanning task") + @RequestMapping(value = "/stopScanning", method = RequestMethod.POST) + public SuccessResponse stopScanning(@PathVariable Long projectId, + @RequestParam String taskId) { + return Responses.success(service.stopScanning(projectId, taskId)); + } + } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java b/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java index 8dedbf4e12..23cd038918 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java @@ -117,9 +117,6 @@ public class SensitiveRuleEntity { @Column(name = "ai_sensitive_types") private List aiSensitiveTypes; - @Column(name = "ai_confidence_threshold") - private Integer aiConfidenceThreshold; - @Column(name = "ai_custom_prompt") private String aiCustomPrompt; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java index 6a256fa0c8..479ed36dca 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java @@ -207,6 +207,21 @@ public Long getDefaultAlgorithmIdByOrganizationId(@NonNull Long organizationId) return entities.get(0).getId(); } + @SkipAuthorize("odc internal usages") + public Optional getAlgorithmIdByName(@NonNull String algorithmName, @NonNull Long organizationId) { + List entities = + algorithmRepository.findByNameAndOrganizationId(algorithmName, organizationId); + if (entities.isEmpty()) { + log.warn("No masking algorithm found with name: {} for organization: {}", algorithmName, organizationId); + return Optional.empty(); + } + if (entities.size() > 1) { + log.warn("Multiple masking algorithms found with name: {} for organization: {}, using the first one", + algorithmName, organizationId); + } + return Optional.of(entities.get(0).getId()); + } + @SkipAuthorize("odc internal usages") public List getMaskingAlgorithmsByOrganizationId(@NonNull Long organizationId) { return organizationId2Algorithms.get(organizationId); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java index 0220d35d7f..362be5f09a 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java @@ -23,6 +23,7 @@ import java.util.Optional; import java.util.Set; import java.util.concurrent.Callable; +import java.util.concurrent.CompletableFuture; import java.util.function.Function; import java.util.stream.Collectors; @@ -38,7 +39,10 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo.ScanningTaskStatus; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnType; import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.model.DefaultSensitiveType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; +import com.oceanbase.odc.service.common.util.SpringContextUtil; /** * @author gaoda.xy @@ -56,8 +60,8 @@ public class SensitiveColumnScanningTask implements Callable { private final Map ruleMap; public SensitiveColumnScanningTask(Database database, List rules, ScanningModeType scanningMode, - SensitiveColumnScanningTaskInfo taskInfo, List existsSensitiveColumns, - Map> table2Columns, Map> view2Columns) { + SensitiveColumnScanningTaskInfo taskInfo, List existsSensitiveColumns, + Map> table2Columns, Map> view2Columns) { this.database = database; // 【修改】接收扫描模式，并创建新的扫描器 this.scanningMode = scanningMode; @@ -76,9 +80,9 @@ public SensitiveColumnScanningTask(Database database, List rules, */ private String getColumnKey(DBTableColumn column) { return String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } @Override @@ -87,58 +91,95 @@ public Void call() { taskInfo.setStatus(ScanningTaskStatus.RUNNING); // 调用重构后的 scanColumns 方法 scanColumns(table2Columns, SensitiveColumnType.TABLE_COLUMN); + if (taskInfo.isCancelled()) { + return null; + } scanColumns(view2Columns, SensitiveColumnType.VIEW_COLUMN); - taskInfo.setStatus(ScanningTaskStatus.SUCCESS); } catch (Exception e) { - taskInfo.setStatus(ScanningTaskStatus.FAILED); - taskInfo.setErrorCode(ErrorCodes.Unexpected); - taskInfo.setErrorMsg(String.format("Error during sensitive column scanning on database=%s, reason=%s", + if (!taskInfo.isCancelled()) { + taskInfo.setStatus(ScanningTaskStatus.FAILED); + taskInfo.setErrorCode(ErrorCodes.Unexpected); + taskInfo.setErrorMsg(String.format("Error during sensitive column scanning on database=%s, reason=%s", database.getName(), e.getMessage())); - } finally { - taskInfo.setCompleteTime(new Date()); + taskInfo.setCompleteTime(new Date()); + } } return null; } - // 【修改】scanColumns 的核心逻辑改为批量扫描 + // 【修改】scanColumns 的核心逻辑改为批量扫描，并支持表级别并发 private void scanColumns(Map> object2Columns, SensitiveColumnType columnType) { - for (Map.Entry> entry : object2Columns.entrySet()) { - String objectName = entry.getKey(); - List columns = entry.getValue(); - - // 【改为批量扫描】一次性扫描整个表的所有列 - Map scanResults = this.scanner.scanBatch(columns, this.scanningMode); - - List sensitiveColumns = new ArrayList<>(); - for (DBTableColumn dbTableColumn : columns) { - String columnKey = getColumnKey(dbTableColumn); - ScanResult scanResult = scanResults.get(columnKey); - - if (scanResult != null) { - // 根据扫描模式获取最终的识别结果 - Optional finalResultOpt = scanResult.getFinalResult(this.scanningMode); - - // 如果最终有识别结果，则处理 - finalResultOpt.ifPresent(finalResult -> { - SensitiveColumnMeta meta = new SensitiveColumnMeta(database.getId(), objectName, - dbTableColumn.getName()); - if (!existsSensitiveColumns.contains(meta)) { - SensitiveColumn column = createSensitiveColumn(columnType, objectName, dbTableColumn, - finalResult); - sensitiveColumns.add(column); - existsSensitiveColumns.add(meta); + if (object2Columns.isEmpty()) { + return; + } + + // 表级别并发处理：为每个表创建异步任务 + List> tableFutures = object2Columns.entrySet().stream() + .map(entry -> CompletableFuture.runAsync(() -> { + String objectName = entry.getKey(); + List columns = entry.getValue(); + + try { + // 检查是否已被中断 + if (taskInfo.isCancelled()) { + return; + } + // 【改为批量扫描】一次性扫描整个表的所有列 + Map scanResults = this.scanner.scanBatch(columns, this.scanningMode); + + // 再次检查是否已被中断 + if (taskInfo.isCancelled()) { + return; + } + + List sensitiveColumns = new ArrayList<>(); + for (DBTableColumn dbTableColumn : columns) { + String columnKey = getColumnKey(dbTableColumn); + ScanResult scanResult = scanResults.get(columnKey); + + if (scanResult != null) { + // 根据扫描模式获取最终的识别结果 + Optional finalResultOpt = scanResult + .getFinalResult(this.scanningMode); + + // 如果最终有识别结果，则处理 + finalResultOpt.ifPresent(finalResult -> { + SensitiveColumnMeta meta = new SensitiveColumnMeta(database.getId(), objectName, + dbTableColumn.getName()); + // 使用同步块保证线程安全 + synchronized (existsSensitiveColumns) { + if (!existsSensitiveColumns.contains(meta)) { + SensitiveColumn column = createSensitiveColumn(columnType, objectName, + dbTableColumn, + finalResult); + sensitiveColumns.add(column); + existsSensitiveColumns.add(meta); + } + } + }); } - }); + } + // 批量添加敏感列结果，使用同步保证线程安全 + if (!sensitiveColumns.isEmpty()) { + taskInfo.addSensitiveColumns(sensitiveColumns); + } + taskInfo.addFinishedTableCount(); + } catch (Exception e) { + System.err.println("表 " + objectName + " 扫描失败: " + e.toString()); + e.printStackTrace(); + // 即使失败也要增加完成计数，避免任务卡住 + taskInfo.addFinishedTableCount(); } - } - taskInfo.addSensitiveColumns(sensitiveColumns); - taskInfo.addFinishedTableCount(); - } + })) + .collect(Collectors.toList()); + + // 等待所有表的扫描任务完成 + CompletableFuture.allOf(tableFutures.toArray(new CompletableFuture[0])).join(); } // 【新增】辅助方法，用于创建 SensitiveColumn 对象，使代码更清晰 private SensitiveColumn createSensitiveColumn(SensitiveColumnType columnType, String objectName, - DBTableColumn dbTableColumn, RecognitionResult result) { + DBTableColumn dbTableColumn, RecognitionResult result) { SensitiveColumn column = new SensitiveColumn(); column.setType(columnType); column.setDatabase(database); @@ -147,11 +188,73 @@ private SensitiveColumn createSensitiveColumn(SensitiveColumnType columnType, St // 从 RecognitionResult 获取 ruleId 和 level column.setSensitiveRuleId(result.getMatchedRuleId()); column.setLevel(result.getLevel()); - // 通过 ruleId 从我们保存的 ruleMap 中找到对应的规则，再获取脱敏算法ID + + // 设置脱敏算法ID + Long maskingAlgorithmId = determineMaskingAlgorithmId(result); + column.setMaskingAlgorithmId(maskingAlgorithmId); + + return column; + } + + /** + * 根据识别结果确定脱敏算法ID + * 对于AI识别的结果，如果是默认敏感类型则自动匹配同名脱敏算法，否则使用系统默认算法 + * 对于传统规则识别的结果，直接使用规则配置的脱敏算法 + */ + private Long determineMaskingAlgorithmId(RecognitionResult result) { + // 通过 ruleId 从我们保存的 ruleMap 中找到对应的规则 SensitiveRule matchedRule = this.ruleMap.get(result.getMatchedRuleId()); - if (matchedRule != null) { - column.setMaskingAlgorithmId(matchedRule.getMaskingAlgorithmId()); + if (matchedRule == null) { + // 如果找不到规则，使用系统默认算法 + return getSystemDefaultAlgorithmId(); + } + + // 如果是AI规则识别的结果，需要根据敏感类型自动匹配算法 + if (SensitiveRuleType.AI.equals(result.getSourceRuleType()) && result.getSensitiveType() != null) { + return handleAiRecognitionResult(result.getSensitiveType()); + } + + // 对于传统规则，直接使用规则配置的脱敏算法ID + return matchedRule.getMaskingAlgorithmId(); + } + + /** + * 处理AI识别结果的脱敏算法匹配 + */ + private Long handleAiRecognitionResult(String sensitiveType) { + // 判断是否为默认敏感类型 + if (DefaultSensitiveType.isDefaultType(sensitiveType)) { + // 通过DefaultSensitiveType获取算法名称，然后根据名称获取当前组织下的算法ID + Optional algorithmNameOpt = DefaultSensitiveType.getAlgorithmNameBySensitiveType(sensitiveType); + if (algorithmNameOpt.isPresent()) { + try { + MaskingAlgorithmService algorithmService = SpringContextUtil.getBean(MaskingAlgorithmService.class); + Optional algorithmIdOpt = algorithmService.getAlgorithmIdByName(algorithmNameOpt.get(), + database.getOrganizationId()); + if (algorithmIdOpt.isPresent()) { + return algorithmIdOpt.get(); + } + } catch (Exception e) { + System.err.println("Failed to get algorithm ID by name: " + e.getMessage()); + } + } + } + + // 如果不是默认类型或获取失败，使用系统默认算法 + return getSystemDefaultAlgorithmId(); + } + + /** + * 获取系统默认脱敏算法ID + */ + private Long getSystemDefaultAlgorithmId() { + try { + MaskingAlgorithmService algorithmService = SpringContextUtil.getBean(MaskingAlgorithmService.class); + return algorithmService.getDefaultAlgorithmIdByOrganizationId(database.getOrganizationId()); + } catch (Exception e) { + // 记录错误日志，但不抛出异常，避免影响整个扫描流程 + System.err.println("Failed to get default masking algorithm ID: " + e.getMessage()); + return null; } - return column; } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java index 1dc3543692..63cca895ab 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java @@ -134,4 +134,17 @@ public SensitiveColumnScanningTaskInfo get(String taskId) { return taskInfo; } + public boolean stop(String taskId) { + SensitiveColumnScanningTaskInfo taskInfo = cache.get(taskId); + if (taskInfo != null && taskInfo.getStatus() == ScanningTaskStatus.RUNNING) { + taskInfo.setCancelled(true); + taskInfo.setStatus(ScanningTaskStatus.CANCELLED); + taskInfo.setCompleteTime(new Date()); + log.info("Sensitive column scanning task stopped, taskId: {}", taskId); + return true; + } + return false; + } + } + diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java index f7129a8fd7..19fa165685 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java @@ -416,6 +416,21 @@ public SensitiveColumnScanningTaskInfo getScanningResults(@NotNull Long projectI return taskInfo; } + @Transactional(rollbackFor = Exception.class) + @PreAuthenticate(hasAnyResourceRole = {"OWNER, DBA, SECURITY_ADMINISTRATOR"}, + actions = {"OWNER", "DBA", "SECURITY_ADMINISTRATOR"}, resourceType = "ODC_PROJECT", + indexOfIdParam = 0) + @StatefulRoute(stateName = StateName.UUID_STATEFUL_ID, stateIdExpression = "#taskId") + public Boolean stopScanning(@NotNull Long projectId, @NotBlank String taskId) { + SensitiveColumnScanningTaskInfo taskInfo = scanningTaskManager.get(taskId); + if (!Objects.equals(taskInfo.getProjectId(), projectId)) { + String errorMsg = String.format("Sensitive column scanning task not exists, taskId=%s", taskId); + throw new NotFoundException(ErrorCodes.IllegalArgument, new Object[] {"taskId", errorMsg}, null); + } + return scanningTaskManager.stop(taskId); + } + + @SkipAuthorize("odc internal usages") public SensitiveColumnEntity nullSafeGet(@NotNull Long id) { return repository.findById(id) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveRuleService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveRuleService.java index 0bc0fb9b05..22b97dd0bd 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveRuleService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveRuleService.java @@ -160,6 +160,8 @@ public SensitiveRule update(@NotNull Long projectId, @NotNull Long id, @NotNull entity.setLevel(rule.getLevel()); entity.setMaskingAlgorithmId(algorithmId); entity.setDescription(rule.getDescription()); + entity.setAiSensitiveTypes(rule.getAiSensitiveTypes()); + entity.setAiCustomPrompt(rule.getAiCustomPrompt()); ruleRepository.saveAndFlush(entity); log.info("Sensitive rule has been updated, id={}, name={}", entity.getId(), entity.getName()); return detail(entity.getProjectId(), entity.getId()); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java index 28898c754c..e21a6afd17 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java @@ -15,11 +15,17 @@ */ package com.oceanbase.odc.service.datasecurity.ai; +import java.util.HashMap; +import java.util.Map; + import org.springframework.context.annotation.Bean; import org.springframework.stereotype.Component; import com.openai.client.OpenAIClient; import com.openai.client.okhttp.OpenAIOkHttpClient; +import com.openai.core.JsonBoolean; +import com.openai.core.JsonNumber; +import com.openai.core.JsonValue; import lombok.Data; @@ -30,36 +36,40 @@ @Component // @ConfigurationProperties(prefix = "datasecurity.ai") public class AIConfig { - /** - * 是否启用AI服务 - */ - private boolean enabled = true;; + private boolean enabled = true; - /** - * API 密钥 - */ //private String apiKey = "sk-c6bbbbde1b7e420b897d0662301c6d7c"; private String apiKey = "token-abc123"; - /** - * API 的基础URL - */ //private String baseUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1"; private String baseUrl = "http://172.25.17.78:8000/v1"; - /** - * 默认使用的模型名称 - */ - //private String model = "qwen2.5-3b-instruct"; + //private String model = "qwen3-8b"; private String model = "nlora"; + private Boolean enableThinking = AIParam.DEFAULT_ENABLE_THINKING; + + private Double temperature = AIParam.DEFAULT_TEMPERATURE; + + private Double topP = AIParam.DEFAULT_TOP_P; + + private Integer topK = AIParam.DEFAULT_TOP_K; + + private Integer minP = AIParam.DEFAULT_MIN_P; + + public Map loadAdditionalParams() { + Map params = new HashMap<>(); + params.put("enable_thinking", JsonBoolean.from(this.enableThinking)); + params.put("top_k", JsonNumber.from(this.topK)); + params.put("min_p", JsonNumber.from(this.minP)); + return params; + } + @Bean public OpenAIClient openAIClient() { return OpenAIOkHttpClient.builder() - .apiKey(this.apiKey) - .baseUrl(this.baseUrl) - .build(); + .apiKey(this.apiKey) + .baseUrl(this.baseUrl) + .build(); } - - } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java index 98342d582f..5265a12b20 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023 OceanBase. + * Copyright (c) 2025 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -29,28 +29,36 @@ @Service public class AIInferenceService { - private final AIConfig aiConfig; + private final AIConfig aiConfig; private final OpenAIClient openAIClient; - public AIInferenceService(AIConfig aiConfig) { + // 直接注入 AIConfig 和 OpenAIClient 两个Bean + public AIInferenceService(AIConfig aiConfig, OpenAIClient openAIClient) { this.aiConfig = aiConfig; - // 根据AIConfig创建OpenAIClient - this.openAIClient = aiConfig.openAIClient(); + this.openAIClient = openAIClient; } - public ChatCompletion chat(String prompt) { - //Map bodyParams = new HashMap<>(); - //bodyParams.put("enable_thinking", JsonBoolean.from(false)); - String model = aiConfig.getModel(); + /** + * 使用系统提示词和用户提示词分别调用AI服务 + * + * @param systemPrompt 系统提示词 + * @param userPrompt 用户提示词 + * @return AI响应 + */ + public ChatCompletion chat(String systemPrompt, String userPrompt) { + try { ChatCompletionCreateParams params = ChatCompletionCreateParams.builder() - .addUserMessage(prompt) - .model(model) - //.additionalBodyProperties(bodyParams) - .build(); + .addSystemMessage(systemPrompt) + .addUserMessage(userPrompt) + .model(aiConfig.getModel()) + .temperature(aiConfig.getTemperature()) + .topP(aiConfig.getTopP()) + .additionalBodyProperties(aiConfig.loadAdditionalParams()) + .build(); return openAIClient.chat().completions().create(params); } catch (Exception e) { - throw new RuntimeException("调用阿里云AI服务失败: " + e.getMessage(), e); + throw new RuntimeException("调用AI服务失败: " + e.getMessage(), e); } } -} +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java new file mode 100644 index 0000000000..3cd387fe7f --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java @@ -0,0 +1,30 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +public class AIParam { + + /** + * Default values for AI configuration + */ + public static final Boolean DEFAULT_ENABLE_THINKING = false; + public static final Double DEFAULT_TEMPERATURE = 0.1; + public static final Double DEFAULT_TOP_P = 1.0; + public static final Integer DEFAULT_TOP_K = 0; + public static final Integer DEFAULT_MIN_P = 0; + + public static final Integer DEFAULT_BATCH_SIZE_IN_TABLE = 30; +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java index d904de9220..92f84a33a5 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java @@ -36,116 +36,55 @@ @Component public class PromptTemplateLoader { - private static final String TEMPLATE_PATH = "/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt"; + private static final String SYSTEM_TEMPLATE_PATH = "/ai-prompt-template/sensitive_column_recognize_system_prompt.txt"; - // 定义新的三个占位符 - private static final String COLUMN_PLACEHOLDER = "{DBTableColumn}"; + // 定义占位符 private static final String TYPES_PLACEHOLDER = "{sensitiveTypes}"; private static final String PROMPT_PLACEHOLDER = "{customPrompt}"; - private String template; + private String systemTemplate; @PostConstruct public void init() { - try (var inputStream = PromptTemplateLoader.class.getResourceAsStream(TEMPLATE_PATH)) { + try (var inputStream = PromptTemplateLoader.class.getResourceAsStream(SYSTEM_TEMPLATE_PATH)) { if (Objects.isNull(inputStream)) { - throw new IllegalStateException("AI prompt template file not found: " + TEMPLATE_PATH); + throw new IllegalStateException("AI system prompt template file not found: " + SYSTEM_TEMPLATE_PATH); } try (var reader = new java.io.BufferedReader(new java.io.InputStreamReader(inputStream))) { - this.template = reader.lines().collect(Collectors.joining(System.lineSeparator())); + this.systemTemplate = reader.lines().collect(Collectors.joining(System.lineSeparator())); } } catch (Exception e) { // 在实际项目中，这里应该使用日志系统 e.printStackTrace(); - throw new IllegalStateException("Failed to load AI prompt template", e); + throw new IllegalStateException("Failed to load AI system prompt template", e); } } /** - * 【新】根据列元数据、敏感类型列表和用户提示，构建最终的 AI 提示词。 + * 构建系统提示词 * - * @param column 数据库表列的元数据对象 * @param sensitiveTypes 用户指定的敏感类型列表 (例如 ["联系方式", "身份信息"]) * @param customPrompt 用户为该规则自定义的补充说明提示 - * @return 填充了所有信息的完整提示词字符串 + * @return 填充了敏感类型和自定义提示的系统提示词字符串 */ - public String buildPrompt(DBTableColumn column, List sensitiveTypes, String customPrompt) { - if (this.template == null || this.template.isEmpty()) { - throw new IllegalStateException("Prompt template is not available. Check loading status."); - } - if (column == null) { - throw new IllegalArgumentException("Input column cannot be null."); - } - - // 1. 格式化列元数据 - String formattedColumn = formatColumnMetadata(column); - - // 2. 格式化敏感类型列表 - String formattedTypes = (sensitiveTypes == null || sensitiveTypes.isEmpty()) - ? "None specified" - : String.join(", ", sensitiveTypes); - - // 3. 格式化用户自定义提示 - String formattedPrompt = (customPrompt == null || customPrompt.trim().isEmpty()) - ? "None" - : customPrompt; - - // 4. 依次替换模板中的三个占位符 - return this.template - .replace(COLUMN_PLACEHOLDER, formattedColumn) - .replace(TYPES_PLACEHOLDER, formattedTypes) - .replace(PROMPT_PLACEHOLDER, formattedPrompt); - } - - /** - * 根据列元数据列表、敏感类型列表和用户提示，构建批量处理的 AI 提示词 - * - * @param columnsJson 数据库表列的元数据对象列表的JSON字符串 - * @param sensitiveTypes 用户指定的敏感类型列表 (例如 ["联系方式", "身份信息"]) - * @param customPrompt 用户为该规则自定义的补充说明提示 - * @return 填充了所有信息的完整提示词字符串 - */ - public String buildPrompt(String columnsJson, List sensitiveTypes, String customPrompt) { - if (this.template == null || this.template.isEmpty()) { - throw new IllegalStateException("Prompt template is not available. Check loading status."); - } - if (columnsJson == null) { - throw new IllegalArgumentException("Input columnsJson cannot be null."); + public String buildSystemPrompt(List sensitiveTypes, String customPrompt) { + if (this.systemTemplate == null || this.systemTemplate.isEmpty()) { + throw new IllegalStateException("System prompt template is not available. Check loading status."); } // 1. 格式化敏感类型列表 String formattedTypes = (sensitiveTypes == null || sensitiveTypes.isEmpty()) - ? "None specified" - : String.join(", ", sensitiveTypes); + ? "No specified category." + : String.join(", ", sensitiveTypes); // 2. 格式化用户自定义提示 String formattedPrompt = (customPrompt == null || customPrompt.trim().isEmpty()) - ? "None" - : customPrompt; + ? "No supplementary rule." + : customPrompt; // 3. 替换模板中的占位符 - return this.template - .replace(COLUMN_PLACEHOLDER, columnsJson) - .replace(TYPES_PLACEHOLDER, formattedTypes) - .replace(PROMPT_PLACEHOLDER, formattedPrompt); + return this.systemTemplate + .replace(TYPES_PLACEHOLDER, formattedTypes) + .replace(PROMPT_PLACEHOLDER, formattedPrompt); } - - /** - * 将列的元数据格式化为对 AI 模型友好的字符串。 - * (此方法逻辑不变) - */ - private String formatColumnMetadata(DBTableColumn column) { - StringBuilder contextBuilder = new StringBuilder(); - contextBuilder.append("Schema Name: ").append(formatValue(column.getSchemaName())).append("\n"); - contextBuilder.append("Table Name: ").append(formatValue(column.getTableName())).append("\n"); - contextBuilder.append("Column Name: ").append(formatValue(column.getName())).append("\n"); - contextBuilder.append("Data Type: ").append(formatValue(column.getTypeName())).append("\n"); - contextBuilder.append("Column Comment: ").append(formatValue(column.getComment())); - return contextBuilder.toString(); - } - - private String formatValue(String value) { - return (value == null || value.trim().isEmpty()) ? "N/A" : value; - } - } \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java index 5bf9aed38d..4ea2e3c6a9 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java @@ -25,7 +25,6 @@ import com.oceanbase.odc.service.datasecurity.model.ScanResult; import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; import com.oceanbase.odc.service.datasecurity.strategy.JointRecognitionStrategy; -import com.oceanbase.odc.service.datasecurity.strategy.RulesAndAiStrategy; import com.oceanbase.odc.service.datasecurity.strategy.RulesOnlyStrategy; import com.oceanbase.odc.service.datasecurity.strategy.ScanningStrategy; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; @@ -46,7 +45,6 @@ public ScanningStrategyFactory() { // 预创建所有策略实例 strategies.put(ScanningModeType.RULES_ONLY, new RulesOnlyStrategy()); strategies.put(ScanningModeType.JOINT_RECOGNITION, new JointRecognitionStrategy()); - strategies.put(ScanningModeType.RULES_AND_AI, new RulesAndAiStrategy()); } /** diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java new file mode 100644 index 0000000000..f1b63279da --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java @@ -0,0 +1,177 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.model; + +import java.util.Arrays; +import java.util.Optional; + +/** + * 默认敏感类型枚举 + * 定义AI识别的13种默认敏感类型，每种类型对应一个同名的脱敏算法 + * 支持多语言匹配和模糊匹配，提高AI识别结果的容错率 + * + * @author AI Assistant + * @date 2024/01/01 + */ +public enum DefaultSensitiveType { + + /** + * 个人姓名（汉字类型） + */ + PERSONAL_NAME_CHINESE("${com.oceanbase.odc.builtin-resource.masking-algorithm.personal-name-chinese.name}"), + + /** + * 个人姓名（字母类型） + */ + PERSONAL_NAME_ALPHABET("${com.oceanbase.odc.builtin-resource.masking-algorithm.personal-name-alphabet.name}"), + + /** + * 昵称 + */ + NICKNAME("${com.oceanbase.odc.builtin-resource.masking-algorithm.nickname.name}"), + + /** + * 邮箱 + */ + EMAIL("${com.oceanbase.odc.builtin-resource.masking-algorithm.email.name}"), + + /** + * 地址 + */ + ADDRESS("${com.oceanbase.odc.builtin-resource.masking-algorithm.address.name}"), + + /** + * 手机号码 + */ + PHONE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.phone-number.name}"), + + /** + * 固定电话 + */ + FIXED_LINE_PHONE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.fixed-line-phone-number.name}"), + + /** + * 证件号码 + */ + CERTIFICATE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.certificate-number.name}"), + + /** + * 银行卡号 + */ + BANK_CARD_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.bank-card-number.name}"), + + /** + * 车牌号 + */ + LICENSE_PLATE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.license-plate-number.name}"), + + /** + * 设备唯一识别号 + */ + DEVICE_ID("${com.oceanbase.odc.builtin-resource.masking-algorithm.device-id.name}"), + + /** + * IP地址 + */ + IP("${com.oceanbase.odc.builtin-resource.masking-algorithm.ip.name}"), + + /** + * MAC地址 + */ + MAC("${com.oceanbase.odc.builtin-resource.masking-algorithm.mac.name}"); + + private final String algorithmName; + + DefaultSensitiveType(String algorithmName) { + this.algorithmName = algorithmName; + } + + public String getAlgorithmName() { + return algorithmName; + } + + /** + * 判断给定的敏感类型是否为默认类型（智能匹配） + * + * @param sensitiveType 敏感类型名称 + * @return 如果是默认类型返回true，否则返回false + */ + public static boolean isDefaultType(String sensitiveType) { + return findBestMatch(sensitiveType).isPresent(); + } + + /** + * 根据敏感类型名称获取对应的脱敏算法名称（智能匹配） + * + * @param sensitiveType 敏感类型名称 + * @return 对应的脱敏算法名称，如果不是默认类型则返回空 + */ + public static Optional getAlgorithmNameBySensitiveType(String sensitiveType) { + return findBestMatch(sensitiveType).map(DefaultSensitiveType::getAlgorithmName); + } + + /** + * 根据敏感类型名称获取对应的枚举值（智能匹配） + * + * @param sensitiveType 敏感类型名称 + * @return 对应的枚举值，如果不是默认类型则返回空 + */ + public static Optional getByDisplayName(String sensitiveType) { + return findBestMatch(sensitiveType); + } + + /** + * 智能匹配敏感类型（支持多种匹配策略） + * + * @param sensitiveType 敏感类型名称 + * @return 匹配到的枚举值 + */ + public static Optional findBestMatch(String sensitiveType) { + if (sensitiveType == null || sensitiveType.trim().isEmpty()) { + return Optional.empty(); + } + + String normalized = sensitiveType.toLowerCase().trim(); + + // 精确匹配枚举名称（下划线格式） + for (DefaultSensitiveType type : values()) { + if (type.name().toLowerCase().equals(normalized)) { + return Optional.of(type); + } + } + + // 精确匹配连字符格式（AI返回的格式） + for (DefaultSensitiveType type : values()) { + String hyphenFormat = type.name().toLowerCase().replace("_", "-"); + if (hyphenFormat.equals(normalized)) { + return Optional.of(type); + } + } + + // 模糊匹配：检查是否包含关键词 + for (DefaultSensitiveType type : values()) { + String enumName = type.name().toLowerCase(); + String hyphenFormat = enumName.replace("_", "-"); + if (enumName.contains(normalized) || normalized.contains(enumName.replace("_", "")) || + hyphenFormat.contains(normalized) || normalized.contains(hyphenFormat.replace("-", ""))) { + return Optional.of(type); + } + } + + return Optional.empty(); + } + +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java index 17c1dbdff8..4d5254a00e 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java @@ -31,5 +31,4 @@ public class RecognitionResult { // AI 规则 private String sensitiveType; // AI 判断出的具体敏感类型 - private Integer confidence; // AI 的置信度 } \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java index 44078f34fb..f0cc9c0263 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java @@ -43,9 +43,6 @@ public Optional getFinalResult(ScanningModeType scanningMode) case JOINT_RECOGNITION: // 对于联合识别，Scanner已经做过决策，直接返回存在的那个结果 return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; - case RULES_AND_AI: - // 对于差异化展示模式，可以根据业务需求调整优先级策略 - return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; default: return Optional.empty(); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java index 3ce04a4be6..33a2740272 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java @@ -23,9 +23,5 @@ public enum ScanningModeType { RULES_ONLY, - RULES_AND_AI, - JOINT_RECOGNITION, - - } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java index 9c38d88b29..03ac7eef62 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java @@ -41,9 +41,10 @@ public class SensitiveColumnScanningTaskInfo { private Date completeTime; private ErrorCode errorCode; private String errorMsg; + private volatile boolean cancelled = false; public SensitiveColumnScanningTaskInfo(@NonNull String taskId, @NonNull Long projectId, - @NonNull Integer allTableCount) { + @NonNull Integer allTableCount) { this.taskId = taskId; this.projectId = projectId; this.status = ScanningTaskStatus.CREATED; @@ -81,14 +82,23 @@ public synchronized void setErrorMsg(String msg) { this.errorMsg = msg; } + public synchronized void setCancelled(boolean cancelled) { + this.cancelled = cancelled; + } + + public boolean isCancelled() { + return this.cancelled; + } + public enum ScanningTaskStatus { CREATED, RUNNING, SUCCESS, - FAILED; + FAILED, + CANCELLED; public boolean isCompleted() { - return this == SUCCESS || this == FAILED; + return this == SUCCESS || this == FAILED || this == CANCELLED; } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java index dec6e18bfa..b812eb59d5 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveRule.java @@ -53,7 +53,6 @@ public class SensitiveRule implements SecurityResource, SingleOrganizationResour @JsonProperty(access = Access.READ_ONLY) private Long projectId; - @NotNull private SensitiveRuleType type; private String databaseRegexExpression; @@ -72,8 +71,6 @@ public class SensitiveRule implements SecurityResource, SingleOrganizationResour private List aiSensitiveTypes = new ArrayList<>(); - private Integer aiConfidenceThreshold; - private String aiCustomPrompt; @NotNull diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index 8e58c838cb..08e8b2cf63 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -30,6 +30,7 @@ import com.google.common.collect.Lists; import com.oceanbase.odc.service.common.util.SpringContextUtil; import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; +import com.oceanbase.odc.service.datasecurity.ai.AIParam; import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; @@ -46,13 +47,10 @@ public class AIColumnRecognizer implements ColumnRecognizer { private final SensitiveRule aiRule; // 直接保存整个规则对象 // @Value() - private static final int BATCH_SIZE = 50; // 这个批次大小需要能全局设置 + private static final int BATCH_SIZE = AIParam.DEFAULT_BATCH_SIZE_IN_TABLE; // 单表内列数超过此值时进行分批处理 private static final ObjectMapper objectMapper = new ObjectMapper(); // 用于解析JSON - // private static final Pattern JSON_PATTERN = - // Pattern.compile("(?s)```json\\s*(\\{.*\\})\\s*```|(\\{.*\\})"); - //// 匹配 {...} 或 [...] private static final Pattern JSON_PATTERN = Pattern - .compile("(?s)```json\\s*([\\{\\[].*[\\}\\]])\\s*```|([\\{\\[].*[\\}\\]])"); + .compile("(?s)```json\\s*([\\{\\[].*[\\}\\]])\\s*```|([\\{\\[].*[\\}\\]])"); public AIColumnRecognizer(SensitiveRule rule) { this.aiRule = rule; @@ -75,74 +73,95 @@ public Map> recognizeBatch(List> batches = Lists.partition(columns, BATCH_SIZE); Map> finalAiResults = new HashMap<>(); - try { - // 3. 遍历每一个小批次，分别调用 AI - for (List batch : batches) { - // a. 为这个小批次构建 Prompt - String prompt = buildBatchPrompt(promptTemplateLoader, batch); - // b. 调用 AI - ChatCompletion completion = aiService.chat(prompt); - String rawContent = completion.choices().get(0).message().content().orElse("[]"); - - // c. 使用正则表达式从AI的返回结果中安全地提取JSON数组字符串 - Matcher matcher = JSON_PATTERN.matcher(rawContent); - String jsonArrayResponse = "[]"; // 提供一个安全的默认值，以防匹配失败 - if (matcher.find()) { - // group(1) 对应被 ```json [...] ``` 包裹的内容, group(2) 对应裸露的 [...] - // 使用 Optional 来优雅地处理可能为null的捕获组 - jsonArrayResponse = Optional.ofNullable(matcher.group(1)).orElse(matcher.group(2)); + // 2. 如果列数超过批次大小，则分批处理；否则直接处理 + if (columns.size() > BATCH_SIZE) { + List> batches = Lists.partition(columns, BATCH_SIZE); + try { + // 3. 遍历每一个小批次，分别调用 AI + for (List batch : batches) { + processBatch(batch, promptTemplateLoader, aiService, finalAiResults); } + } catch (Exception e) { + e.printStackTrace(); + return finalAiResults; + } + } else { + // 直接处理单批次 + try { + processBatch(columns, promptTemplateLoader, aiService, finalAiResults); + } catch (Exception e) { + e.printStackTrace(); + return finalAiResults; + } + } + return finalAiResults; + } - // d. 解析提取出的、更纯净的 JSON 数组 - List batchResults = objectMapper.readValue(jsonArrayResponse, - new TypeReference>() { - }); - - // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 - int maxIndex = Math.min(batch.size(), batchResults.size()); - for (int i = 0; i < maxIndex; i++) { - DBTableColumn column = batch.get(i); - String columnKey = getColumnKey(column); - AiResponseDto dto = batchResults.get(i); - - if (dto.isSensitive()) { - RecognitionResult result = RecognitionResult.builder() - .matched(true) - .matchedRuleId(this.aiRule.getId()) - .level(dto.getRiskLevel()) - .sourceRuleType(SensitiveRuleType.AI) - .sensitiveType(dto.getSensitiveType()) - .confidence(dto.getConfidence()) - .build(); - finalAiResults.put(columnKey, Optional.of(result)); - } else { - finalAiResults.put(columnKey, Optional.empty()); - } - } + /** + * 处理单个批次的列数据 + */ + private void processBatch(List batch, PromptTemplateLoader promptTemplateLoader, + AIInferenceService aiService, Map> finalAiResults) throws IOException { + // a. 构建系统提示词 + String systemPrompt = promptTemplateLoader.buildSystemPrompt(aiRule.getAiSensitiveTypes(), aiRule.getAiCustomPrompt()); + + // b. 构建用户提示词（列数据的JSON数组） + String userPrompt = buildUserPrompt(batch); + + + // c. 调用 AI + ChatCompletion completion = aiService.chat(systemPrompt, userPrompt); + String rawContent = completion.choices().get(0).message().content().orElse("[]"); + + // c. 使用正则表达式从AI的返回结果中安全地提取JSON数组字符串 + Matcher matcher = JSON_PATTERN.matcher(rawContent); + String jsonArrayResponse = "[]"; // 提供一个安全的默认值，以防匹配失败 + if (matcher.find()) { + // group(1) 对应被 ```json [...] ``` 包裹的内容, group(2) 对应裸露的 [...] + // 使用 Optional 来优雅地处理可能为null的捕获组 + jsonArrayResponse = Optional.ofNullable(matcher.group(1)).orElse(matcher.group(2)); + } - // 如果AI返回的结果数量与输入不匹配，记录警告信息 - if (batchResults.size() != batch.size()) { - System.err.println("警告: AI返回结果数量(" + batchResults.size() + - ")与输入列数量(" + batch.size() + ")不匹配"); - } + // d. 解析提取出的、更纯净的 JSON 数组 + List batchResults = objectMapper.readValue(jsonArrayResponse, + new TypeReference>() { + }); + + + // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 + int maxIndex = Math.min(batch.size(), batchResults.size()); + for (int i = 0; i < maxIndex; i++) { + DBTableColumn column = batch.get(i); + String columnKey = getColumnKey(column); + AiResponseDto dto = batchResults.get(i); + + if (dto.isSensitive()) { + RecognitionResult result = RecognitionResult.builder() + .matched(true) + .matchedRuleId(this.aiRule.getId()) + .level(dto.getRiskLevel()) + .sourceRuleType(SensitiveRuleType.AI) + .sensitiveType(dto.getSensitiveCategory()) + .build(); + finalAiResults.put(columnKey, Optional.of(result)); + } else { + finalAiResults.put(columnKey, Optional.empty()); } - } catch (Exception e) { - // 在实际项目中，应使用日志系统记录详细错误 - e.printStackTrace(); - // 出现异常时返回当前已成功识别的结果，或返回空map - return finalAiResults; } - // 4. 返回包含所有 AI 识别结果的完整 Map - return finalAiResults; + + // 如果AI返回的结果数量与输入不匹配，记录警告信息 + if (batchResults.size() != batch.size()) { + System.err.println("警告: AI返回结果数量(" + batchResults.size() + + ")与输入列数量(" + batch.size() + ")不匹配"); + } } - // 构建批量 Prompt 的新逻辑 - private String buildBatchPrompt(PromptTemplateLoader promptTemplateLoader, List batch) - throws IOException { + /** + * 构建用户提示词（列数据的JSON数组） + */ + private String buildUserPrompt(List batch) throws IOException { // 将一批列的元数据转换为 JSON 数组字符串 List> columnMetadataList = batch.stream().map(c -> { Map meta = new HashMap<>(); @@ -153,19 +172,14 @@ private String buildBatchPrompt(PromptTemplateLoader promptTemplateLoader, List< meta.put("dataType", c.getTypeName()); return meta; }).collect(Collectors.toList()); - String columnsJson = objectMapper.writeValueAsString(columnMetadataList); - // 调用 PromptTemplateLoader 来填充 - return promptTemplateLoader.buildPrompt(columnsJson, aiRule.getAiSensitiveTypes(), aiRule.getAiCustomPrompt()); + return objectMapper.writeValueAsString(columnMetadataList); } - // 移除私有方法，直接使用接口的默认实现 - // 用于承载 AI 返回的 JSON 数据的内部类 @Data private static class AiResponseDto { private boolean sensitive; private SensitiveLevel riskLevel; - private Integer confidence; - private String sensitiveType; + private String sensitiveCategory; } } \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java deleted file mode 100644 index ac441b2b76..0000000000 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesAndAiStrategy.java +++ /dev/null @@ -1,72 +0,0 @@ -/* - * Copyright (c) 2025 OceanBase. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -package com.oceanbase.odc.service.datasecurity.strategy; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.Optional; -import java.util.concurrent.CompletableFuture; - -import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; -import com.oceanbase.odc.service.datasecurity.model.ScanResult; -import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; -import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - -/** - * 规则+AI策略实现 - * 同时执行基础规则和AI识别，用于差异化展示两种结果 - * - * @author Assistant - * @date 2025/1/27 - */ -public class RulesAndAiStrategy extends AbstractScanningStrategy { - - @Override - public ScanResult scan(DBTableColumn column, List basicRecognizers, - List aiRecognizers) { - Optional basicResult = findFirstMatch(basicRecognizers, column); - Optional aiResult = findFirstMatch(aiRecognizers, column); - return new ScanResult(basicResult, aiResult); - } - - @Override - public Map scanBatch(List columns, List basicRecognizers, - List aiRecognizers) { - // 并行执行基础规则和AI识别 - CompletableFuture>> basicFuture = CompletableFuture - .supplyAsync(() -> findAllFirstMatches(basicRecognizers, columns)); - CompletableFuture>> aiFuture = CompletableFuture - .supplyAsync(() -> findAllFirstMatches(aiRecognizers, columns)); - - // 等待两个任务完成并合并结果 - CompletableFuture.allOf(basicFuture, aiFuture).join(); - - Map> basicResults = basicFuture.join(); - Map> aiResults = aiFuture.join(); - - Map results = new HashMap<>(); - for (DBTableColumn column : columns) { - String columnKey = getColumnKey(column); - Optional basicResult = basicResults.getOrDefault(columnKey, Optional.empty()); - Optional aiResult = aiResults.getOrDefault(columnKey, Optional.empty()); - results.put(columnKey, new ScanResult(basicResult, aiResult)); - } - - return results; - } -} \ No newline at end of file diff --git a/server/odc-service/src/main/resources/ai-prompt-template/sensitive_column_recognize_system_prompt.txt b/server/odc-service/src/main/resources/ai-prompt-template/sensitive_column_recognize_system_prompt.txt new file mode 100644 index 0000000000..aeec7f1e50 --- /dev/null +++ b/server/odc-service/src/main/resources/ai-prompt-template/sensitive_column_recognize_system_prompt.txt @@ -0,0 +1,81 @@ +You are a data classification expert specializing in data privacy and security. Your task is to determine if a database column's metadata likely belongs to sensitive data . + +Analyze the JSON array of database columns that I will provide in your next message, based on the following criteria: + +**1. Input data structure:** +You will receive a JSON array, where each object represents a database column. Each column originates from a table in the actual backend business. The column information includes schemaName, tableName, columnName, comment, and dataType. + +**2. Sensitive Categories Check Rule:** +If there is "No specified category", you should independently determine whether it is sensitive. In this case, the sensitiveCategory should return "null". If there are some specified categories, you MUST and can ONLY select categories from the following "Sensitive Categories to Check" list as the identification result. If the sensitive columns you identified do not fall within the provided categories, simply classify them as non-sensitive. +Sensitive Categories to Check: {sensitiveTypes} + +**3. Supplementary Rule:** +This task may involve a "Supplementary Recognition Rule". If it exists, please ensure that the content of "Supplementary Recognition Rule" is given the highest priority. If not, "No supplementary rule" will be displayed here. Please ignore this item. +Supplementary Recognition Rule: {customPrompt} + +**CRITICAL CONSTRAINT - STRICTLY ENFORCE:** + +1. You are STRICTLY FORBIDDEN from outputting any sensitive type that is NOT explicitly listed in the "Sensitive Categories to Check" section above. +If a column is sensitive but doesn't match ANY of the specified categories, you MUST set "sensitiveCategory" to null. +Violating this constraint will result in system errors. + +2. Based on all the information above, return a JSON array where each object corresponds to a column from the input array in the SAME order. +For each column, respond with a JSON object in the following format ONLY. Do not add any other text or explanations. +The final output should be a single, valid JSON array. Please do not output Markdown code blocks, such as ```json. + +3. The number of columns in your output MUST be exactly equal to the number of columns in the input. +Do not omit any columns, and do not add extra columns. Each input column must have a corresponding output object in the same order. + +Format for each object in the output array: +{ + "sensitive": boolean, + "riskLevel": "HIGH" | "MEDIUM" | "LOW", + "sensitiveCategory": string (MUST be one of the categories listed in "Sensitive Categories to Check" section, or null) +} + +**Example:** +Assuming the specified categories include "address, phone-number, email, ip and fixed-line-phone-number" and input json array is: +[ + { + "schemaName": "education_db", + "tableName": "students", + "columnName": "student_address" + "comment": "The student's home address" + "dataType": "TEXT" + }, + { + "schemaName": "education_db", + "tableName": "students", + "columnName": "student_home_phone" + "comment": "The student's fixed home phone" + "dataType": "VARCHAR(20)" + }, + { + "schemaName": "education_db", + "tableName": "students", + "columnName": "math_grade" + "comment": "The students' math exam scores" + "dataType": "TINYINT" + } +] + +The expected output format for 3 columns: +[ + { + "sensitive": true, + "riskLevel": "HIGH", + "sensitiveCategory": "address" + }, + { + "sensitive": true, + "riskLevel": "MEDIUM", + "sensitiveCategory": "fixed-line-phone-number" + }, + { + "sensitive": false, + "riskLevel": "LOW", + "sensitiveCategory": null + } +] + + diff --git a/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt b/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt deleted file mode 100644 index b38b562000..0000000000 --- a/server/odc-service/src/main/resources/ai-prompt-templete/sensitive_column_recognize_prompt_templete.txt +++ /dev/null @@ -1,46 +0,0 @@ -You are a data classification expert specializing in data privacy and security. Your task is to determine if a database column's content likely belongs to one of the user-specified sensitive data categories. - -Analyze the JSON array of database columns provided below. For EACH column in the array, determine if it belongs to any of the specified sensitive data categories. - -1. **Sensitive Categories to Check:** -{sensitiveTypes} - -2. **JSON Array of Columns to Analyze:** -{DBTableColumn} - -3. **Additional User-Provided Hint:** -{customPrompt} - -Based on all the information above, return a JSON array where each object corresponds to a column from the input array in the SAME order. -For each column, respond with a JSON object in the following format ONLY. Do not add any other text or explanations. -The final output should be a single, valid JSON array. Please do not output Markdown code blocks, such as ```json. - -Example of expected output format for 3 columns: -[ - { - "sensitive": true, - "riskLevel": "HIGH", - "confidence": 95, - "sensitiveType": "Financial Information" - }, - { - "sensitive": true, - "riskLevel": "MEDIUM", - "confidence": 80, - "sensitiveType": "Personal Identification" - }, - { - "sensitive": false, - "riskLevel": "LOW", - "confidence": 30, - "sensitiveType": null - } -] - -Format for each object in the output array: -{ - "sensitive": boolean, - "riskLevel": "HIGH" | "MEDIUM" | "LOW", - "confidence": number (an int value between 0 and 100), - "sensitiveType": string (the specific category it belongs to, or null) -} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java deleted file mode 100644 index 8bd122b809..0000000000 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScannerTest.java +++ /dev/null @@ -1,202 +0,0 @@ -/* - * Copyright (c) 2025 OceanBase. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.oceanbase.odc.service.datasecurity; - -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import java.util.Map; - -import org.junit.Assert; -import org.junit.Before; -import org.junit.Test; - -import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; -import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; -import com.oceanbase.odc.service.datasecurity.model.ScanResult; -import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; -import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; -import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - -/** - * 敏感列扫描器单元测试类 - */ -public class SensitiveColumnScannerTest { - - private List rules; - private SensitiveColumnScanner scanner; - - @Before - public void setUp() { - // 创建测试规则 - SensitiveRule regexRule = createTestRegexRule(1L); - SensitiveRule aiRule = createTestAIRule(2L); - rules = Arrays.asList(regexRule, aiRule); - ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); - scanner = new SensitiveColumnScanner(rules, strategyFactory); - } - - @Test - public void scan_rulesOnlyMode_returnBasicResult() { - // 准备阶段：创建能被正则规则匹配的测试列 - DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); - - // 执行阶段：使用仅规则模式扫描 - ScanResult result = scanner.scan(column, ScanningModeType.RULES_ONLY); - - // 验证阶段：确认返回基础规则结果 - Assert.assertNotNull("结果不应为null", result); - Assert.assertTrue("基础规则结果应存在", result.getBasicRuleResult().isPresent()); - Assert.assertFalse("AI规则结果应不存在", result.getAiRuleResult().isPresent()); - - RecognitionResult basicResult = result.getBasicRuleResult().get(); - Assert.assertTrue("应匹配成功", basicResult.isMatched()); - Assert.assertEquals("匹配的规则ID应为1", Long.valueOf(1), basicResult.getMatchedRuleId()); - } - - @Test - public void scan_jointRecognitionMode_ruleMatched_returnBasicResultOnly() { - // 准备阶段：创建能被正则规则匹配的测试列 - DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); - - // 执行阶段：使用联合识别模式扫描 - ScanResult scanResult = scanner.scan(column, ScanningModeType.JOINT_RECOGNITION); - - // 验证阶段：确认只返回基础规则结果 - Assert.assertNotNull("结果不应为null", scanResult); - Assert.assertTrue("基础规则结果应存在", scanResult.getBasicRuleResult().isPresent()); - Assert.assertFalse("AI规则结果应不存在", scanResult.getAiRuleResult().isPresent()); - } - - @Test - public void scan_jointRecognitionMode_ruleNotMatched_returnAiResult() { - // 准备阶段：创建不能被正则规则匹配的测试列 - DBTableColumn column = createTestColumn("test_db", "products", "name", "product name"); - - // 执行阶段：使用联合识别模式扫描 - ScanResult result = scanner.scan(column, ScanningModeType.JOINT_RECOGNITION); - - // 验证阶段：确认基础规则结果不存在（因为没有mock AI识别器，所以AI结果也不存在） - Assert.assertNotNull("结果不应为null", result); - Assert.assertFalse("基础规则结果应不存在", result.getBasicRuleResult().isPresent()); - } - - @Test - public void scan_rulesAndAiMode_returnBothResults() { - // 准备阶段：创建能被正则规则匹配的测试列 - DBTableColumn column = createTestColumn("test_db", "users", "email", "user email address"); - - // 执行阶段：使用规则+AI模式扫描 - ScanResult result = scanner.scan(column, ScanningModeType.RULES_AND_AI); - - // 验证阶段：确认返回基础规则结果（AI结果不存在因为没有mock） - Assert.assertNotNull("结果不应为null", result); - Assert.assertTrue("基础规则结果应存在", result.getBasicRuleResult().isPresent()); - } - - @Test - public void scanBatch_emptyList_returnEmptyMap() { - // 执行阶段：使用仅规则模式批量扫描空列表 - Map result = scanner.scanBatch(Collections.emptyList(), ScanningModeType.RULES_ONLY); - - // 验证阶段：确认返回空Map - Assert.assertNotNull("结果不应为null", result); - Assert.assertTrue("空列表应返回空Map", result.isEmpty()); - } - - @Test - public void scanBatch_rulesOnlyMode_returnBasicResults() { - // 准备阶段：创建多个测试列，其中一列能被正则规则匹配 - DBTableColumn column1 = createTestColumn("test_db", "users", "email", "user email address"); // 能匹配 - DBTableColumn column2 = createTestColumn("test_db", "products", "name", "product name"); // 不能匹配 - List columns = Arrays.asList(column1, column2); - - // 执行阶段：使用仅规则模式批量扫描 - Map results = scanner.scanBatch(columns, ScanningModeType.RULES_ONLY); - - // 验证阶段：确认返回正确的批量扫描结果 - Assert.assertNotNull("结果不应为null", results); - Assert.assertEquals("应返回两个扫描结果", 2, results.size()); - Assert.assertTrue("应包含users.email的结果", results.containsKey("users.email")); - Assert.assertTrue("应包含products.name的结果", results.containsKey("products.name")); - - // 验证email列的扫描结果（应该匹配） - ScanResult emailResult = results.get("users.email"); - Assert.assertTrue("email的基础规则结果应存在", emailResult.getBasicRuleResult().isPresent()); - Assert.assertEquals("email匹配的规则ID应为1", Long.valueOf(1), emailResult.getBasicRuleResult().get().getMatchedRuleId()); - - // 验证name列的扫描结果（不应该匹配） - ScanResult nameResult = results.get("products.name"); - Assert.assertFalse("name的基础规则结果应不存在", nameResult.getBasicRuleResult().isPresent()); - } - - // --- 辅助方法 --- - - /** - * 创建测试用的正则规则 - * @param id 规则ID - * @return 正则规则对象 - */ - private SensitiveRule createTestRegexRule(Long id) { - SensitiveRule rule = new SensitiveRule(); - rule.setId(id); - rule.setType(SensitiveRuleType.REGEX); - rule.setEnabled(true); - rule.setLevel(SensitiveLevel.HIGH); - rule.setDatabaseRegexExpression("^\\S+$"); - rule.setTableRegexExpression("^\\S+$"); - rule.setColumnRegexExpression("^\\S*email\\S*$"); - rule.setColumnCommentRegexExpression("^[\\S\\s]*email[\\S\\s]*$"); - return rule; - } - - /** - * 创建测试用的AI规则 - * @param id 规则ID - * @return AI规则对象 - */ - private SensitiveRule createTestAIRule(Long id) { - SensitiveRule rule = new SensitiveRule(); - rule.setId(id); - rule.setType(SensitiveRuleType.AI); - rule.setEnabled(true); - rule.setLevel(SensitiveLevel.HIGH); - rule.setAiSensitiveTypes(Arrays.asList("联系方式", "财务信息", "身份信息")); - rule.setAiConfidenceThreshold(80); - rule.setAiCustomPrompt("请识别敏感数据列"); - return rule; - } - - /** - * 创建测试用的数据库列 - * @param schemaName 数据库名 - * @param tableName 表名 - * @param columnName 列名 - * @param comment 列注释 - * @return 数据库列对象 - */ - private DBTableColumn createTestColumn(String schemaName, String tableName, String columnName, String comment) { - DBTableColumn column = new DBTableColumn(); - column.setSchemaName(schemaName); - column.setTableName(tableName); - column.setName(columnName); - column.setComment(comment); - column.setTypeName("VARCHAR"); - return column; - } -} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java deleted file mode 100644 index ad0398bcef..0000000000 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java +++ /dev/null @@ -1,218 +0,0 @@ -/* - * Copyright (c) 2025 OceanBase. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package com.oceanbase.odc.service.datasecurity.recognizer; - -import static org.mockito.Mockito.*; - -import java.util.Arrays; -import java.util.Collections; -import java.util.List; -import java.util.Map; -import java.util.Optional; - -import org.junit.Assert; -import org.junit.Before; -import org.junit.Test; -import org.mockito.Mock; -import org.mockito.MockitoAnnotations; - -import com.oceanbase.odc.service.common.util.SpringContextUtil; -import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; -import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; -import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; -import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; -import com.oceanbase.tools.dbbrowser.model.DBTableColumn; -import com.openai.models.chat.completions.ChatCompletion; - -/** - * AI列识别器单元测试类 - */ -public class AIColumnRecognizerTest { - - @Mock - private PromptTemplateLoader promptTemplateLoader; - - @Mock - private AIInferenceService aiInferenceService; - - @Mock - private SpringContextUtil springContextUtil; - - private SensitiveRule aiRule; - private AIColumnRecognizer aiColumnRecognizer; - - @Before - public void setUp() { - MockitoAnnotations.initMocks(this); - - // Mock SpringContextUtil - mockStatic(SpringContextUtil.class); - when(SpringContextUtil.getBean(PromptTemplateLoader.class)).thenReturn(promptTemplateLoader); - when(SpringContextUtil.getBean(AIInferenceService.class)).thenReturn(aiInferenceService); - - // 创建测试用的AI规则 - aiRule = createTestAIRule(1L); - aiColumnRecognizer = new AIColumnRecognizer(aiRule); - } - - @Test - public void recognize_singleColumn_returnEmpty() { - // 准备阶段：创建测试列 - DBTableColumn column = createTestColumn("test_db", "users", "email", "用户邮箱"); - - // 执行阶段：调用单列识别方法 - Optional result = aiColumnRecognizer.recognize(column); - - // 验证阶段：确认单列识别返回空结果 - Assert.assertFalse("单列识别应返回空结果", result.isPresent()); - } - - @Test - public void recognizeBatch_emptyList_returnEmptyMap() { - // 执行阶段：调用批量识别方法，传入空列表 - Map result = aiColumnRecognizer.recognizeBatch(Collections.emptyList()); - - // 验证阶段：确认返回空Map - Assert.assertNotNull("结果不应为null", result); - Assert.assertTrue("空列表应返回空Map", result.isEmpty()); - } - - @Test - public void recognizeBatch_nullList_returnEmptyMap() { - // 执行阶段：调用批量识别方法，传入null - Map result = aiColumnRecognizer.recognizeBatch(null); - - // 验证阶段：确认返回空Map - Assert.assertNotNull("结果不应为null", result); - Assert.assertTrue("null列表应返回空Map", result.isEmpty()); - } - - @Test - public void recognizeBatch_singleColumn_returnResult() throws Exception { - // 准备阶段：创建测试数据和模拟AI服务响应 - DBTableColumn column = createTestColumn("test_db", "users", "email", "用户邮箱"); - List columns = Arrays.asList(column); - - String prompt = "test prompt"; - when(promptTemplateLoader.buildPrompt(anyString(), anyList(), anyString())).thenReturn(prompt); - - String aiResponse = "[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 95, \"sensitiveType\": \"联系方式\"}]"; - ChatCompletion chatCompletion = mock(ChatCompletion.class); - when(chatCompletion.choices()).thenReturn(Arrays.asList(mock(com.openai.models.chat.completions.ChatCompletion.Choice.class))); - when(chatCompletion.choices().get(0).message()).thenReturn(mock(com.openai.models.chat.completions.ChatCompletionMessage.class)); - when(chatCompletion.choices().get(0).message().content()).thenReturn(Optional.of(aiResponse)); - when(aiInferenceService.chat(prompt)).thenReturn(chatCompletion); - - // 执行阶段：调用批量识别方法 - Map result = aiColumnRecognizer.recognizeBatch(columns); - - // 验证阶段：确认返回正确的识别结果 - Assert.assertNotNull("结果不应为null", result); - Assert.assertEquals("应返回一个识别结果", 1, result.size()); - Assert.assertTrue("应包含指定列的结果", result.containsKey("users.email")); - - RecognitionResult recognitionResult = result.get("users.email"); - Assert.assertTrue("应识别为敏感列", recognitionResult.isMatched()); - Assert.assertEquals("规则ID应匹配", aiRule.getId(), recognitionResult.getMatchedRuleId()); - Assert.assertEquals("规则类型应为AI", SensitiveRuleType.AI, recognitionResult.getSourceRuleType()); - Assert.assertEquals("风险等级应为HIGH", SensitiveLevel.HIGH, recognitionResult.getLevel()); - Assert.assertEquals("置信度应为95", Double.valueOf(95), recognitionResult.getConfidence()); - Assert.assertEquals("敏感类型应为联系方式", "联系方式", recognitionResult.getSensitiveType()); - } - - @Test - public void recognizeBatch_multipleColumns_returnResults() throws Exception { - // 准备阶段：创建多个测试列和模拟AI服务响应 - DBTableColumn column1 = createTestColumn("test_db", "users", "email", "用户邮箱"); - DBTableColumn column2 = createTestColumn("test_db", "employees", "salary", "员工薪资"); - DBTableColumn column3 = createTestColumn("test_db", "products", "name", "产品名称"); - List columns = Arrays.asList(column1, column2, column3); - - String prompt = "test prompt"; - when(promptTemplateLoader.buildPrompt(anyString(), anyList(), anyString())).thenReturn(prompt); - - // 模拟AI响应：前两列为敏感列，第三列（name）为非敏感列 - String aiResponse = "[" - + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 95, \"sensitiveType\": \"联系方式\"}," - + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"confidence\": 90, \"sensitiveType\": \"财务信息\"}," - + "{\"sensitive\": false, \"riskLevel\": \"LOW\", \"confidence\": 20, \"sensitiveType\": null}" - + "]"; - ChatCompletion chatCompletion = mock(ChatCompletion.class); - when(chatCompletion.choices()).thenReturn(Arrays.asList(mock(com.openai.models.chat.completions.ChatCompletion.Choice.class))); - when(chatCompletion.choices().get(0).message()).thenReturn(mock(com.openai.models.chat.completions.ChatCompletionMessage.class)); - when(chatCompletion.choices().get(0).message().content()).thenReturn(Optional.of(aiResponse)); - when(aiInferenceService.chat(prompt)).thenReturn(chatCompletion); - - // 执行阶段：调用批量识别方法 - Map result = aiColumnRecognizer.recognizeBatch(columns); - - // 验证阶段：确认返回正确的识别结果数量（只有敏感列会被返回） - Assert.assertNotNull("结果不应为null", result); - Assert.assertEquals("应返回两个识别结果（只有敏感列会被返回）", 2, result.size()); - Assert.assertTrue("应包含users.email的结果", result.containsKey("users.email")); - Assert.assertTrue("应包含employees.salary的结果", result.containsKey("employees.salary")); - Assert.assertFalse("不应包含products.name的结果（非敏感列）", result.containsKey("products.name")); - - // 验证email列的识别结果 - RecognitionResult emailResult = result.get("users.email"); - Assert.assertTrue("email应识别为敏感列", emailResult.isMatched()); - Assert.assertEquals("email敏感类型应为联系方式", "联系方式", emailResult.getSensitiveType()); - - // 验证salary列的识别结果 - RecognitionResult salaryResult = result.get("employees.salary"); - Assert.assertTrue("salary应识别为敏感列", salaryResult.isMatched()); - Assert.assertEquals("salary敏感类型应为财务信息", "财务信息", salaryResult.getSensitiveType()); - } - - // --- 辅助方法 --- - - /** - * 创建测试用的AI规则 - * @param id 规则ID - * @return AI规则对象 - */ - private SensitiveRule createTestAIRule(Long id) { - SensitiveRule rule = new SensitiveRule(); - rule.setId(id); - rule.setType(SensitiveRuleType.AI); - rule.setEnabled(true); - rule.setLevel(SensitiveLevel.HIGH); - rule.setAiSensitiveTypes(Arrays.asList("联系方式", "财务信息", "身份信息")); - rule.setAiConfidenceThreshold(80); - rule.setAiCustomPrompt("请识别敏感数据列"); - return rule; - } - - /** - * 创建测试用的数据库列 - * @param schemaName 数据库名 - * @param tableName 表名 - * @param columnName 列名 - * @param comment 列注释 - * @return 数据库列对象 - */ - private DBTableColumn createTestColumn(String schemaName, String tableName, String columnName, String comment) { - DBTableColumn column = new DBTableColumn(); - column.setSchemaName(schemaName); - column.setTableName(tableName); - column.setName(columnName); - column.setComment(comment); - column.setTypeName("VARCHAR"); - return column; - } -} \ No newline at end of file From cb1c373873eaf93eaa07dfc546ee8df58809c737 Mon Sep 17 00:00:00 2001 From: fenyf Date: Wed, 13 Aug 2025 20:38:20 +0800 Subject: [PATCH 03/10] feature(ai_recognition): Add single-table recognition functionality for passive scanning and refactor the code in the sole AI mode --- .../v2/SensitiveColumnController.java | 14 ++ .../datasecurity/SensitiveColumnService.java | 169 ++++++++++++++++++ .../SingleTableScanTaskManager.java | 116 ++++++++++++ .../factory/ScanningStrategyFactory.java | 2 + .../datasecurity/model/ScanResult.java | 19 +- .../datasecurity/model/ScanningModeType.java | 2 + .../model/SingleTableScanReq.java | 52 ++++++ .../recognizer/AIColumnRecognizer.java | 1 - .../datasecurity/strategy/AIOnlyStrategy.java | 59 ++++++ 9 files changed, 416 insertions(+), 18 deletions(-) create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java index 381760f830..d48cbc4654 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java @@ -35,12 +35,14 @@ import com.oceanbase.odc.service.common.response.Responses; import com.oceanbase.odc.service.common.response.SuccessResponse; import com.oceanbase.odc.service.datasecurity.SensitiveColumnService; +import com.oceanbase.odc.service.datasecurity.SingleTableScanTaskManager; import com.oceanbase.odc.service.datasecurity.model.DatabaseWithAllColumns; import com.oceanbase.odc.service.datasecurity.model.QuerySensitiveColumnParams; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumn; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningReq; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnStats; +import com.oceanbase.odc.service.datasecurity.model.SingleTableScanReq; import com.oceanbase.odc.service.datasecurity.model.UpdateSensitiveColumnsReq; import io.swagger.annotations.ApiOperation; @@ -150,4 +152,16 @@ public SuccessResponse stopScanning(@PathVariable Long projectId, return Responses.success(service.stopScanning(projectId, taskId)); } + @ApiOperation(value = "getSingleTableScanResult", notes = "Get single table scan result") + @RequestMapping(value = "/singleTableScan/{taskId}/result", method = RequestMethod.GET) + public SuccessResponse getSingleTableScanResult(@PathVariable Long projectId, + @PathVariable String taskId) { + return Responses.success(service.getSingleTableScanResult(projectId, taskId)); + } + @ApiOperation(value = "scanSingleTableAsync", notes = "Start an asynchronous single table scan") + @RequestMapping(value = "/scanSingleTableAsync", method = RequestMethod.POST) + public SuccessResponse scanSingleTableAsync(@PathVariable Long projectId, + @RequestBody SingleTableScanReq req) { + return Responses.success(service.scanSingleTableAsync(projectId, req)); + } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java index 19fa165685..8130d77ef4 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java @@ -23,7 +23,9 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; +import java.util.UUID; import java.util.stream.Collectors; import javax.validation.Valid; @@ -61,15 +63,22 @@ import com.oceanbase.odc.service.connection.database.model.Database; import com.oceanbase.odc.service.connection.model.ConnectionConfig; import com.oceanbase.odc.service.datasecurity.extractor.model.DBColumn; +import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; import com.oceanbase.odc.service.datasecurity.model.DatabaseWithAllColumns; import com.oceanbase.odc.service.datasecurity.model.MaskingAlgorithm; import com.oceanbase.odc.service.datasecurity.model.QuerySensitiveColumnParams; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumn; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnMeta; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningReq; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnStats; +import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnType; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.model.SingleTableScanReq; import com.oceanbase.odc.service.datasecurity.util.SensitiveColumnMapper; import com.oceanbase.odc.service.db.browser.DBSchemaAccessors; import com.oceanbase.odc.service.feature.VersionDiffConfigService; @@ -115,6 +124,8 @@ public class SensitiveColumnService { private HorizontalDataPermissionValidator permissionValidator; @Autowired private VersionDiffConfigService versionDiffConfigService; + @Autowired + private SingleTableScanTaskManager singleTableScanTaskManager; @Transactional(rollbackFor = Exception.class) @PreAuthenticate(hasAnyResourceRole = {"OWNER, DBA, SECURITY_ADMINISTRATOR"}, @@ -542,4 +553,162 @@ private Map> getFilteringExistColumns(Long databaseI return filtered; } + /** + * 扫描单个表的敏感列 + */ + @Transactional(rollbackFor = Exception.class) + @PreAuthenticate(hasAnyResourceRole = { "OWNER, DBA, SECURITY_ADMINISTRATOR" }, actions = { "OWNER", "DBA", + "SECURITY_ADMINISTRATOR" }, resourceType = "ODC_PROJECT", indexOfIdParam = 0) + public String scanSingleTableAsync(@NotNull Long projectId, @NotNull @Valid SingleTableScanReq req) { + // 获取当前用户信息，用于在异步任务中设置认证上下文 + final Long currentUserId = authenticationFacade.currentUserId(); + final Long currentOrganizationId = authenticationFacade.currentOrganizationId(); + final String currentUserAccountName = authenticationFacade.currentUserAccountName(); + + // 启动异步任务 + String taskId = UUID.randomUUID().toString(); + singleTableScanTaskManager.startTask(taskId, () -> { + try { + // 在异步任务中设置用户认证上下文 + com.oceanbase.odc.service.iam.util.SecurityContextUtils.setCurrentUser( + currentUserId, currentOrganizationId, currentUserAccountName); + + List result = performSingleTableScan(projectId, req); + singleTableScanTaskManager.setTaskResult(taskId, result); + } catch (Exception e) { + log.error("Single table scan failed for taskId: {}, projectId: {}, databaseId: {}, tableName: {}", + taskId, projectId, req.getDatabaseId(), req.getTableName(), e); + String errorMessage = e.getMessage() != null ? e.getMessage() + : "扫描过程中发生未知错误: " + e.getClass().getSimpleName(); + singleTableScanTaskManager.setTaskError(taskId, errorMessage); + } + }); + return taskId; + } + + /** + * 获取单表扫描结果 + */ + @PreAuthenticate(hasAnyResourceRole = { "OWNER, DBA, SECURITY_ADMINISTRATOR" }, actions = { "OWNER", "DBA", + "SECURITY_ADMINISTRATOR" }, resourceType = "ODC_PROJECT", indexOfIdParam = 0) + public SingleTableScanTaskManager.SingleTableScanTask getSingleTableScanResult(@NotNull Long projectId, + @NotBlank String taskId) { + return singleTableScanTaskManager.getTask(taskId); + } + + /** + * 执行单表扫描的具体逻辑 + */ + private List performSingleTableScan(@NotNull Long projectId, + @NotNull @Valid SingleTableScanReq req) { + // 1. 获取数据库信息 + Database database = databaseService.detail(req.getDatabaseId()); + PreConditions.notNull(database, "database"); + checkProjectDatabases(projectId, Collections.singletonList(req.getDatabaseId())); + + // 2. 获取连接配置 + ConnectionConfig connectionConfig = connectionService + .getForConnectionSkipPermissionCheck(database.getDataSource().getId()); + + // 3. 获取表列信息 + List tableColumns = getTableColumns(connectionConfig, database.getName(), req.getTableName()); + if (CollectionUtils.isEmpty(tableColumns)) { + return Collections.emptyList(); + } + + // 4. 获取扫描规则（使用预置的系统规则） + List rules = getScanningRules(projectId, null); + if (CollectionUtils.isEmpty(rules)) { + return Collections.emptyList(); + } + + // 5. 执行扫描 - 使用批量扫描以保持表级别上下文 + ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); + SensitiveColumnScanner scanner = new SensitiveColumnScanner(rules, strategyFactory); + + // 使用批量扫描而不是逐个扫描，这样AI可以看到整个表的所有列 + Map scanResults = scanner.scanBatch(tableColumns, req.getScanningMode()); + + List results = new ArrayList<>(); + for (DBTableColumn column : tableColumns) { + String columnKey = String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + + ScanResult scanResult = scanResults.get(columnKey); + if (scanResult != null) { + Optional finalResult = scanResult.getFinalResult(req.getScanningMode()); + + if (finalResult.isPresent()) { + RecognitionResult result = finalResult.get(); + SensitiveColumn sensitiveColumn = new SensitiveColumn(); + sensitiveColumn.setDatabase(database); + sensitiveColumn.setTableName(column.getTableName()); + sensitiveColumn.setColumnName(column.getName()); + sensitiveColumn.setType(SensitiveColumnType.TABLE_COLUMN); + sensitiveColumn.setEnabled(true); + sensitiveColumn.setSensitiveRuleId(result.getMatchedRuleId()); + sensitiveColumn.setLevel(result.getLevel()); + // 设置默认脱敏算法ID，可以从规则中获取 + SensitiveRule matchedRule = rules.stream() + .filter(r -> r.getId().equals(result.getMatchedRuleId())) + .findFirst() + .orElse(null); + if (matchedRule != null && matchedRule.getMaskingAlgorithmId() != null) { + sensitiveColumn.setMaskingAlgorithmId(matchedRule.getMaskingAlgorithmId()); + } else { + // 设置默认脱敏算法ID + sensitiveColumn.setMaskingAlgorithmId(1L); + } + results.add(sensitiveColumn); + } + } + } + + return results; + } + + /** + * 获取指定表的列信息 + */ + private List getTableColumns(ConnectionConfig connectionConfig, String databaseName, + String tableName) { + ConnectionSession session = new DefaultConnectSessionFactory(connectionConfig).generateSession(); + try { + DBSchemaAccessor accessor = DBSchemaAccessors.create(session); + return accessor.listTableColumns(databaseName, tableName); + } finally { + session.expire(); + } + } + + /** + * 获取单表扫描的预置规则 + */ + private List getScanningRules(Long projectId, List sensitiveRuleIds) { + // 单表扫描使用预置的系统规则，不依赖用户配置的规则 + SensitiveRule defaultRule = createDefaultScanningRule(); + return Collections.singletonList(defaultRule); + } + + /** + * 创建单表扫描的默认规则 + */ + private SensitiveRule createDefaultScanningRule() { + SensitiveRule rule = new SensitiveRule(); + rule.setId(-1L); // 使用负数ID表示系统预置规则 + rule.setName("Single Table Scan Default Rule"); + rule.setEnabled(true); + rule.setType(SensitiveRuleType.AI); // 使用AI类型 + rule.setAiSensitiveTypes(null); // 设置为null，对应提示词中的"No specified category" + rule.setAiCustomPrompt(null); // 设置为null，对应提示词中的"No supplementary rule" + Long organizationId = authenticationFacade.currentOrganizationId(); + Long defaultAlgorithmId = algorithmService.getDefaultAlgorithmIdByOrganizationId(organizationId); + rule.setMaskingAlgorithmId(defaultAlgorithmId); + rule.setLevel(SensitiveLevel.MEDIUM); // 设置默认敏感级别 + rule.setBuiltin(true); + return rule; + } } + diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java new file mode 100644 index 0000000000..e727a78050 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java @@ -0,0 +1,116 @@ +package com.oceanbase.odc.service.datasecurity; + +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CompletableFuture; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.beans.factory.annotation.Qualifier; +import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor; +import org.springframework.stereotype.Component; + +import com.oceanbase.odc.service.datasecurity.model.SensitiveColumn; + +import lombok.Data; +import lombok.extern.slf4j.Slf4j; + +/** + * 单表扫描任务管理器 + * 提供轻量级的异步任务管理功能 + */ +@Slf4j +@Component +public class SingleTableScanTaskManager { + + private final Map tasks = new ConcurrentHashMap<>(); + + @Autowired + @Qualifier("scanSensitiveColumnExecutor") + private ThreadPoolTaskExecutor executor; + + /** + * 启动单表扫描任务 + */ + public String startTask(String taskId, Runnable scanTask) { + SingleTableScanTask task = new SingleTableScanTask(taskId); + tasks.put(taskId, task); + + // 使用Spring的ThreadPoolTaskExecutor，它会自动传递Spring Security上下文 + executor.submit(() -> { + try { + task.setStatus(TaskStatus.RUNNING); + scanTask.run(); + task.setStatus(TaskStatus.COMPLETED); + } catch (Exception e) { + log.error("Single table scan task failed, taskId={}", taskId, e); + task.setStatus(TaskStatus.FAILED); + task.setErrorMessage(e.getMessage()); + } + }); + + return taskId; + } + + /** + * 启动单表扫描任务（自动生成taskId） + */ + public String startTask(Runnable scanTask) { + String taskId = UUID.randomUUID().toString(); + return startTask(taskId, scanTask); + } + + /** + * 获取任务状态 + */ + public SingleTableScanTask getTask(String taskId) { + return tasks.get(taskId); + } + + /** + * 设置任务结果 + */ + public void setTaskResult(String taskId, List result) { + SingleTableScanTask task = tasks.get(taskId); + if (task != null) { + task.setResult(result); + } + } + + /** + * 设置任务错误 + */ + public void setTaskError(String taskId, String errorMessage) { + SingleTableScanTask task = tasks.get(taskId); + if (task != null) { + task.setStatus(TaskStatus.FAILED); + task.setErrorMessage(errorMessage); + } + } + + /** + * 清理已完成的任务（可选的清理机制） + */ + public void cleanupTask(String taskId) { + tasks.remove(taskId); + } + + /** + * 任务状态枚举 + */ + public enum TaskStatus { + PENDING, RUNNING, COMPLETED, FAILED + } + + /** + * 单表扫描任务信息 + */ + @Data + public static class SingleTableScanTask { + private final String taskId; + private TaskStatus status = TaskStatus.PENDING; + private List result; + private String errorMessage; + } +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java index 4ea2e3c6a9..f2ab1db487 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java @@ -24,6 +24,7 @@ import com.oceanbase.odc.service.datasecurity.model.ScanResult; import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; +import com.oceanbase.odc.service.datasecurity.strategy.AIOnlyStrategy; import com.oceanbase.odc.service.datasecurity.strategy.JointRecognitionStrategy; import com.oceanbase.odc.service.datasecurity.strategy.RulesOnlyStrategy; import com.oceanbase.odc.service.datasecurity.strategy.ScanningStrategy; @@ -45,6 +46,7 @@ public ScanningStrategyFactory() { // 预创建所有策略实例 strategies.put(ScanningModeType.RULES_ONLY, new RulesOnlyStrategy()); strategies.put(ScanningModeType.JOINT_RECOGNITION, new JointRecognitionStrategy()); + strategies.put(ScanningModeType.AI_ONLY, new AIOnlyStrategy()); } /** diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java index f0cc9c0263..f81ec26dda 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java @@ -40,6 +40,8 @@ public Optional getFinalResult(ScanningModeType scanningMode) switch (scanningMode) { case RULES_ONLY: return basicRuleResult; + case AI_ONLY: + return aiRuleResult; case JOINT_RECOGNITION: // 对于联合识别，Scanner已经做过决策，直接返回存在的那个结果 return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; @@ -47,21 +49,4 @@ public Optional getFinalResult(ScanningModeType scanningMode) return Optional.empty(); } } - - /** - * 判断是否有任何识别结果 - */ - public boolean hasAnyResult() { - return basicRuleResult.isPresent() || aiRuleResult.isPresent(); - } - - /** - * 获取所有可用的结果（用于差异化展示场景） - */ - public List getAllResults() { - List results = new ArrayList<>(); - basicRuleResult.ifPresent(results::add); - aiRuleResult.ifPresent(results::add); - return results; - } } \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java index 33a2740272..5355078bf1 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanningModeType.java @@ -24,4 +24,6 @@ public enum ScanningModeType { RULES_ONLY, JOINT_RECOGNITION, + + AI_ONLY; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java new file mode 100644 index 0000000000..f5a93092b7 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java @@ -0,0 +1,52 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.model; + +import java.util.List; + +import javax.validation.constraints.NotBlank; +import javax.validation.constraints.NotNull; + +import lombok.Data; + +/** + * 单表敏感列扫描请求 + * + * @author Assistant + * @date 2025/1/27 + */ +@Data +public class SingleTableScanReq { + + /** + * 数据库ID + */ + @NotNull + private Long databaseId; + + /** + * 表名 + */ + @NotBlank + private String tableName; + + /** + * 扫描模式，默认为AI识别 + */ + @NotNull + private ScanningModeType scanningMode = ScanningModeType.JOINT_RECOGNITION; + +} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index 08e8b2cf63..e8d7011778 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -46,7 +46,6 @@ public class AIColumnRecognizer implements ColumnRecognizer { private final SensitiveRule aiRule; // 直接保存整个规则对象 - // @Value() private static final int BATCH_SIZE = AIParam.DEFAULT_BATCH_SIZE_IN_TABLE; // 单表内列数超过此值时进行分批处理 private static final ObjectMapper objectMapper = new ObjectMapper(); // 用于解析JSON private static final Pattern JSON_PATTERN = Pattern diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java new file mode 100644 index 0000000000..c69207be8a --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +/** + * 仅AI扫描策略实现 + * 只使用AI识别器进行识别，忽略基础规则 + * + * @author Assistant + * @date 2025/1/27 + */ +public class AIOnlyStrategy extends AbstractScanningStrategy { + + @Override + public ScanResult scan(DBTableColumn column, List basicRecognizers, + List aiRecognizers) { + Optional aiResult = findFirstMatch(aiRecognizers, column); + return new ScanResult(Optional.empty(), aiResult); + } + + @Override + public Map scanBatch(List columns, List basicRecognizers, + List aiRecognizers) { + Map> aiResults = findAllFirstMatches(aiRecognizers, columns); + Map results = new HashMap<>(); + + for (DBTableColumn column : columns) { + String columnKey = getColumnKey(column); + Optional aiResult = aiResults.getOrDefault(columnKey, Optional.empty()); + results.put(columnKey, new ScanResult(Optional.empty(), aiResult)); + } + + return results; + } +} \ No newline at end of file From 3cc6a8881fc9b3800d2fb16037126529eb3ce2d4 Mon Sep 17 00:00:00 2001 From: fenyf Date: Tue, 2 Sep 2025 11:04:07 +0800 Subject: [PATCH 04/10] feature(ai_recognition): Add the function of querying AI status --- .../V_4_3_4_21__add_ai_system_config.sql | 16 +++++ .../web/controller/v2/AIController.java | 63 +++++++++++++++++++ .../odc/config/CommonSecurityProperties.java | 3 +- .../SingleTableScanTaskManager.java | 1 - .../datasecurity/ai/AIInferenceService.java | 42 ++++++++++--- .../datasecurity/ai/AIStatusResponse.java | 50 +++++++++++++++ 6 files changed, 164 insertions(+), 11 deletions(-) create mode 100644 server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_21__add_ai_system_config.sql create mode 100644 server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java create mode 100644 server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java diff --git a/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_21__add_ai_system_config.sql b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_21__add_ai_system_config.sql new file mode 100644 index 0000000000..9219e0d17d --- /dev/null +++ b/server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_21__add_ai_system_config.sql @@ -0,0 +1,16 @@ + +INSERT INTO config_system_configuration(`key`, `value`, `description`) +VALUES('odc.ai.enabled', 'false', 'Whether AI feature is enabled, disabled by default') +ON DUPLICATE KEY UPDATE `id`=`id`; + +INSERT INTO config_system_configuration(`key`, `value`, `description`) +VALUES('odc.ai.api-key', '', 'AI API key, required when AI feature is enabled') +ON DUPLICATE KEY UPDATE `id`=`id`; + +INSERT INTO config_system_configuration(`key`, `value`, `description`) +VALUES('odc.ai.base-url', 'https://api.openai.com', 'AI API base URL, defaults to OpenAI official API endpoint') +ON DUPLICATE KEY UPDATE `id`=`id`; + +INSERT INTO config_system_configuration(`key`, `value`, `description`) +VALUES('odc.ai.model', 'gpt-3.5-turbo', 'AI model to use, defaults to gpt-3.5-turbo') +ON DUPLICATE KEY UPDATE `id`=`id`; \ No newline at end of file diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java new file mode 100644 index 0000000000..ce4960d2d6 --- /dev/null +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java @@ -0,0 +1,63 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.server.web.controller.v2; + +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; + +import com.oceanbase.odc.core.authority.util.SkipAuthorize; +import com.oceanbase.odc.service.common.response.Responses; +import com.oceanbase.odc.service.common.response.SuccessResponse; +import com.oceanbase.odc.service.datasecurity.ai.AIConfig; +import com.oceanbase.odc.service.datasecurity.ai.AIStatusResponse; + +import io.swagger.annotations.Api; +import io.swagger.annotations.ApiOperation; + +/** + * AI功能控制器 + * 提供AI功能状态查询接口 + */ +@Api(tags = "AI功能") +@RestController +@RequestMapping("/api/v2/ai") +public class AIController { + + @Autowired + private AIConfig aiConfig; + + /** + * 查询AI功能状态 + * @return AI功能状态信息 + */ + @ApiOperation(value = "查询AI功能状态", notes = "返回AI功能是否启用以及配置状态") + @SkipAuthorize("AI status is safe to query for authenticated users") + @GetMapping("/status") + public SuccessResponse getAIStatus() { + AIStatusResponse response = new AIStatusResponse(); + response.setEnabled(aiConfig.isEnabled()); + response.setAvailable(aiConfig.isAIAvailable()); + response.setModel(aiConfig.getModel()); + response.setBaseUrl(aiConfig.getBaseUrl()); + // 不返回敏感信息如API密钥 + response.setApiKeyConfigured(aiConfig.getApiKey() != null && !aiConfig.getApiKey().trim().isEmpty()); + + return Responses.success(response); + } +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/config/CommonSecurityProperties.java b/server/odc-service/src/main/java/com/oceanbase/odc/config/CommonSecurityProperties.java index b2e8bbfb4f..86a78a66db 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/config/CommonSecurityProperties.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/config/CommonSecurityProperties.java @@ -56,7 +56,8 @@ public class CommonSecurityProperties { "/api/v2/internal/file/downloadImportFile", "/api/v2/info", "/api/v2/sso/state", - "/api/v2/encryption/publicKey"}; + "/api/v2/encryption/publicKey", + "/api/v2/ai/status"}; private static final String[] STATIC_RESOURCES = new String[] { "/", diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java index e727a78050..4c1dc60f67 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java @@ -4,7 +4,6 @@ import java.util.Map; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; -import java.util.concurrent.CompletableFuture; import org.springframework.beans.factory.annotation.Autowired; import org.springframework.beans.factory.annotation.Qualifier; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java index 5265a12b20..926127308e 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -15,37 +15,53 @@ */ package com.oceanbase.odc.service.datasecurity.ai; -import java.util.HashMap; -import java.util.Map; +import java.util.Optional; import org.springframework.stereotype.Service; import com.openai.client.OpenAIClient; -import com.openai.core.JsonBoolean; -import com.openai.core.JsonValue; import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionCreateParams; @Service public class AIInferenceService { - private final AIConfig aiConfig; - private final OpenAIClient openAIClient; + private final AIConfig aiConfig; + private final Optional openAIClient; - // 直接注入 AIConfig 和 OpenAIClient 两个Bean - public AIInferenceService(AIConfig aiConfig, OpenAIClient openAIClient) { + // 注入 AIConfig 和可选的 OpenAIClient Bean + public AIInferenceService(AIConfig aiConfig, Optional openAIClient) { this.aiConfig = aiConfig; this.openAIClient = openAIClient; } + /** + * 检查AI功能是否可用 + * @throws IllegalStateException 如果AI功能不可用 + */ + private void checkAIAvailability() { + if (!aiConfig.isEnabled()) { + throw new IllegalStateException("AI功能未启用。请联系管理员启用AI功能。"); + } + if (!aiConfig.isAIAvailable()) { + throw new IllegalStateException("AI功能配置不完整。请联系管理员配置AI相关参数。"); + } + if (!openAIClient.isPresent()) { + throw new IllegalStateException("AI客户端未初始化。请检查AI配置并重启服务。"); + } + } + /** * 使用系统提示词和用户提示词分别调用AI服务 * * @param systemPrompt 系统提示词 * @param userPrompt 用户提示词 * @return AI响应 + * @throws IllegalStateException 如果AI功能不可用 */ public ChatCompletion chat(String systemPrompt, String userPrompt) { + // 检查AI功能可用性 + checkAIAvailability(); try { ChatCompletionCreateParams params = ChatCompletionCreateParams.builder() @@ -56,9 +72,17 @@ public ChatCompletion chat(String systemPrompt, String userPrompt) { .topP(aiConfig.getTopP()) .additionalBodyProperties(aiConfig.loadAdditionalParams()) .build(); - return openAIClient.chat().completions().create(params); + return openAIClient.get().chat().completions().create(params); } catch (Exception e) { throw new RuntimeException("调用AI服务失败: " + e.getMessage(), e); } } + + /** + * 检查AI功能是否可用（不抛出异常） + * @return true if AI功能可用 + */ + public boolean isAIAvailable() { + return aiConfig.isEnabled() && aiConfig.isAIAvailable() && openAIClient.isPresent(); + } } \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java new file mode 100644 index 0000000000..1f59358382 --- /dev/null +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java @@ -0,0 +1,50 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.oceanbase.odc.service.datasecurity.ai; + +import lombok.Data; + +/** + * AI功能状态响应 + */ + @Data + public class AIStatusResponse { + /** + * AI功能是否启用 + */ + private boolean enabled; + + /** + * AI功能是否可用（启用且配置完整） + */ + private boolean available; + + /** + * 使用的AI模型 + */ + private String model; + + /** + * API基础URL + */ + private String baseUrl; + + /** + * API密钥是否已配置 + */ + private boolean apiKeyConfigured; + } \ No newline at end of file From 5785dda8a2a0cb6f22aeb89e8844ed4c5c49f8c2 Mon Sep 17 00:00:00 2001 From: fenyf Date: Tue, 2 Sep 2025 11:19:18 +0800 Subject: [PATCH 05/10] feature(ai_recognition): Define exception codes and standardize logs --- .../odc/core/shared/constant/ErrorCodes.java | 10 +++++ .../SensitiveColumnScanningTask.java | 10 +++-- .../odc/service/datasecurity/ai/AIConfig.java | 38 ++++++++++++++----- .../datasecurity/ai/AIInferenceService.java | 10 +++-- .../datasecurity/ai/PromptTemplateLoader.java | 3 ++ .../recognizer/AIColumnRecognizer.java | 33 ++++++++++++---- 6 files changed, 79 insertions(+), 25 deletions(-) diff --git a/server/odc-core/src/main/java/com/oceanbase/odc/core/shared/constant/ErrorCodes.java b/server/odc-core/src/main/java/com/oceanbase/odc/core/shared/constant/ErrorCodes.java index d7954a985e..64e73603c4 100644 --- a/server/odc-core/src/main/java/com/oceanbase/odc/core/shared/constant/ErrorCodes.java +++ b/server/odc-core/src/main/java/com/oceanbase/odc/core/shared/constant/ErrorCodes.java @@ -334,6 +334,16 @@ public enum ErrorCodes implements ErrorCode { ExtractFileFailed, InvalidSignature, + /** + * AI Service + */ + AIServiceNotAvailable, + AIConfigurationIncomplete, + AIClientNotInitialized, + AIInferenceServiceError, + AIResponseFormatError, + AIResponseCountMismatch, + ; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java index 362be5f09a..2665991fec 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java @@ -44,10 +44,13 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import com.oceanbase.odc.service.common.util.SpringContextUtil; +import lombok.extern.slf4j.Slf4j; + /** * @author gaoda.xy * @date 2023/5/25 14:43 */ +@Slf4j public class SensitiveColumnScanningTask implements Callable { private final Database database; @@ -165,8 +168,7 @@ private void scanColumns(Map> object2Columns, Sensit } taskInfo.addFinishedTableCount(); } catch (Exception e) { - System.err.println("表 " + objectName + " 扫描失败: " + e.toString()); - e.printStackTrace(); + log.error("Failed to scan table {}: {}", objectName, e.getMessage(), e); // 即使失败也要增加完成计数，避免任务卡住 taskInfo.addFinishedTableCount(); } @@ -235,7 +237,7 @@ private Long handleAiRecognitionResult(String sensitiveType) { return algorithmIdOpt.get(); } } catch (Exception e) { - System.err.println("Failed to get algorithm ID by name: " + e.getMessage()); + log.error("Failed to get algorithm ID by name: {}", e.getMessage(), e); } } } @@ -253,7 +255,7 @@ private Long getSystemDefaultAlgorithmId() { return algorithmService.getDefaultAlgorithmIdByOrganizationId(database.getOrganizationId()); } catch (Exception e) { // 记录错误日志，但不抛出异常，避免影响整个扫描流程 - System.err.println("Failed to get default masking algorithm ID: " + e.getMessage()); + log.error("Failed to get default masking algorithm ID: {}", e.getMessage(), e); return null; } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java index e21a6afd17..b61ad5e640 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java @@ -18,9 +18,13 @@ import java.util.HashMap; import java.util.Map; +import org.springframework.beans.factory.annotation.Value; +import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty; import org.springframework.context.annotation.Bean; import org.springframework.stereotype.Component; +import com.oceanbase.odc.core.shared.constant.ErrorCodes; +import com.oceanbase.odc.core.shared.exception.BadRequestException; import com.openai.client.OpenAIClient; import com.openai.client.okhttp.OpenAIOkHttpClient; import com.openai.core.JsonBoolean; @@ -30,22 +34,26 @@ import lombok.Data; /** - * 后期需要弄成动态配置 + * AI功能配置类，支持从系统配置中动态读取配置 */ @Data @Component -// @ConfigurationProperties(prefix = "datasecurity.ai") public class AIConfig { - private boolean enabled = true; + @Value("${odc.ai.enabled:false}") + private boolean enabled; - //private String apiKey = "sk-c6bbbbde1b7e420b897d0662301c6d7c"; - private String apiKey = "token-abc123"; + @Value("${odc.ai.api-key:}") + private String apiKey; - //private String baseUrl = "https://dashscope.aliyuncs.com/compatible-mode/v1"; - private String baseUrl = "http://172.25.17.78:8000/v1"; + @Value("${odc.ai.base-url:https://api.openai.com}") + private String baseUrl; - //private String model = "qwen3-8b"; - private String model = "nlora"; + @Value("${odc.ai.model:gpt-3.5-turbo}") + private String model; + + // 硬编码超时和重试配置 + private static final int TIMEOUT_SECONDS = 30; + private static final int MAX_RETRIES = 3; private Boolean enableThinking = AIParam.DEFAULT_ENABLE_THINKING; @@ -66,10 +74,22 @@ public Map loadAdditionalParams() { } @Bean + @ConditionalOnProperty(name = "odc.ai.enabled", havingValue = "true") public OpenAIClient openAIClient() { + if (apiKey == null || apiKey.trim().isEmpty()) { + throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, new Object[]{"API key is not configured"}, "AI service is enabled but API key is not configured. Please set odc.ai.api-key configuration."); + } return OpenAIOkHttpClient.builder() .apiKey(this.apiKey) .baseUrl(this.baseUrl) .build(); } + + /** + * 检查AI功能是否可用 + * @return true if AI功能启用且配置完整 + */ + public boolean isAIAvailable() { + return enabled && apiKey != null && !apiKey.trim().isEmpty(); + } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java index 926127308e..02a1e0d6ed 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -19,6 +19,8 @@ import org.springframework.stereotype.Service; +import com.oceanbase.odc.core.shared.constant.ErrorCodes; +import com.oceanbase.odc.core.shared.exception.BadRequestException; import com.openai.client.OpenAIClient; import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionCreateParams; @@ -41,13 +43,13 @@ public AIInferenceService(AIConfig aiConfig, Optional openAIClient */ private void checkAIAvailability() { if (!aiConfig.isEnabled()) { - throw new IllegalStateException("AI功能未启用。请联系管理员启用AI功能。"); + throw new BadRequestException(ErrorCodes.AIServiceNotAvailable, new Object[]{"AI service is not enabled"}, "AI service is not enabled. Please contact administrator to enable AI service."); } if (!aiConfig.isAIAvailable()) { - throw new IllegalStateException("AI功能配置不完整。请联系管理员配置AI相关参数。"); + throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, new Object[]{"AI configuration is incomplete"}, "AI configuration is incomplete. Please contact administrator to configure AI parameters."); } if (!openAIClient.isPresent()) { - throw new IllegalStateException("AI客户端未初始化。请检查AI配置并重启服务。"); + throw new BadRequestException(ErrorCodes.AIClientNotInitialized, new Object[]{"AI client is not initialized"}, "AI client is not initialized. Please check AI configuration and restart service."); } } @@ -74,7 +76,7 @@ public ChatCompletion chat(String systemPrompt, String userPrompt) { .build(); return openAIClient.get().chat().completions().create(params); } catch (Exception e) { - throw new RuntimeException("调用AI服务失败: " + e.getMessage(), e); + throw new BadRequestException(ErrorCodes.AIInferenceServiceError, new Object[]{e.getMessage()}, "Failed to call AI inference service: " + e.getMessage(), e); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java index 92f84a33a5..c0fae884da 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java @@ -25,6 +25,7 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; +import lombok.extern.slf4j.Slf4j; import lombok.var; /** @@ -33,6 +34,7 @@ * 该类负责加载结构化的 AI 提示词模板，并根据列元数据、指定的敏感类型和用户自定义提示来构建最终的提示词。 *

*/ +@Slf4j @Component public class PromptTemplateLoader { @@ -56,6 +58,7 @@ public void init() { } catch (Exception e) { // 在实际项目中，这里应该使用日志系统 e.printStackTrace(); + log.error("Failed to load AI system prompt template: {}", e.getMessage(), e); throw new IllegalStateException("Failed to load AI system prompt template", e); } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index e8d7011778..d8d74db1ed 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -28,6 +28,8 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.Lists; +import com.oceanbase.odc.core.shared.constant.ErrorCodes; +import com.oceanbase.odc.core.shared.exception.BadRequestException; import com.oceanbase.odc.service.common.util.SpringContextUtil; import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; import com.oceanbase.odc.service.datasecurity.ai.AIParam; @@ -39,10 +41,12 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import com.openai.models.chat.completions.ChatCompletion; import lombok.Data; +import lombok.extern.slf4j.Slf4j; /** * AI 列识别器（最终版） */ +@Slf4j public class AIColumnRecognizer implements ColumnRecognizer { private final SensitiveRule aiRule; // 直接保存整个规则对象 @@ -83,7 +87,7 @@ public Map> recognizeBatch(List> recognizeBatch(List batch, PromptTemplateLoader prompt // c. 使用正则表达式从AI的返回结果中安全地提取JSON数组字符串 Matcher matcher = JSON_PATTERN.matcher(rawContent); - String jsonArrayResponse = "[]"; // 提供一个安全的默认值，以防匹配失败 + String jsonArrayResponse = null; if (matcher.find()) { // group(1) 对应被 ```json [...] ``` 包裹的内容, group(2) 对应裸露的 [...] // 使用 Optional 来优雅地处理可能为null的捕获组 jsonArrayResponse = Optional.ofNullable(matcher.group(1)).orElse(matcher.group(2)); } + if (jsonArrayResponse == null) { + throw new BadRequestException(ErrorCodes.AIResponseFormatError, + new Object[]{"No valid JSON array found in AI response"}, + "AI response does not contain valid JSON format: " + rawContent); + } + // d. 解析提取出的、更纯净的 JSON 数组 - List batchResults = objectMapper.readValue(jsonArrayResponse, - new TypeReference>() { - }); + List batchResults; + try { + batchResults = objectMapper.readValue(jsonArrayResponse, + new TypeReference>() { + }); + } catch (Exception e) { + throw new BadRequestException(ErrorCodes.AIResponseFormatError, + new Object[]{"Failed to parse JSON: " + e.getMessage()}, + "Failed to parse AI response JSON: " + jsonArrayResponse, e); + } // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 @@ -152,8 +169,8 @@ private void processBatch(List batch, PromptTemplateLoader prompt // 如果AI返回的结果数量与输入不匹配，记录警告信息 if (batchResults.size() != batch.size()) { - System.err.println("警告: AI返回结果数量(" + batchResults.size() + - ")与输入列数量(" + batch.size() + ")不匹配"); + log.warn("AI response count ({}) does not match input column count ({})", + batchResults.size(), batch.size()); } } From 0431ac2defdfbccb95be3e6678fbb7d5707d1cb9 Mon Sep 17 00:00:00 2001 From: fenyf Date: Tue, 2 Sep 2025 14:06:30 +0800 Subject: [PATCH 06/10] feature(ai_recognition): Test modification MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit feature(ai_recognition): 集成测试 --- ...ensitiveColumnScanningTaskManagerTest.java | 13 +- .../recognizer/AIColumnRecognizer.java | 30 +- .../service/datasecurity/ai/AIConfigTest.java | 228 ++++++++++++ .../ai/AIInferenceServiceTest.java | 259 ++++++++++++++ .../ai/PromptTemplateLoaderTest.java | 221 ++++++++++++ .../factory/ScanningStrategyFactoryTest.java | 203 +++++++++++ .../recognizer/AIColumnRecognizerTest.java | 324 ++++++++++++++++++ .../GroovyColumnRecognizerTest.java | 25 +- .../recognizer/PathColumnRecognizerTest.java | 14 - .../recognizer/RegexColumnRecognizerTest.java | 15 +- .../strategy/AIOnlyStrategyTest.java | 229 +++++++++++++ .../AbstractScanningStrategyTest.java | 249 ++++++++++++++ .../JointRecognitionStrategyTest.java | 295 ++++++++++++++++ .../strategy/RulesOnlyStrategyTest.java | 206 +++++++++++ 14 files changed, 2244 insertions(+), 67 deletions(-) create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java create mode 100644 server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java diff --git a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java index 2d8b7e1980..a75b8325e6 100644 --- a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java +++ b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java @@ -47,6 +47,7 @@ import com.oceanbase.odc.service.collaboration.project.model.Project; import com.oceanbase.odc.service.connection.database.model.Database; import com.oceanbase.odc.service.connection.model.ConnectionConfig; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo; import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnScanningTaskInfo.ScanningTaskStatus; import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; @@ -107,7 +108,7 @@ public static void tearDown() { public void test_start_groovyRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createGroovySensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -117,7 +118,7 @@ public void test_start_groovyRule_OBMySQL() { public void test_start_groovyRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createGroovySensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -127,7 +128,7 @@ public void test_start_groovyRule_OBOracle() { public void test_start_pathRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createPathSensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -137,7 +138,7 @@ public void test_start_pathRule_OBMySQL() { public void test_start_pathRule_OBMOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createPathSensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -147,7 +148,7 @@ public void test_start_pathRule_OBMOracle() { public void test_start_RegexRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_MYSQL)); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -157,7 +158,7 @@ public void test_start_RegexRule_OBMySQL() { public void test_start_RegexRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_ORACLE)); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index d8d74db1ed..be1cf87864 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -28,9 +28,11 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.Lists; +import com.oceanbase.odc.service.common.util.SpringContextUtil; + +import lombok.extern.slf4j.Slf4j; import com.oceanbase.odc.core.shared.constant.ErrorCodes; import com.oceanbase.odc.core.shared.exception.BadRequestException; -import com.oceanbase.odc.service.common.util.SpringContextUtil; import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; import com.oceanbase.odc.service.datasecurity.ai.AIParam; import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; @@ -41,10 +43,9 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import com.openai.models.chat.completions.ChatCompletion; import lombok.Data; -import lombok.extern.slf4j.Slf4j; /** - * AI 列识别器（最终版） + * AI Column Recognizer */ @Slf4j public class AIColumnRecognizer implements ColumnRecognizer { @@ -86,6 +87,9 @@ public Map> recognizeBatch(List batch : batches) { processBatch(batch, promptTemplateLoader, aiService, finalAiResults); } + } catch (BadRequestException e) { + // 重新抛出BadRequestException，让调用方处理 + throw e; } catch (Exception e) { log.error("Failed to process AI column recognition batch", e); return finalAiResults; @@ -94,6 +98,9 @@ public Map> recognizeBatch(List> recognizeBatch(List batch, PromptTemplateLoader promptTemplateLoader, AIInferenceService aiService, Map> finalAiResults) throws IOException { // a. 构建系统提示词 - String systemPrompt = promptTemplateLoader.buildSystemPrompt(aiRule.getAiSensitiveTypes(), aiRule.getAiCustomPrompt()); + String systemPrompt = promptTemplateLoader.buildSystemPrompt(aiRule.getAiSensitiveTypes(), + aiRule.getAiCustomPrompt()); // b. 构建用户提示词（列数据的JSON数组） String userPrompt = buildUserPrompt(batch); - // c. 调用 AI ChatCompletion completion = aiService.chat(systemPrompt, userPrompt); String rawContent = completion.choices().get(0).message().content().orElse("[]"); @@ -129,7 +136,7 @@ private void processBatch(List batch, PromptTemplateLoader prompt if (jsonArrayResponse == null) { throw new BadRequestException(ErrorCodes.AIResponseFormatError, - new Object[]{"No valid JSON array found in AI response"}, + new Object[] { "No valid JSON array found in AI response" }, "AI response does not contain valid JSON format: " + rawContent); } @@ -141,11 +148,10 @@ private void processBatch(List batch, PromptTemplateLoader prompt }); } catch (Exception e) { throw new BadRequestException(ErrorCodes.AIResponseFormatError, - new Object[]{"Failed to parse JSON: " + e.getMessage()}, + new Object[] { "Failed to parse JSON: " + e.getMessage() }, "Failed to parse AI response JSON: " + jsonArrayResponse, e); } - // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 int maxIndex = Math.min(batch.size(), batchResults.size()); for (int i = 0; i < maxIndex; i++) { @@ -167,7 +173,7 @@ private void processBatch(List batch, PromptTemplateLoader prompt } } - // 如果AI返回的结果数量与输入不匹配，记录警告信息 + // Log warning if AI response count doesn't match input count if (batchResults.size() != batch.size()) { log.warn("AI response count ({}) does not match input column count ({})", batchResults.size(), batch.size()); @@ -175,7 +181,7 @@ private void processBatch(List batch, PromptTemplateLoader prompt } /** - * 构建用户提示词（列数据的JSON数组） + * Build user prompt (JSON array of column data) */ private String buildUserPrompt(List batch) throws IOException { // 将一批列的元数据转换为 JSON 数组字符串 @@ -191,7 +197,7 @@ private String buildUserPrompt(List batch) throws IOException { return objectMapper.writeValueAsString(columnMetadataList); } - // 用于承载 AI 返回的 JSON 数据的内部类 + // Inner class for holding AI response JSON data @Data private static class AiResponseDto { private boolean sensitive; diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java new file mode 100644 index 0000000000..257f88ebb4 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java @@ -0,0 +1,228 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import java.util.Map; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.mockito.junit.MockitoJUnitRunner; +import org.springframework.test.util.ReflectionTestUtils; + +import com.oceanbase.odc.core.shared.exception.BadRequestException; +import com.openai.client.OpenAIClient; +import com.openai.core.JsonValue; + +@RunWith(MockitoJUnitRunner.class) +public class AIConfigTest { + + @Rule + public ExpectedException thrown = ExpectedException.none(); + + private AIConfig aiConfig; + + @Before + public void setUp() { + aiConfig = new AIConfig(); + } + + @Test + public void test_isAIAvailable_enabledAndApiKeySet_returnsTrue() { + // Given + ReflectionTestUtils.setField(aiConfig, "enabled", true); + ReflectionTestUtils.setField(aiConfig, "apiKey", "test-api-key"); + + // When + boolean result = aiConfig.isAIAvailable(); + + // Then + Assert.assertTrue(result); + } + + @Test + public void test_isAIAvailable_disabledWithApiKey_returnsFalse() { + // Given + ReflectionTestUtils.setField(aiConfig, "enabled", false); + ReflectionTestUtils.setField(aiConfig, "apiKey", "test-api-key"); + + // When + boolean result = aiConfig.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_enabledWithoutApiKey_returnsFalse() { + // Given + ReflectionTestUtils.setField(aiConfig, "enabled", true); + ReflectionTestUtils.setField(aiConfig, "apiKey", ""); + + // When + boolean result = aiConfig.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_enabledWithNullApiKey_returnsFalse() { + // Given + ReflectionTestUtils.setField(aiConfig, "enabled", true); + ReflectionTestUtils.setField(aiConfig, "apiKey", null); + + // When + boolean result = aiConfig.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_enabledWithWhitespaceApiKey_returnsFalse() { + // Given + ReflectionTestUtils.setField(aiConfig, "enabled", true); + ReflectionTestUtils.setField(aiConfig, "apiKey", " "); + + // When + boolean result = aiConfig.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_loadAdditionalParams_defaultValues_returnsCorrectMap() { + // Given - using default values + + // When + Map params = aiConfig.loadAdditionalParams(); + + // Then + Assert.assertNotNull(params); + Assert.assertEquals(3, params.size()); + Assert.assertTrue(params.containsKey("enable_thinking")); + Assert.assertTrue(params.containsKey("top_k")); + Assert.assertTrue(params.containsKey("min_p")); + } + + @Test + public void test_loadAdditionalParams_customValues_returnsCorrectMap() { + // Given + ReflectionTestUtils.setField(aiConfig, "enableThinking", false); + ReflectionTestUtils.setField(aiConfig, "topK", 50); + ReflectionTestUtils.setField(aiConfig, "minP", 10); + + // When + Map params = aiConfig.loadAdditionalParams(); + + // Then + Assert.assertNotNull(params); + Assert.assertEquals(3, params.size()); + Assert.assertTrue(params.containsKey("enable_thinking")); + Assert.assertTrue(params.containsKey("top_k")); + Assert.assertTrue(params.containsKey("min_p")); + } + + @Test + public void test_openAIClient_validApiKey_returnsClient() { + // Given + ReflectionTestUtils.setField(aiConfig, "apiKey", "test-api-key"); + ReflectionTestUtils.setField(aiConfig, "baseUrl", "https://api.openai.com"); + + // When + OpenAIClient client = aiConfig.openAIClient(); + + // Then + Assert.assertNotNull(client); + } + + @Test + public void test_openAIClient_nullApiKey_throwsException() { + // Given + ReflectionTestUtils.setField(aiConfig, "apiKey", null); + thrown.expect(BadRequestException.class); + thrown.expectMessage("API key is not configured"); + + // When + aiConfig.openAIClient(); + } + + @Test + public void test_openAIClient_emptyApiKey_throwsException() { + // Given + ReflectionTestUtils.setField(aiConfig, "apiKey", ""); + thrown.expect(BadRequestException.class); + thrown.expectMessage("API key is not configured"); + + // When + aiConfig.openAIClient(); + } + + @Test + public void test_openAIClient_whitespaceApiKey_throwsException() { + // Given + ReflectionTestUtils.setField(aiConfig, "apiKey", " "); + thrown.expect(BadRequestException.class); + thrown.expectMessage("API key is not configured"); + + // When + aiConfig.openAIClient(); + } + + @Test + public void test_gettersAndSetters_workCorrectly() { + // Test enabled + aiConfig.setEnabled(true); + Assert.assertTrue(aiConfig.isEnabled()); + + // Test apiKey + aiConfig.setApiKey("test-key"); + Assert.assertEquals("test-key", aiConfig.getApiKey()); + + // Test baseUrl + aiConfig.setBaseUrl("https://test.com"); + Assert.assertEquals("https://test.com", aiConfig.getBaseUrl()); + + // Test model + aiConfig.setModel("gpt-4"); + Assert.assertEquals("gpt-4", aiConfig.getModel()); + + // Test temperature + aiConfig.setTemperature(0.8); + Assert.assertEquals(Double.valueOf(0.8), aiConfig.getTemperature()); + + // Test topP + aiConfig.setTopP(0.9); + Assert.assertEquals(Double.valueOf(0.9), aiConfig.getTopP()); + + // Test enableThinking + aiConfig.setEnableThinking(false); + Assert.assertFalse(aiConfig.getEnableThinking()); + + // Test topK + aiConfig.setTopK(100); + Assert.assertEquals(Integer.valueOf(100), aiConfig.getTopK()); + + // Test minP + aiConfig.setMinP(20); + Assert.assertEquals(Integer.valueOf(20), aiConfig.getMinP()); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java new file mode 100644 index 0000000000..b0a7810f77 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java @@ -0,0 +1,259 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.core.shared.exception.BadRequestException; +import com.openai.client.OpenAIClient; +import com.openai.core.JsonValue; +import com.openai.models.chat.completions.ChatCompletion; + +@RunWith(MockitoJUnitRunner.class) +public class AIInferenceServiceTest { + + @Rule + public ExpectedException thrown = ExpectedException.none(); + + @Mock + private AIConfig aiConfig; + + @Mock + private OpenAIClient openAIClient; + + @Mock + private ChatCompletion chatCompletion; + + private AIInferenceService aiInferenceService; + + @Before + public void setUp() { + // Setup mock for openAIClient.chat().completions().create() call chain + // We'll mock this in individual test methods as needed + } + + @Test + public void test_chat_allConditionsMet_callsConfigMethods() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + setupValidAIConfig(); + String systemPrompt = "You are a helpful assistant"; + String userPrompt = "Hello, world!"; + + // When - This will throw an exception due to actual OpenAI call, but we can verify config calls + try { + aiInferenceService.chat(systemPrompt, userPrompt); + } catch (Exception e) { + // Expected - we can't mock the OpenAI client chain easily + } + + // Then - Verify that config methods were called + Mockito.verify(aiConfig).isEnabled(); + Mockito.verify(aiConfig).isAIAvailable(); + } + + @Test + public void test_chat_aiNotEnabled_throwsException() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + Mockito.when(aiConfig.isEnabled()).thenReturn(false); + thrown.expect(BadRequestException.class); + thrown.expectMessage("AI service is not enabled"); + + // When + aiInferenceService.chat("system", "user"); + } + + @Test + public void test_chat_aiNotAvailable_throwsException() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(false); + thrown.expect(BadRequestException.class); + thrown.expectMessage("AI configuration is incomplete"); + + // When + aiInferenceService.chat("system", "user"); + } + + @Test + public void test_chat_clientNotPresent_throwsException() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.empty()); + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(true); + thrown.expect(BadRequestException.class); + thrown.expectMessage("AI client is not initialized"); + + // When + aiInferenceService.chat("system", "user"); + } + + @Test + public void test_chat_clientThrowsException_wrapsException() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + setupValidAIConfig(); + // Note: This test verifies exception wrapping behavior + // The actual OpenAI client will throw an exception due to invalid configuration + thrown.expect(BadRequestException.class); + thrown.expectMessage("Failed to call AI inference service"); + + // When + aiInferenceService.chat("system", "user"); + } + + @Test + public void test_chat_verifyParametersPassedCorrectly() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + setupValidAIConfig(); + String systemPrompt = "You are a helpful assistant"; + String userPrompt = "Hello, world!"; + String model = "gpt-3.5-turbo"; + Double temperature = 0.7; + Double topP = 0.9; + Map additionalParams = new HashMap<>(); + + Mockito.when(aiConfig.getModel()).thenReturn(model); + Mockito.when(aiConfig.getTemperature()).thenReturn(temperature); + Mockito.when(aiConfig.getTopP()).thenReturn(topP); + Mockito.when(aiConfig.loadAdditionalParams()).thenReturn(additionalParams); + + // When - This will throw an exception but we can verify config method calls + try { + aiInferenceService.chat(systemPrompt, userPrompt); + } catch (Exception e) { + // Expected - actual OpenAI call will fail + } + + // Then - Verify config methods were called + Mockito.verify(aiConfig).getModel(); + Mockito.verify(aiConfig).getTemperature(); + Mockito.verify(aiConfig).getTopP(); + Mockito.verify(aiConfig).loadAdditionalParams(); + } + + @Test + public void test_isAIAvailable_allConditionsMet_returnsTrue() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(true); + + // When + boolean result = aiInferenceService.isAIAvailable(); + + // Then + Assert.assertTrue(result); + } + + @Test + public void test_isAIAvailable_aiNotEnabled_returnsFalse() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + Mockito.when(aiConfig.isEnabled()).thenReturn(false); + // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit evaluation + + // When + boolean result = aiInferenceService.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_aiNotAvailable_returnsFalse() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(false); + + // When + boolean result = aiInferenceService.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_clientNotPresent_returnsFalse() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.empty()); + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(true); + + // When + boolean result = aiInferenceService.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_isAIAvailable_noConditionsMet_returnsFalse() { + // Given + aiInferenceService = new AIInferenceService(aiConfig, Optional.empty()); + Mockito.when(aiConfig.isEnabled()).thenReturn(false); + // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit evaluation + + // When + boolean result = aiInferenceService.isAIAvailable(); + + // Then + Assert.assertFalse(result); + } + + @Test + public void test_constructor_withValidParameters_createsInstance() { + // Given & When + AIInferenceService service = new AIInferenceService(aiConfig, Optional.of(openAIClient)); + + // Then + Assert.assertNotNull(service); + } + + @Test + public void test_constructor_withEmptyClient_createsInstance() { + // Given & When + AIInferenceService service = new AIInferenceService(aiConfig, Optional.empty()); + + // Then + Assert.assertNotNull(service); + } + + private void setupValidAIConfig() { + Mockito.when(aiConfig.isEnabled()).thenReturn(true); + Mockito.when(aiConfig.isAIAvailable()).thenReturn(true); + Mockito.when(aiConfig.getModel()).thenReturn("gpt-3.5-turbo"); + Mockito.when(aiConfig.getTemperature()).thenReturn(0.7); + Mockito.when(aiConfig.getTopP()).thenReturn(0.9); + Mockito.when(aiConfig.loadAdditionalParams()).thenReturn(new HashMap<>()); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java new file mode 100644 index 0000000000..9af485d4e6 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java @@ -0,0 +1,221 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.ai; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.junit.MockitoJUnitRunner; +import org.springframework.test.util.ReflectionTestUtils; + +@RunWith(MockitoJUnitRunner.class) +public class PromptTemplateLoaderTest { + + private PromptTemplateLoader promptTemplateLoader; + private static final String MOCK_TEMPLATE = "Sensitive types: {sensitiveTypes}\nCustom prompt: {customPrompt}"; + + @Before + public void setUp() { + promptTemplateLoader = new PromptTemplateLoader(); + // Set mock template to avoid file loading issues in test + ReflectionTestUtils.setField(promptTemplateLoader, "systemTemplate", MOCK_TEMPLATE); + } + + @Test + public void test_buildSystemPrompt_withValidSensitiveTypesAndCustomPrompt_returnsFormattedPrompt() { + // Given + List sensitiveTypes = Arrays.asList("contact_info", "identity_info"); + String customPrompt = "Additional rules for identification"; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain formatted sensitive types", + result.contains("contact_info, identity_info")); + Assert.assertTrue("Should contain custom prompt", + result.contains("Additional rules for identification")); + } + + @Test + public void test_buildSystemPrompt_withEmptySensitiveTypes_returnsDefaultMessage() { + // Given + List sensitiveTypes = Collections.emptyList(); + String customPrompt = "Custom rules"; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain default message for empty types", + result.contains("No specified category.")); + Assert.assertTrue("Should contain custom prompt", result.contains("Custom rules")); + } + + @Test + public void test_buildSystemPrompt_withNullSensitiveTypes_returnsDefaultMessage() { + // Given + List sensitiveTypes = null; + String customPrompt = "Custom rules"; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain default message for null types", + result.contains("No specified category.")); + Assert.assertTrue("Should contain custom prompt", result.contains("Custom rules")); + } + + @Test + public void test_buildSystemPrompt_withEmptyCustomPrompt_returnsDefaultMessage() { + // Given + List sensitiveTypes = Arrays.asList("email", "phone"); + String customPrompt = ""; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain formatted sensitive types", + result.contains("email, phone")); + Assert.assertTrue("Should contain default message for empty prompt", + result.contains("No supplementary rule.")); + } + + @Test + public void test_buildSystemPrompt_withNullCustomPrompt_returnsDefaultMessage() { + // Given + List sensitiveTypes = Arrays.asList("address"); + String customPrompt = null; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain formatted sensitive types", result.contains("address")); + Assert.assertTrue("Should contain default message for null prompt", + result.contains("No supplementary rule.")); + } + + @Test + public void test_buildSystemPrompt_withWhitespaceCustomPrompt_returnsDefaultMessage() { + // Given + List sensitiveTypes = Arrays.asList("name"); + String customPrompt = " \t\n "; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain formatted sensitive types", result.contains("name")); + Assert.assertTrue("Should contain default message for whitespace prompt", + result.contains("No supplementary rule.")); + } + + @Test + public void test_buildSystemPrompt_withSingleSensitiveType_returnsCorrectFormat() { + // Given + List sensitiveTypes = Arrays.asList("credit_card"); + String customPrompt = "Strict validation required"; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain single sensitive type", result.contains("credit_card")); + Assert.assertFalse("Should not contain comma for single type", + result.contains("credit_card,")); + Assert.assertTrue("Should contain custom prompt", + result.contains("Strict validation required")); + } + + @Test(expected = IllegalStateException.class) + public void test_buildSystemPrompt_withNullTemplate_throwsException() { + // Given + ReflectionTestUtils.setField(promptTemplateLoader, "systemTemplate", null); + List sensitiveTypes = Arrays.asList("test"); + String customPrompt = "test"; + + // When + promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then - exception should be thrown + } + + @Test(expected = IllegalStateException.class) + public void test_buildSystemPrompt_withEmptyTemplate_throwsException() { + // Given + ReflectionTestUtils.setField(promptTemplateLoader, "systemTemplate", ""); + List sensitiveTypes = Arrays.asList("test"); + String customPrompt = "test"; + + // When + promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then - exception should be thrown + } + + @Test + public void test_buildSystemPrompt_withMultipleSensitiveTypes_returnsCommaSeparated() { + // Given + List sensitiveTypes = Arrays.asList("email", "phone", "address", "name"); + String customPrompt = "Multiple type validation"; + + // When + String result = promptTemplateLoader.buildSystemPrompt(sensitiveTypes, customPrompt); + + // Then + Assert.assertNotNull("Result should not be null", result); + Assert.assertTrue("Should contain all types comma-separated", + result.contains("email, phone, address, name")); + Assert.assertTrue("Should contain custom prompt", + result.contains("Multiple type validation")); + } + + @Test + public void test_buildSystemPrompt_preservesOriginalTemplate_afterMultipleCalls() { + // Given + List sensitiveTypes1 = Arrays.asList("type1"); + List sensitiveTypes2 = Arrays.asList("type2"); + String customPrompt1 = "prompt1"; + String customPrompt2 = "prompt2"; + + // When + String result1 = promptTemplateLoader.buildSystemPrompt(sensitiveTypes1, customPrompt1); + String result2 = promptTemplateLoader.buildSystemPrompt(sensitiveTypes2, customPrompt2); + + // Then + Assert.assertNotNull("First result should not be null", result1); + Assert.assertNotNull("Second result should not be null", result2); + Assert.assertTrue("First result should contain type1", result1.contains("type1")); + Assert.assertTrue("Second result should contain type2", result2.contains("type2")); + Assert.assertFalse("First result should not contain type2", result1.contains("type2")); + Assert.assertFalse("Second result should not contain type1", result2.contains("type1")); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java new file mode 100644 index 0000000000..ea10732fc7 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java @@ -0,0 +1,203 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.factory; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.odc.service.datasecurity.strategy.AIOnlyStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.JointRecognitionStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.RulesOnlyStrategy; +import com.oceanbase.odc.service.datasecurity.strategy.ScanningStrategy; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +@RunWith(MockitoJUnitRunner.class) +public class ScanningStrategyFactoryTest { + + private ScanningStrategyFactory factory; + private DBTableColumn testColumn; + private List emptyRecognizers; + + @Before + public void setUp() { + factory = new ScanningStrategyFactory(); + testColumn = createTestColumn("user_phone", "varchar", "user phone number"); + emptyRecognizers = Collections.emptyList(); + } + + @Test + public void test_getStrategy_withRulesOnly_returnsRulesOnlyStrategy() { + // When + ScanningStrategy strategy = factory.getStrategy(ScanningModeType.RULES_ONLY); + + // Then + Assert.assertNotNull("Should return a strategy", strategy); + Assert.assertTrue("Should return RulesOnlyStrategy", strategy instanceof RulesOnlyStrategy); + } + + @Test + public void test_getStrategy_withAiOnly_returnsAiOnlyStrategy() { + // When + ScanningStrategy strategy = factory.getStrategy(ScanningModeType.AI_ONLY); + + // Then + Assert.assertNotNull("Should return a strategy", strategy); + Assert.assertTrue("Should return AIOnlyStrategy", strategy instanceof AIOnlyStrategy); + } + + @Test + public void test_getStrategy_withJointRecognition_returnsJointRecognitionStrategy() { + // When + ScanningStrategy strategy = factory.getStrategy(ScanningModeType.JOINT_RECOGNITION); + + // Then + Assert.assertNotNull("Should return a strategy", strategy); + Assert.assertTrue("Should return JointRecognitionStrategy", strategy instanceof JointRecognitionStrategy); + } + + @Test + public void test_getStrategy_withNullMode_returnsNoOpStrategy() { + // When + ScanningStrategy strategy = factory.getStrategy(null); + + // Then + Assert.assertNotNull("Should return a strategy", strategy); + + // Test that it behaves like NoOpStrategy + ScanResult result = strategy.scan(testColumn, emptyRecognizers, emptyRecognizers); + Assert.assertFalse("NoOp strategy should not have basic result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("NoOp strategy should not have AI result", result.getAiRuleResult().isPresent()); + } + + @Test + public void test_getStrategy_returnsSameInstanceForSameMode() { + // When + ScanningStrategy strategy1 = factory.getStrategy(ScanningModeType.RULES_ONLY); + ScanningStrategy strategy2 = factory.getStrategy(ScanningModeType.RULES_ONLY); + + // Then + Assert.assertSame("Should return same instance for same mode", strategy1, strategy2); + } + + @Test + public void test_getStrategy_returnsDifferentInstancesForDifferentModes() { + // When + ScanningStrategy rulesOnlyStrategy = factory.getStrategy(ScanningModeType.RULES_ONLY); + ScanningStrategy aiOnlyStrategy = factory.getStrategy(ScanningModeType.AI_ONLY); + ScanningStrategy jointStrategy = factory.getStrategy(ScanningModeType.JOINT_RECOGNITION); + + // Then + Assert.assertNotSame("Rules and AI strategies should be different", rulesOnlyStrategy, aiOnlyStrategy); + Assert.assertNotSame("Rules and Joint strategies should be different", rulesOnlyStrategy, jointStrategy); + Assert.assertNotSame("AI and Joint strategies should be different", aiOnlyStrategy, jointStrategy); + } + + @Test + public void test_allStrategies_implementScanningStrategy() { + // Given + ScanningModeType[] allModes = {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; + + // When & Then + for (ScanningModeType mode : allModes) { + ScanningStrategy strategy = factory.getStrategy(mode); + Assert.assertNotNull("Strategy should not be null for mode: " + mode, strategy); + Assert.assertTrue("Strategy should implement ScanningStrategy for mode: " + mode, + strategy instanceof ScanningStrategy); + } + } + + @Test + public void test_allStrategies_canHandleBasicOperations() { + // Given + ScanningModeType[] allModes = {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; + List columns = Arrays.asList(testColumn); + + // When & Then + for (ScanningModeType mode : allModes) { + ScanningStrategy strategy = factory.getStrategy(mode); + + // Test single scan + ScanResult singleResult = strategy.scan(testColumn, emptyRecognizers, emptyRecognizers); + Assert.assertNotNull("Single scan result should not be null for mode: " + mode, singleResult); + + // Test batch scan + Map batchResults = strategy.scanBatch(columns, emptyRecognizers, emptyRecognizers); + Assert.assertNotNull("Batch scan results should not be null for mode: " + mode, batchResults); + Assert.assertEquals("Batch scan should return result for each column for mode: " + mode, + 1, batchResults.size()); + } + } + + @Test + public void test_noOpStrategy_handlesBatchScanCorrectly() { + // Given + ScanningStrategy noOpStrategy = factory.getStrategy(null); + DBTableColumn column1 = createTestColumn("col1", "varchar", "comment1"); + DBTableColumn column2 = createTestColumn("col2", "varchar", "comment2"); + List columns = Arrays.asList(column1, column2); + + // When + Map results = noOpStrategy.scanBatch(columns, emptyRecognizers, emptyRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 2, results.size()); + + for (ScanResult result : results.values()) { + Assert.assertFalse("NoOp strategy should not have basic result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("NoOp strategy should not have AI result", result.getAiRuleResult().isPresent()); + } + } + + @Test + public void test_noOpStrategy_handlesNullColumnsGracefully() { + // Given + ScanningStrategy noOpStrategy = factory.getStrategy(null); + DBTableColumn columnWithNulls = new DBTableColumn(); + columnWithNulls.setName(null); + columnWithNulls.setSchemaName(null); + columnWithNulls.setTableName(null); + List columns = Arrays.asList(columnWithNulls); + + // When + Map results = noOpStrategy.scanBatch(columns, emptyRecognizers, emptyRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + Assert.assertTrue("Should contain key for unknown column", + results.containsKey("unknown_schema.unknown_table.unknown_column")); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + return column; + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java new file mode 100644 index 0000000000..50775763c3 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java @@ -0,0 +1,324 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.recognizer; + +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.ExpectedException; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockedStatic; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.core.shared.exception.BadRequestException; +import com.oceanbase.odc.service.common.util.SpringContextUtil; +import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; +import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; +import com.openai.models.chat.completions.ChatCompletion; + + +@RunWith(MockitoJUnitRunner.class) +public class AIColumnRecognizerTest { + + @Rule + public ExpectedException thrown = ExpectedException.none(); + + @Mock + private PromptTemplateLoader promptTemplateLoader; + + @Mock + private AIInferenceService aiInferenceService; + + private AIColumnRecognizer recognizer; + private SensitiveRule aiRule; + + @Before + public void setUp() { + aiRule = createAIRule(); + recognizer = new AIColumnRecognizer(aiRule); + } + + @Test + public void test_recognize_singleColumn_returnsSensitive() { + // Given + DBTableColumn column = createTestColumn("user_phone", "varchar", "user phone number"); + String systemPrompt = "System prompt for AI"; + String aiResponse = "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]\n```"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, aiResponse); + + // When + Optional result = recognizer.recognize(column); + + // Then + Assert.assertTrue(result.isPresent()); + Assert.assertTrue(result.get().isMatched()); + Assert.assertEquals(SensitiveLevel.HIGH, result.get().getLevel()); + Assert.assertEquals("contact_info", result.get().getSensitiveType()); + Assert.assertEquals(SensitiveRuleType.AI, result.get().getSourceRuleType()); + } + } + + @Test + public void test_recognize_singleColumn_returnsNotSensitive() { + // Given + DBTableColumn column = createTestColumn("id", "bigint", "primary key"); + String systemPrompt = "System prompt for AI"; + String aiResponse = "```json\n[{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}]\n```"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, aiResponse); + + // When + Optional result = recognizer.recognize(column); + + // Then + Assert.assertFalse(result.isPresent()); + } + } + + @Test + public void test_recognizeBatch_multipleColumns_returnsMixedResults() { + // Given + List columns = Arrays.asList( + createTestColumn("user_phone", "varchar", "user phone number"), + createTestColumn("id", "bigint", "primary key"), + createTestColumn("email", "varchar", "user email address") + ); + String systemPrompt = "System prompt for AI"; + String aiResponse = "```json\n[" + + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + + "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}," + + "{\"sensitive\": true, \"riskLevel\": \"MEDIUM\", \"sensitiveCategory\": \"contact_info\"}" + + "]```"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, aiResponse); + + // When + Map> results = recognizer.recognizeBatch(columns); + + // Then + Assert.assertEquals(3, results.size()); + + // Check phone column + String phoneKey = getColumnKey(columns.get(0)); + Assert.assertTrue(results.get(phoneKey).isPresent()); + Assert.assertEquals(SensitiveLevel.HIGH, results.get(phoneKey).get().getLevel()); + + // Check id column + String idKey = getColumnKey(columns.get(1)); + Assert.assertFalse(results.get(idKey).isPresent()); + + // Check email column + String emailKey = getColumnKey(columns.get(2)); + Assert.assertTrue(results.get(emailKey).isPresent()); + Assert.assertEquals(SensitiveLevel.MEDIUM, results.get(emailKey).get().getLevel()); + } + } + + @Test + public void test_recognizeBatch_emptyList_returnsEmptyMap() { + // When + Map> results = recognizer.recognizeBatch(Collections.emptyList()); + + // Then + Assert.assertTrue(results.isEmpty()); + } + + @Test + public void test_recognizeBatch_nullList_returnsEmptyMap() { + // When + Map> results = recognizer.recognizeBatch(null); + + // Then + Assert.assertTrue(results.isEmpty()); + } + + @Test + public void test_recognizeBatch_aiServiceThrowsException_returnsEmptyResults() { + // Given + List columns = Arrays.asList(createTestColumn("test", "varchar", "test")); + String systemPrompt = "System prompt for AI"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + mockedSpringContext.when(() -> SpringContextUtil.getBean(PromptTemplateLoader.class)) + .thenReturn(promptTemplateLoader); + mockedSpringContext.when(() -> SpringContextUtil.getBean(AIInferenceService.class)) + .thenReturn(aiInferenceService); + + Mockito.when(promptTemplateLoader.buildSystemPrompt(Mockito.anyList(), Mockito.anyString())) + .thenReturn(systemPrompt); + Mockito.when(aiInferenceService.chat(Mockito.anyString(), Mockito.anyString())) + .thenThrow(new RuntimeException("AI service error")); + + // When + Map> results = recognizer.recognizeBatch(columns); + + // Then + Assert.assertTrue(results.isEmpty()); + } + } + + @Test + public void test_recognizeBatch_invalidJsonResponse_throwsException() { + // Given + List columns = Arrays.asList(createTestColumn("test", "varchar", "test")); + String systemPrompt = "System prompt for AI"; + String invalidResponse = "This is not a valid JSON response"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, invalidResponse); + + thrown.expect(BadRequestException.class); + thrown.expectMessage("AI response does not contain valid JSON format"); + + // When + recognizer.recognizeBatch(columns); + } + } + + @Test + public void test_recognizeBatch_malformedJsonResponse_throwsException() { + // Given + List columns = Arrays.asList(createTestColumn("test", "varchar", "test")); + String systemPrompt = "System prompt for AI"; + String malformedResponse = "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\" missing_comma \"field\": \"value\"}]```"; // Invalid JSON syntax + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, malformedResponse); + + thrown.expect(BadRequestException.class); + thrown.expectMessage("Failed to parse AI response JSON"); + + // When + recognizer.recognizeBatch(columns); + } + } + + @Test + public void test_recognizeBatch_responseWithoutJsonWrapper_parsesCorrectly() { + // Given + DBTableColumn column = createTestColumn("user_phone", "varchar", "user phone number"); + String systemPrompt = "System prompt for AI"; + String aiResponse = "[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]"; // No ```json wrapper + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, aiResponse); + + // When + Optional result = recognizer.recognize(column); + + // Then + Assert.assertTrue(result.isPresent()); + Assert.assertTrue(result.get().isMatched()); + Assert.assertEquals(SensitiveLevel.HIGH, result.get().getLevel()); + } + } + + @Test + public void test_recognizeBatch_mismatchedResponseCount_handlesGracefully() { + // Given + List columns = Arrays.asList( + createTestColumn("col1", "varchar", "test1"), + createTestColumn("col2", "varchar", "test2"), + createTestColumn("col3", "varchar", "test3") + ); + String systemPrompt = "System prompt for AI"; + // AI returns only 2 results for 3 columns + String aiResponse = "```json\n[" + + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + + "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}" + + "]```"; + + try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { + setupMocks(mockedSpringContext, systemPrompt, aiResponse); + + // When + Map> results = recognizer.recognizeBatch(columns); + + // Then + Assert.assertEquals(2, results.size()); // Only 2 results processed + + String col1Key = getColumnKey(columns.get(0)); + String col2Key = getColumnKey(columns.get(1)); + String col3Key = getColumnKey(columns.get(2)); + + Assert.assertTrue(results.containsKey(col1Key)); + Assert.assertTrue(results.containsKey(col2Key)); + Assert.assertFalse(results.containsKey(col3Key)); // Third column not processed + } + } + + // Helper methods + private void setupMocks(MockedStatic mockedSpringContext, String systemPrompt, String aiResponse) { + mockedSpringContext.when(() -> SpringContextUtil.getBean(PromptTemplateLoader.class)) + .thenReturn(promptTemplateLoader); + mockedSpringContext.when(() -> SpringContextUtil.getBean(AIInferenceService.class)) + .thenReturn(aiInferenceService); + + Mockito.when(promptTemplateLoader.buildSystemPrompt(Mockito.anyList(), Mockito.anyString())) + .thenReturn(systemPrompt); + + // Mock the chain: completion.choices().get(0).message().content().orElse("[]") + // Use Mockito's deep stubbing with RETURNS_DEEP_STUBS + ChatCompletion mockCompletion = Mockito.mock(ChatCompletion.class, Mockito.RETURNS_DEEP_STUBS); + Mockito.when(aiInferenceService.chat(Mockito.anyString(), Mockito.anyString())) + .thenReturn(mockCompletion); + Mockito.when(mockCompletion.choices().get(0).message().content()) + .thenReturn(Optional.of(aiResponse)); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + return column; + } + + private SensitiveRule createAIRule() { + SensitiveRule rule = new SensitiveRule(); + rule.setId(1L); + rule.setType(SensitiveRuleType.AI); + rule.setLevel(SensitiveLevel.HIGH); + rule.setEnabled(true); + rule.setAiSensitiveTypes(Arrays.asList("contact_info", "identity_info")); + rule.setAiCustomPrompt("Custom AI prompt for testing"); + return rule; + } + + private String getColumnKey(DBTableColumn column) { + return column.getSchemaName() + "." + column.getTableName() + "." + column.getName(); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java index f13c0c4b57..4eb1c2ebbb 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java @@ -29,27 +29,19 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; -/** - * @author gaoda.xy - * @date 2023/5/23 19:35 - */ + public class GroovyColumnRecognizerTest { - // 原作者使用的 Junit 4 异常测试方式，我们予以保留 + @Rule public ExpectedException thrown = ExpectedException.none(); @Test public void test_recognize_true() { - // 【修改】通过辅助方法创建规则和识别器 SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); DBTableColumn dbTableColumn = createTestColumn(); - - // 【修改】调用新的 recognize 方法并检查 Optional 返回值 Optional resultOpt = recognizer.recognize(dbTableColumn); - - // 【修改】断言结果存在，并验证内容 Assert.assertTrue("脚本匹配成功，应返回有值的 Optional", resultOpt.isPresent()); RecognitionResult result = resultOpt.get(); Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), result.getMatchedRuleId()); @@ -61,11 +53,8 @@ public void test_recognize_false() { SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); DBTableColumn dbTableColumn = createTestColumn(); - dbTableColumn.setTableName("unmatched_table"); // 使脚本匹配失败 - + dbTableColumn.setTableName("unmatched_table"); Optional resultOpt = recognizer.recognize(dbTableColumn); - - // 【修改】断言结果为空 Optional Assert.assertFalse("脚本匹配失败，应返回空的 Optional", resultOpt.isPresent()); } @@ -74,22 +63,16 @@ public void test_recognize_nullColumnName() { SensitiveRule rule = createGroovyRule(1L, buildDefaultGroovyScript()); ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); DBTableColumn dbTableColumn = createTestColumn(); - dbTableColumn.setName(null); // 脚本内部会因为 null.equals(...) 抛异常 - + dbTableColumn.setName(null); Optional resultOpt = recognizer.recognize(dbTableColumn); - - // 【修改】断言结果为空 Optional，因为脚本执行异常被捕获 Assert.assertFalse("脚本执行异常，应返回空的 Optional", resultOpt.isPresent()); } - // --- 以下所有安全校验测试，保留原作者的意图和断言方式 --- - @Test public void test_securityInterceptor_systemExit() { thrown.expect(Exception.class); thrown.expectMessage("Method call is not security"); String script = "System.exit(-1);"; - // 【修改】调用新的构造函数 new GroovyColumnRecognizer(createGroovyRule(1L, script)); } diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java index 21dec644e1..877a6da112 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java @@ -28,23 +28,15 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; -/** - * @author gaoda.xy - * @date 2023/5/24 15:05 - */ public class PathColumnRecognizerTest { @Test public void test_recognize_true() { - // 【修改】通过辅助方法创建规则，并用规则创建识别器 SensitiveRule rule = createPathRule(1L, Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); ColumnRecognizer recognizer = new PathColumnRecognizer(rule); - // 【修改】断言方式改为检查 Optional.isPresent() - // 原作者的每个测试用例都予以保留 Optional result1 = recognizer.recognize(createDBTableColumn("a", "b12", "c")); Assert.assertTrue("路径 'a.b12.c' 应匹配成功", result1.isPresent()); - // 增加对返回内容的校验，使测试更严谨 Assert.assertEquals(rule.getId(), result1.get().getMatchedRuleId()); Optional result2 = recognizer.recognize(createDBTableColumn("a12", "34b56", "c")); @@ -56,12 +48,9 @@ public void test_recognize_true() { @Test public void test_recognize_false() { - // 【修改】通过辅助方法创建规则 SensitiveRule rule = createPathRule(1L, Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); ColumnRecognizer recognizer = new PathColumnRecognizer(rule); - // 【修改】断言方式改为检查 !Optional.isPresent() - // 原作者的每个测试用例都予以保留 Assert.assertFalse("路径 'a.b.c' 应被排除，匹配失败", recognizer.recognize(createDBTableColumn("a", "b", "c")).isPresent()); @@ -72,7 +61,6 @@ public void test_recognize_false() { recognizer.recognize(createDBTableColumn("a12", "b34", null)).isPresent()); } - // --- 辅助方法 --- private DBTableColumn createDBTableColumn(String schemaName, String tableName, String columnName) { DBTableColumn column = new DBTableColumn(); @@ -81,8 +69,6 @@ private DBTableColumn createDBTableColumn(String schemaName, String tableName, S column.setName(columnName); return column; } - - // 【新增】辅助方法，用于快速创建测试用的 Path 规则 private SensitiveRule createPathRule(Long id, List includes, List excludes) { SensitiveRule rule = new SensitiveRule(); rule.setId(id); diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java index 2e19ca422f..3a2cdad4d2 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java @@ -26,35 +26,26 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; -/** - * @author gaoda.xy - * @date 2023/5/24 16:21 - */ + public class RegexColumnRecognizerTest { @Test public void recognize_returnTrue() { - // 【修改】通过辅助方法创建规则，并用规则创建识别器 SensitiveRule rule = createTestRegexRule(1L); ColumnRecognizer recognizer = new RegexColumnRecognizer(rule); DBTableColumn column = createDBTableColumn("xxx", "xxx", "user_email", "email of user"); - // 【修改】调用新的 recognize 方法并检查 Optional 返回值 Optional resultOpt = recognizer.recognize(column); - // 【修改】断言结果存在，并验证内容 Assert.assertTrue("正则表达式应匹配成功", resultOpt.isPresent()); Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), resultOpt.get().getMatchedRuleId()); } @Test public void recognize_returnFalse() { - // 【修改】通过辅助方法创建规则 SensitiveRule rule = createTestRegexRule(1L); ColumnRecognizer recognizer = new RegexColumnRecognizer(rule); - // 【修改】断言方式改为检查 !Optional.isPresent() - // 原作者的每个测试用例都予以保留 Assert.assertFalse("Comment 为 null 时不应匹配", recognizer.recognize(createDBTableColumn("xxx", "xxx", "user_email", null)).isPresent()); @@ -66,16 +57,12 @@ public void recognize_returnFalse() { } - // --- 辅助方法 --- - - // 【修改】原 createRegexColumnRecognizer 替换为 createTestRegexRule private SensitiveRule createTestRegexRule(Long id) { SensitiveRule rule = new SensitiveRule(); rule.setId(id); rule.setType(SensitiveRuleType.REGEX); rule.setEnabled(true); rule.setLevel(SensitiveLevel.HIGH); - // 保留原作者的正则表达式 rule.setDatabaseRegexExpression("^\\S+$"); rule.setTableRegexExpression("^\\S+$"); rule.setColumnRegexExpression("^\\S*email\\S*$"); diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java new file mode 100644 index 0000000000..3a27836924 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java @@ -0,0 +1,229 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +@RunWith(MockitoJUnitRunner.class) +public class AIOnlyStrategyTest { + + @Mock + private ColumnRecognizer mockBasicRecognizer; + + @Mock + private ColumnRecognizer mockAiRecognizer; + + private AIOnlyStrategy strategy; + private List basicRecognizers; + private List aiRecognizers; + private DBTableColumn testColumn; + + @Before + public void setUp() { + strategy = new AIOnlyStrategy(); + basicRecognizers = Arrays.asList(mockBasicRecognizer); + aiRecognizers = Arrays.asList(mockAiRecognizer); + testColumn = createTestColumn("user_phone", "varchar", "user phone number"); + } + + @Test + public void test_scan_withAiRecognizerMatch_returnsAiResultOnly() { + // Given + RecognitionResult aiResult = createRecognitionResult(2L, SensitiveLevel.MEDIUM, SensitiveRuleType.AI); + + Mockito.when(mockAiRecognizer.recognize(testColumn)).thenReturn(Optional.of(aiResult)); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertTrue("Should have AI rule result", result.getAiRuleResult().isPresent()); + Assert.assertEquals("Should return AI result", aiResult, result.getAiRuleResult().get()); + + // Verify basic recognizer is not called (AI only strategy) + Mockito.verify(mockBasicRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scan_withNoAiRecognizerMatch_returnsEmptyResult() { + // Given + Mockito.when(mockAiRecognizer.recognize(testColumn)).thenReturn(Optional.empty()); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + + // Verify basic recognizer is not called + Mockito.verify(mockBasicRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scan_withEmptyAiRecognizers_returnsEmptyResult() { + // Given + List emptyAiRecognizers = Collections.emptyList(); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, emptyAiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + } + + @Test + public void test_scanBatch_withMixedResults_returnsCorrectMapping() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + DBTableColumn column3 = createTestColumn("user_name", "varchar", "name"); + List columns = Arrays.asList(column1, column2, column3); + + RecognitionResult aiResult1 = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.AI); + RecognitionResult aiResult3 = createRecognitionResult(3L, SensitiveLevel.LOW, SensitiveRuleType.AI); + + // Mock recognizeBatch for single recognizer scenario + Map> batchResults = new HashMap<>(); + batchResults.put(getColumnKey(column1), Optional.of(aiResult1)); + batchResults.put(getColumnKey(column2), Optional.empty()); + batchResults.put(getColumnKey(column3), Optional.of(aiResult3)); + Mockito.when(mockAiRecognizer.recognizeBatch(columns)).thenReturn(batchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 3, results.size()); + + String key1 = getColumnKey(column1); + String key2 = getColumnKey(column2); + String key3 = getColumnKey(column3); + + Assert.assertFalse("Column1 should not have basic result", results.get(key1).getBasicRuleResult().isPresent()); + Assert.assertTrue("Column1 should have AI result", results.get(key1).getAiRuleResult().isPresent()); + + Assert.assertFalse("Column2 should not have basic result", results.get(key2).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column2 should not have AI result", results.get(key2).getAiRuleResult().isPresent()); + + Assert.assertFalse("Column3 should not have basic result", results.get(key3).getBasicRuleResult().isPresent()); + Assert.assertTrue("Column3 should have AI result", results.get(key3).getAiRuleResult().isPresent()); + + // Verify basic recognizer is never called + Mockito.verify(mockBasicRecognizer, Mockito.never()).recognize(Mockito.any()); + Mockito.verify(mockBasicRecognizer, Mockito.never()).recognizeBatch(Mockito.any()); + } + + @Test + public void test_scanBatch_withEmptyColumns_returnsEmptyMap() { + // Given + List emptyColumns = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(emptyColumns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertTrue("Should return empty map", results.isEmpty()); + } + + @Test + public void test_scanBatch_withEmptyAiRecognizers_returnsEmptyResults() { + // Given + List columns = Arrays.asList(testColumn); + List emptyAiRecognizers = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, emptyAiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + String key = getColumnKey(testColumn); + Assert.assertFalse("Should not have basic result", results.get(key).getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI result", results.get(key).getAiRuleResult().isPresent()); + } + + @Test + public void test_scanBatch_withSingleAiRecognizer_usesBatchRecognition() { + // Given + List columns = Arrays.asList(testColumn); + Map> batchResults = Collections.singletonMap( + getColumnKey(testColumn), + Optional.of(createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.AI)) + ); + + Mockito.when(mockAiRecognizer.recognizeBatch(columns)).thenReturn(batchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + String key = getColumnKey(testColumn); + Assert.assertTrue("Should have AI result", results.get(key).getAiRuleResult().isPresent()); + + // Verify batch recognition is used + Mockito.verify(mockAiRecognizer).recognizeBatch(columns); + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + return column; + } + + private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { + return RecognitionResult.builder() + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); + } + + private String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java new file mode 100644 index 0000000000..69e9fc3062 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java @@ -0,0 +1,249 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +@RunWith(MockitoJUnitRunner.class) +public class AbstractScanningStrategyTest { + + @Mock + private ColumnRecognizer mockRecognizer1; + + @Mock + private ColumnRecognizer mockRecognizer2; + + private TestableAbstractScanningStrategy strategy; + private DBTableColumn testColumn; + + @Before + public void setUp() { + strategy = new TestableAbstractScanningStrategy(); + testColumn = createTestColumn("user_phone", "varchar", "user phone number"); + } + + @Test + public void test_findFirstMatch_withMatchingRecognizer_returnsResult() { + // Given + RecognitionResult expectedResult = createRecognitionResult(1L, SensitiveLevel.HIGH); + List recognizers = Arrays.asList(mockRecognizer1, mockRecognizer2); + + Mockito.when(mockRecognizer1.recognize(testColumn)).thenReturn(Optional.empty()); + Mockito.when(mockRecognizer2.recognize(testColumn)).thenReturn(Optional.of(expectedResult)); + + // When + Optional result = strategy.findFirstMatch(recognizers, testColumn); + + // Then + Assert.assertTrue("Should find matching result", result.isPresent()); + Assert.assertEquals("Should return expected result", expectedResult, result.get()); + } + + @Test + public void test_findFirstMatch_withNoMatchingRecognizer_returnsEmpty() { + // Given + List recognizers = Arrays.asList(mockRecognizer1, mockRecognizer2); + + Mockito.when(mockRecognizer1.recognize(testColumn)).thenReturn(Optional.empty()); + Mockito.when(mockRecognizer2.recognize(testColumn)).thenReturn(Optional.empty()); + + // When + Optional result = strategy.findFirstMatch(recognizers, testColumn); + + // Then + Assert.assertFalse("Should not find any result", result.isPresent()); + } + + @Test + public void test_findFirstMatch_withEmptyRecognizers_returnsEmpty() { + // Given + List recognizers = Collections.emptyList(); + + // When + Optional result = strategy.findFirstMatch(recognizers, testColumn); + + // Then + Assert.assertFalse("Should not find any result with empty recognizers", result.isPresent()); + } + + @Test + public void test_findAllFirstMatches_withMultipleColumns_returnsCorrectMapping() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + List columns = Arrays.asList(column1, column2); + List recognizers = Arrays.asList(mockRecognizer1); + + RecognitionResult result1 = createRecognitionResult(1L, SensitiveLevel.HIGH); + + // Mock recognizeBatch method for single recognizer scenario + Map> batchResults = new HashMap<>(); + batchResults.put(strategy.getColumnKey(column1), Optional.of(result1)); + batchResults.put(strategy.getColumnKey(column2), Optional.empty()); + Mockito.when(mockRecognizer1.recognizeBatch(columns)).thenReturn(batchResults); + + // When + Map> results = strategy.findAllFirstMatches(recognizers, columns); + + // Then + Assert.assertEquals("Should return results for all columns", 2, results.size()); + Assert.assertTrue("Should find result for column1", results.get(strategy.getColumnKey(column1)).isPresent()); + Assert.assertFalse("Should not find result for column2", results.get(strategy.getColumnKey(column2)).isPresent()); + } + + @Test + public void test_findAllFirstMatches_withEmptyColumns_returnsEmptyMap() { + // Given + List columns = Collections.emptyList(); + List recognizers = Arrays.asList(mockRecognizer1); + + // When + Map> results = strategy.findAllFirstMatches(recognizers, columns); + + // Then + Assert.assertTrue("Should return empty map", results.isEmpty()); + } + + @Test + public void test_findAllFirstMatches_withEmptyRecognizers_returnsEmptyResults() { + // Given + List columns = Arrays.asList(testColumn); + List recognizers = Collections.emptyList(); + + // When + Map> results = strategy.findAllFirstMatches(recognizers, columns); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + Assert.assertFalse("Should not find any result", results.get(strategy.getColumnKey(testColumn)).isPresent()); + } + + @Test + public void test_findAllFirstMatches_withMultipleRecognizers_usesIndividualRecognize() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + List columns = Arrays.asList(column1, column2); + List recognizers = Arrays.asList(mockRecognizer1, mockRecognizer2); + + RecognitionResult result1 = createRecognitionResult(1L, SensitiveLevel.HIGH); + + // Mock individual recognize calls for multiple recognizers scenario + Mockito.when(mockRecognizer1.recognize(column1)).thenReturn(Optional.empty()); + Mockito.when(mockRecognizer2.recognize(column1)).thenReturn(Optional.of(result1)); + Mockito.when(mockRecognizer1.recognize(column2)).thenReturn(Optional.empty()); + Mockito.when(mockRecognizer2.recognize(column2)).thenReturn(Optional.empty()); + + // When + Map> results = strategy.findAllFirstMatches(recognizers, columns); + + // Then + Assert.assertEquals("Should return results for all columns", 2, results.size()); + Assert.assertTrue("Should find result for column1", results.get(strategy.getColumnKey(column1)).isPresent()); + Assert.assertFalse("Should not find result for column2", results.get(strategy.getColumnKey(column2)).isPresent()); + + // Verify that recognizeBatch was not called for multiple recognizers + Mockito.verify(mockRecognizer1, Mockito.never()).recognizeBatch(Mockito.any()); + Mockito.verify(mockRecognizer2, Mockito.never()).recognizeBatch(Mockito.any()); + } + + @Test + public void test_getColumnKey_returnsCorrectFormat() { + // Given + DBTableColumn column = createTestColumn("user_phone", "varchar", "phone"); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + + // When + String key = strategy.getColumnKey(column); + + // Then + Assert.assertEquals("Should return correct column key format", "test_schema.test_table.user_phone", key); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + return column; + } + + private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level) { + return RecognitionResult.builder() + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(SensitiveRuleType.REGEX) + .build(); + } + + // Testable implementation of AbstractScanningStrategy for testing protected methods + private static class TestableAbstractScanningStrategy extends AbstractScanningStrategy { + @Override + public com.oceanbase.odc.service.datasecurity.model.ScanResult scan( + DBTableColumn column, + List basicRecognizers, + List aiRecognizers) { + return null; // Not used in these tests + } + + @Override + public Map scanBatch( + List columns, + List basicRecognizers, + List aiRecognizers) { + return null; // Not used in these tests + } + + // Expose protected methods for testing + @Override + public Optional findFirstMatch(List recognizers, DBTableColumn column) { + return super.findFirstMatch(recognizers, column); + } + + @Override + public Map> findAllFirstMatches(List recognizers, List columns) { + return super.findAllFirstMatches(recognizers, columns); + } + + @Override + public String getColumnKey(DBTableColumn column) { + return super.getColumnKey(column); + } + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java new file mode 100644 index 0000000000..1e1f492937 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java @@ -0,0 +1,295 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + +@RunWith(MockitoJUnitRunner.class) +public class JointRecognitionStrategyTest { + + @Mock + private ColumnRecognizer mockBasicRecognizer; + + @Mock + private ColumnRecognizer mockAiRecognizer; + + private JointRecognitionStrategy strategy; + private List basicRecognizers; + private List aiRecognizers; + private DBTableColumn testColumn; + + @Before + public void setUp() { + strategy = new JointRecognitionStrategy(); + basicRecognizers = Arrays.asList(mockBasicRecognizer); + aiRecognizers = Arrays.asList(mockAiRecognizer); + testColumn = createTestColumn("user_phone", "varchar", "user phone number"); + } + + @Test + public void test_scan_withBasicRecognizerMatch_returnsBasicResultAndSkipsAi() { + // Given + RecognitionResult basicResult = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.REGEX); + + Mockito.when(mockBasicRecognizer.recognize(testColumn)).thenReturn(Optional.of(basicResult)); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertTrue("Should have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + Assert.assertEquals("Should return basic result", basicResult, result.getBasicRuleResult().get()); + + // Verify AI recognizer is not called when basic recognizer matches + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scan_withNoBasicMatchButAiMatch_returnsAiResult() { + // Given + RecognitionResult aiResult = createRecognitionResult(2L, SensitiveLevel.MEDIUM, SensitiveRuleType.AI); + + Mockito.when(mockBasicRecognizer.recognize(testColumn)).thenReturn(Optional.empty()); + Mockito.when(mockAiRecognizer.recognize(testColumn)).thenReturn(Optional.of(aiResult)); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertTrue("Should have AI rule result", result.getAiRuleResult().isPresent()); + Assert.assertEquals("Should return AI result", aiResult, result.getAiRuleResult().get()); + + // Verify AI recognizer is called as fallback + Mockito.verify(mockAiRecognizer).recognize(testColumn); + } + + @Test + public void test_scan_withNoMatches_returnsEmptyResult() { + // Given + Mockito.when(mockBasicRecognizer.recognize(testColumn)).thenReturn(Optional.empty()); + Mockito.when(mockAiRecognizer.recognize(testColumn)).thenReturn(Optional.empty()); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + + // Verify both recognizers are called + Mockito.verify(mockBasicRecognizer).recognize(testColumn); + Mockito.verify(mockAiRecognizer).recognize(testColumn); + } + + @Test + public void test_scanBatch_withMixedResults_optimizesAiCalls() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + DBTableColumn column3 = createTestColumn("user_name", "varchar", "name"); + DBTableColumn column4 = createTestColumn("user_id", "bigint", "id"); + List columns = Arrays.asList(column1, column2, column3, column4); + + // Basic recognizer matches column1 and column3 + RecognitionResult basicResult1 = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.REGEX); + RecognitionResult basicResult3 = createRecognitionResult(3L, SensitiveLevel.LOW, SensitiveRuleType.GROOVY); + + // Mock recognizeBatch for basic recognizer (single recognizer scenario) + Map> basicBatchResults = new HashMap<>(); + basicBatchResults.put(getColumnKey(column1), Optional.of(basicResult1)); + basicBatchResults.put(getColumnKey(column2), Optional.empty()); + basicBatchResults.put(getColumnKey(column3), Optional.of(basicResult3)); + basicBatchResults.put(getColumnKey(column4), Optional.empty()); + Mockito.when(mockBasicRecognizer.recognizeBatch(columns)).thenReturn(basicBatchResults); + + // AI recognizer only matches column2 (column4 has no match) + RecognitionResult aiResult2 = createRecognitionResult(2L, SensitiveLevel.MEDIUM, SensitiveRuleType.AI); + + // Mock recognizeBatch for AI recognizer on remaining columns (column2, column4) + List remainingColumns = Arrays.asList(column2, column4); + Map> aiBatchResults = new HashMap<>(); + aiBatchResults.put(getColumnKey(column2), Optional.of(aiResult2)); + aiBatchResults.put(getColumnKey(column4), Optional.empty()); + Mockito.when(mockAiRecognizer.recognizeBatch(remainingColumns)).thenReturn(aiBatchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 4, results.size()); + + String key1 = getColumnKey(column1); + String key2 = getColumnKey(column2); + String key3 = getColumnKey(column3); + String key4 = getColumnKey(column4); + + // Column1: basic match, no AI call + Assert.assertTrue("Column1 should have basic result", results.get(key1).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column1 should not have AI result", results.get(key1).getAiRuleResult().isPresent()); + + // Column2: no basic match, AI match + Assert.assertFalse("Column2 should not have basic result", results.get(key2).getBasicRuleResult().isPresent()); + Assert.assertTrue("Column2 should have AI result", results.get(key2).getAiRuleResult().isPresent()); + + // Column3: basic match, no AI call + Assert.assertTrue("Column3 should have basic result", results.get(key3).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column3 should not have AI result", results.get(key3).getAiRuleResult().isPresent()); + + // Column4: no matches + Assert.assertFalse("Column4 should not have basic result", results.get(key4).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column4 should not have AI result", results.get(key4).getAiRuleResult().isPresent()); + + // Verify AI recognizer is only called for columns without basic matches + Mockito.verify(mockAiRecognizer).recognizeBatch(remainingColumns); + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scanBatch_withAllBasicMatches_skipsAiCompletely() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + List columns = Arrays.asList(column1, column2); + + RecognitionResult basicResult1 = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.REGEX); + RecognitionResult basicResult2 = createRecognitionResult(2L, SensitiveLevel.MEDIUM, SensitiveRuleType.GROOVY); + + // Mock recognizeBatch for basic recognizer (all matches) + Map> basicBatchResults = new HashMap<>(); + basicBatchResults.put(getColumnKey(column1), Optional.of(basicResult1)); + basicBatchResults.put(getColumnKey(column2), Optional.of(basicResult2)); + Mockito.when(mockBasicRecognizer.recognizeBatch(columns)).thenReturn(basicBatchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 2, results.size()); + + // Verify AI recognizer is never called + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + Mockito.verify(mockAiRecognizer, Mockito.never()).recognizeBatch(Mockito.any()); + } + + @Test + public void test_scanBatch_withNoBasicMatches_callsAiForAll() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + List columns = Arrays.asList(column1, column2); + + // Mock recognizeBatch for basic recognizer (no matches) + Map> basicBatchResults = new HashMap<>(); + basicBatchResults.put(getColumnKey(column1), Optional.empty()); + basicBatchResults.put(getColumnKey(column2), Optional.empty()); + Mockito.when(mockBasicRecognizer.recognizeBatch(columns)).thenReturn(basicBatchResults); + + RecognitionResult aiResult1 = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.AI); + + // Mock recognizeBatch for AI recognizer (all columns since no basic matches) + Map> aiBatchResults = new HashMap<>(); + aiBatchResults.put(getColumnKey(column1), Optional.of(aiResult1)); + aiBatchResults.put(getColumnKey(column2), Optional.empty()); + Mockito.when(mockAiRecognizer.recognizeBatch(columns)).thenReturn(aiBatchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 2, results.size()); + + // Verify AI recognizer is called for all columns + Mockito.verify(mockAiRecognizer).recognizeBatch(columns); + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scanBatch_withEmptyColumns_returnsEmptyMap() { + // Given + List emptyColumns = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(emptyColumns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertTrue("Should return empty map", results.isEmpty()); + } + + @Test + public void test_scanBatch_withEmptyRecognizers_returnsEmptyResults() { + // Given + List columns = Arrays.asList(testColumn); + List emptyBasicRecognizers = Collections.emptyList(); + List emptyAiRecognizers = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(columns, emptyBasicRecognizers, emptyAiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + String key = getColumnKey(testColumn); + Assert.assertFalse("Should not have basic result", results.get(key).getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI result", results.get(key).getAiRuleResult().isPresent()); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + return column; + } + + private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { + return RecognitionResult.builder() + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); + } + + private String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + } +} \ No newline at end of file diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java new file mode 100644 index 0000000000..e9a20e2120 --- /dev/null +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java @@ -0,0 +1,206 @@ +/* + * Copyright (c) 2025 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.oceanbase.odc.service.datasecurity.strategy; + +import java.util.Arrays; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; + +import org.junit.Assert; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.junit.MockitoJUnitRunner; + +import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; +import com.oceanbase.odc.service.datasecurity.model.ScanResult; +import com.oceanbase.odc.service.datasecurity.model.SensitiveLevel; +import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; +import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; +import com.oceanbase.tools.dbbrowser.model.DBTableColumn; + + +@RunWith(MockitoJUnitRunner.class) +public class RulesOnlyStrategyTest { + + @Mock + private ColumnRecognizer mockBasicRecognizer; + + @Mock + private ColumnRecognizer mockAiRecognizer; + + private RulesOnlyStrategy strategy; + private List basicRecognizers; + private List aiRecognizers; + private DBTableColumn testColumn; + + @Before + public void setUp() { + strategy = new RulesOnlyStrategy(); + basicRecognizers = Arrays.asList(mockBasicRecognizer); + aiRecognizers = Arrays.asList(mockAiRecognizer); + testColumn = createTestColumn("user_phone", "varchar", "user phone number"); + } + + @Test + public void test_scan_withBasicRecognizerMatch_returnsBasicResultOnly() { + // Given + RecognitionResult basicResult = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.REGEX); + + Mockito.when(mockBasicRecognizer.recognize(testColumn)).thenReturn(Optional.of(basicResult)); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertTrue("Should have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + Assert.assertEquals("Should return basic result", basicResult, result.getBasicRuleResult().get()); + + // Verify AI recognizer is not called (rules only strategy) + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scan_withNoBasicRecognizerMatch_returnsEmptyResult() { + // Given + Mockito.when(mockBasicRecognizer.recognize(testColumn)).thenReturn(Optional.empty()); + + // When + ScanResult result = strategy.scan(testColumn, basicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + + // Verify AI recognizer is not called + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + } + + @Test + public void test_scan_withEmptyBasicRecognizers_returnsEmptyResult() { + // Given + List emptyBasicRecognizers = Collections.emptyList(); + + // When + ScanResult result = strategy.scan(testColumn, emptyBasicRecognizers, aiRecognizers); + + // Then + Assert.assertFalse("Should not have basic rule result", result.getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI rule result", result.getAiRuleResult().isPresent()); + } + + @Test + public void test_scanBatch_withMixedResults_returnsCorrectMapping() { + // Given + DBTableColumn column1 = createTestColumn("user_phone", "varchar", "phone"); + DBTableColumn column2 = createTestColumn("user_email", "varchar", "email"); + DBTableColumn column3 = createTestColumn("user_name", "varchar", "name"); + List columns = Arrays.asList(column1, column2, column3); + + RecognitionResult result1 = createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.REGEX); + RecognitionResult result3 = createRecognitionResult(3L, SensitiveLevel.LOW, SensitiveRuleType.GROOVY); + + // Mock recognizeBatch for single recognizer scenario + Map> batchResults = new HashMap<>(); + batchResults.put(getColumnKey(column1), Optional.of(result1)); + batchResults.put(getColumnKey(column2), Optional.empty()); + batchResults.put(getColumnKey(column3), Optional.of(result3)); + Mockito.when(mockBasicRecognizer.recognizeBatch(columns)).thenReturn(batchResults); + + // When + Map results = strategy.scanBatch(columns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 3, results.size()); + + String key1 = getColumnKey(column1); + String key2 = getColumnKey(column2); + String key3 = getColumnKey(column3); + + Assert.assertTrue("Column1 should have basic result", results.get(key1).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column1 should not have AI result", results.get(key1).getAiRuleResult().isPresent()); + + Assert.assertFalse("Column2 should not have basic result", results.get(key2).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column2 should not have AI result", results.get(key2).getAiRuleResult().isPresent()); + + Assert.assertTrue("Column3 should have basic result", results.get(key3).getBasicRuleResult().isPresent()); + Assert.assertFalse("Column3 should not have AI result", results.get(key3).getAiRuleResult().isPresent()); + + // Verify AI recognizer is never called + Mockito.verify(mockAiRecognizer, Mockito.never()).recognize(Mockito.any()); + Mockito.verify(mockAiRecognizer, Mockito.never()).recognizeBatch(Mockito.any()); + } + + @Test + public void test_scanBatch_withEmptyColumns_returnsEmptyMap() { + // Given + List emptyColumns = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(emptyColumns, basicRecognizers, aiRecognizers); + + // Then + Assert.assertTrue("Should return empty map", results.isEmpty()); + } + + @Test + public void test_scanBatch_withEmptyBasicRecognizers_returnsEmptyResults() { + // Given + List columns = Arrays.asList(testColumn); + List emptyBasicRecognizers = Collections.emptyList(); + + // When + Map results = strategy.scanBatch(columns, emptyBasicRecognizers, aiRecognizers); + + // Then + Assert.assertEquals("Should return results for all columns", 1, results.size()); + String key = getColumnKey(testColumn); + Assert.assertFalse("Should not have basic result", results.get(key).getBasicRuleResult().isPresent()); + Assert.assertFalse("Should not have AI result", results.get(key).getAiRuleResult().isPresent()); + } + + private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { + DBTableColumn column = new DBTableColumn(); + column.setName(columnName); + column.setTypeName(typeName); + column.setComment(comment); + column.setSchemaName("test_schema"); + column.setTableName("test_table"); + return column; + } + + private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { + return RecognitionResult.builder() + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); + } + + private String getColumnKey(DBTableColumn column) { + return String.format("%s.%s.%s", + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); + } +} \ No newline at end of file From 10e1315fa6c1c1a0ef036b5325ce18401be5b640 Mon Sep 17 00:00:00 2001 From: fenyf Date: Mon, 8 Sep 2025 13:20:47 +0800 Subject: [PATCH 07/10] fix(ai_recognition): Code formatting, correction of comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(ai_recognition):AI状态查询bug修复 --- .../web/controller/v2/AIController.java | 21 +-- .../datasecurity/SensitiveColumnScanner.java | 23 +-- .../SensitiveColumnScanningTask.java | 140 ++++++------------ .../SensitiveColumnScanningTaskManager.java | 2 +- .../datasecurity/SensitiveColumnService.java | 87 ++++------- .../SingleTableScanTaskManager.java | 37 +---- .../odc/service/datasecurity/ai/AIConfig.java | 21 +-- .../datasecurity/ai/AIInferenceService.java | 52 +++---- .../odc/service/datasecurity/ai/AIParam.java | 6 +- .../datasecurity/ai/AIStatusResponse.java | 36 ++--- .../datasecurity/ai/PromptTemplateLoader.java | 38 ++--- .../factory/ScanningStrategyFactory.java | 20 +-- .../model/DefaultSensitiveType.java | 79 +--------- .../datasecurity/model/ScanResult.java | 20 +-- .../model/SingleTableScanReq.java | 20 +-- .../recognizer/AIColumnRecognizer.java | 69 ++++----- .../recognizer/ColumnRecognizer.java | 9 +- .../recognizer/GroovyColumnRecognizer.java | 2 - .../recognizer/PathColumnRecognizer.java | 2 - .../recognizer/RegexColumnRecognizer.java | 19 ++- .../datasecurity/strategy/AIOnlyStrategy.java | 9 +- .../strategy/AbstractScanningStrategy.java | 38 +---- .../strategy/JointRecognitionStrategy.java | 18 +-- .../strategy/RulesOnlyStrategy.java | 9 +- .../strategy/ScanningStrategy.java | 24 +-- .../service/datasecurity/ai/AIConfigTest.java | 2 +- .../ai/AIInferenceServiceTest.java | 10 +- .../ai/PromptTemplateLoaderTest.java | 26 ++-- .../factory/ScanningStrategyFactoryTest.java | 14 +- .../recognizer/AIColumnRecognizerTest.java | 63 ++++---- .../GroovyColumnRecognizerTest.java | 43 +++--- .../recognizer/PathColumnRecognizerTest.java | 13 +- .../recognizer/RegexColumnRecognizerTest.java | 10 +- .../strategy/AIOnlyStrategyTest.java | 23 ++- .../AbstractScanningStrategyTest.java | 33 +++-- .../JointRecognitionStrategyTest.java | 18 +-- .../strategy/RulesOnlyStrategyTest.java | 18 +-- 37 files changed, 376 insertions(+), 698 deletions(-) diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java index ce4960d2d6..f8d702ae6f 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java @@ -25,16 +25,17 @@ import com.oceanbase.odc.service.common.response.Responses; import com.oceanbase.odc.service.common.response.SuccessResponse; import com.oceanbase.odc.service.datasecurity.ai.AIConfig; +import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; import com.oceanbase.odc.service.datasecurity.ai.AIStatusResponse; import io.swagger.annotations.Api; import io.swagger.annotations.ApiOperation; /** - * AI功能控制器 - * 提供AI功能状态查询接口 + * @author fenyf + * @date 2025/8/10 12:41 */ -@Api(tags = "AI功能") +@Api(tags = "AI") @RestController @RequestMapping("/api/v2/ai") public class AIController { @@ -42,20 +43,20 @@ public class AIController { @Autowired private AIConfig aiConfig; - /** - * 查询AI功能状态 - * @return AI功能状态信息 - */ - @ApiOperation(value = "查询AI功能状态", notes = "返回AI功能是否启用以及配置状态") + @Autowired + private AIInferenceService aiInferenceService; + + + @ApiOperation(value = "Query the status of the AI function", + notes = "Return the status of whether the AI function is enabled and its configuration status") @SkipAuthorize("AI status is safe to query for authenticated users") @GetMapping("/status") public SuccessResponse getAIStatus() { AIStatusResponse response = new AIStatusResponse(); response.setEnabled(aiConfig.isEnabled()); - response.setAvailable(aiConfig.isAIAvailable()); + response.setAvailable(aiInferenceService.isAIAvailable()); response.setModel(aiConfig.getModel()); response.setBaseUrl(aiConfig.getBaseUrl()); - // 不返回敏感信息如API密钥 response.setApiKeyConfigured(aiConfig.getApiKey() != null && !aiConfig.getApiKey().trim().isEmpty()); return Responses.success(response); diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java index 8cc51db195..4788d4c7cc 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java @@ -31,17 +31,16 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** - * 敏感列识别的编排器，负责根据不同的扫描模式执行识别策略。 - * 使用策略模式重构，消除重复代码。 + * @author fenyf + * @date 2025/8/10 12:41 */ public class SensitiveColumnScanner { private final List basicRecognizers; - private final List aiRecognizers; + private final List aiRecognizers; private final ScanningStrategyFactory strategyFactory; public SensitiveColumnScanner(List rules, ScanningStrategyFactory strategyFactory) { - // 在构造时，就将规则分好类，并创建对应的识别器 this.basicRecognizers = rules.stream() .filter(r -> r.getType() != SensitiveRuleType.AI) .map(ColumnRecognizerFactory::create) @@ -53,27 +52,13 @@ public SensitiveColumnScanner(List rules, ScanningStrategyFactory this.strategyFactory = strategyFactory; } - /** - * 核心扫描方法 - * - * @param column 待扫描的列 - * @param mode 用户选择的扫描模式 - * @return 包含一个或两个结果的最终扫描报告 - */ public ScanResult scan(DBTableColumn column, ScanningModeType mode) { ScanningStrategy strategy = strategyFactory.getStrategy(mode); return strategy.scan(column, basicRecognizers, aiRecognizers); } - /** - * 批量扫描方法 - * - * @param columns 待扫描的列列表 - * @param mode 用户选择的扫描模式 - * @return 扫描结果映射，key为列名，value为扫描结果 - */ public Map scanBatch(List columns, ScanningModeType mode) { ScanningStrategy strategy = strategyFactory.getStrategy(mode); return strategy.scanBatch(columns, basicRecognizers, aiRecognizers); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java index 2665991fec..25b1551c2d 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java @@ -63,10 +63,9 @@ public class SensitiveColumnScanningTask implements Callable { private final Map ruleMap; public SensitiveColumnScanningTask(Database database, List rules, ScanningModeType scanningMode, - SensitiveColumnScanningTaskInfo taskInfo, List existsSensitiveColumns, - Map> table2Columns, Map> view2Columns) { + SensitiveColumnScanningTaskInfo taskInfo, List existsSensitiveColumns, + Map> table2Columns, Map> view2Columns) { this.database = database; - // 【修改】接收扫描模式，并创建新的扫描器 this.scanningMode = scanningMode; ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); this.scanner = new SensitiveColumnScanner(rules, strategyFactory); @@ -74,25 +73,20 @@ public SensitiveColumnScanningTask(Database database, List rules, this.view2Columns = view2Columns; this.taskInfo = taskInfo; this.existsSensitiveColumns = new HashSet<>(existsSensitiveColumns); - // 【修改】将规则列表转换为 Map，方便通过 ID 快速查找 this.ruleMap = rules.stream().collect(Collectors.toMap(SensitiveRule::getId, Function.identity())); } - /** - * 生成列的唯一标识符 - */ private String getColumnKey(DBTableColumn column) { return String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } @Override public Void call() { try { taskInfo.setStatus(ScanningTaskStatus.RUNNING); - // 调用重构后的 scanColumns 方法 scanColumns(table2Columns, SensitiveColumnType.TABLE_COLUMN); if (taskInfo.isCancelled()) { return null; @@ -103,136 +97,105 @@ public Void call() { taskInfo.setStatus(ScanningTaskStatus.FAILED); taskInfo.setErrorCode(ErrorCodes.Unexpected); taskInfo.setErrorMsg(String.format("Error during sensitive column scanning on database=%s, reason=%s", - database.getName(), e.getMessage())); + database.getName(), e.getMessage())); taskInfo.setCompleteTime(new Date()); } } return null; } - // 【修改】scanColumns 的核心逻辑改为批量扫描，并支持表级别并发 private void scanColumns(Map> object2Columns, SensitiveColumnType columnType) { if (object2Columns.isEmpty()) { return; } - // 表级别并发处理：为每个表创建异步任务 List> tableFutures = object2Columns.entrySet().stream() - .map(entry -> CompletableFuture.runAsync(() -> { - String objectName = entry.getKey(); - List columns = entry.getValue(); + .map(entry -> CompletableFuture.runAsync(() -> { + String objectName = entry.getKey(); + List columns = entry.getValue(); - try { - // 检查是否已被中断 - if (taskInfo.isCancelled()) { - return; - } - // 【改为批量扫描】一次性扫描整个表的所有列 - Map scanResults = this.scanner.scanBatch(columns, this.scanningMode); - - // 再次检查是否已被中断 - if (taskInfo.isCancelled()) { - return; - } - - List sensitiveColumns = new ArrayList<>(); - for (DBTableColumn dbTableColumn : columns) { - String columnKey = getColumnKey(dbTableColumn); - ScanResult scanResult = scanResults.get(columnKey); - - if (scanResult != null) { - // 根据扫描模式获取最终的识别结果 - Optional finalResultOpt = scanResult - .getFinalResult(this.scanningMode); + try { + if (taskInfo.isCancelled()) { + return; + } + Map scanResults = this.scanner.scanBatch(columns, this.scanningMode); + if (taskInfo.isCancelled()) { + return; + } - // 如果最终有识别结果，则处理 - finalResultOpt.ifPresent(finalResult -> { - SensitiveColumnMeta meta = new SensitiveColumnMeta(database.getId(), objectName, - dbTableColumn.getName()); - // 使用同步块保证线程安全 - synchronized (existsSensitiveColumns) { - if (!existsSensitiveColumns.contains(meta)) { - SensitiveColumn column = createSensitiveColumn(columnType, objectName, - dbTableColumn, - finalResult); - sensitiveColumns.add(column); - existsSensitiveColumns.add(meta); + List sensitiveColumns = new ArrayList<>(); + for (DBTableColumn dbTableColumn : columns) { + String columnKey = getColumnKey(dbTableColumn); + ScanResult scanResult = scanResults.get(columnKey); + + if (scanResult != null) { + Optional finalResultOpt = scanResult + .getFinalResult(this.scanningMode); + finalResultOpt.ifPresent(finalResult -> { + SensitiveColumnMeta meta = new SensitiveColumnMeta(database.getId(), objectName, + dbTableColumn.getName()); + synchronized (existsSensitiveColumns) { + if (!existsSensitiveColumns.contains(meta)) { + SensitiveColumn column = createSensitiveColumn(columnType, objectName, + dbTableColumn, + finalResult); + sensitiveColumns.add(column); + existsSensitiveColumns.add(meta); + } } - } - }); + }); + } } + if (!sensitiveColumns.isEmpty()) { + taskInfo.addSensitiveColumns(sensitiveColumns); + } + taskInfo.addFinishedTableCount(); + } catch (Exception e) { + log.error("Failed to scan table {}: {}", objectName, e.getMessage(), e); + taskInfo.addFinishedTableCount(); } - // 批量添加敏感列结果，使用同步保证线程安全 - if (!sensitiveColumns.isEmpty()) { - taskInfo.addSensitiveColumns(sensitiveColumns); - } - taskInfo.addFinishedTableCount(); - } catch (Exception e) { - log.error("Failed to scan table {}: {}", objectName, e.getMessage(), e); - // 即使失败也要增加完成计数，避免任务卡住 - taskInfo.addFinishedTableCount(); - } - })) - .collect(Collectors.toList()); + })) + .collect(Collectors.toList()); - // 等待所有表的扫描任务完成 CompletableFuture.allOf(tableFutures.toArray(new CompletableFuture[0])).join(); } - // 【新增】辅助方法，用于创建 SensitiveColumn 对象，使代码更清晰 private SensitiveColumn createSensitiveColumn(SensitiveColumnType columnType, String objectName, - DBTableColumn dbTableColumn, RecognitionResult result) { + DBTableColumn dbTableColumn, RecognitionResult result) { SensitiveColumn column = new SensitiveColumn(); column.setType(columnType); column.setDatabase(database); column.setTableName(objectName); column.setColumnName(dbTableColumn.getName()); - // 从 RecognitionResult 获取 ruleId 和 level column.setSensitiveRuleId(result.getMatchedRuleId()); column.setLevel(result.getLevel()); - - // 设置脱敏算法ID Long maskingAlgorithmId = determineMaskingAlgorithmId(result); column.setMaskingAlgorithmId(maskingAlgorithmId); return column; } - /** - * 根据识别结果确定脱敏算法ID - * 对于AI识别的结果，如果是默认敏感类型则自动匹配同名脱敏算法，否则使用系统默认算法 - * 对于传统规则识别的结果，直接使用规则配置的脱敏算法 - */ private Long determineMaskingAlgorithmId(RecognitionResult result) { - // 通过 ruleId 从我们保存的 ruleMap 中找到对应的规则 SensitiveRule matchedRule = this.ruleMap.get(result.getMatchedRuleId()); if (matchedRule == null) { - // 如果找不到规则，使用系统默认算法 return getSystemDefaultAlgorithmId(); } - // 如果是AI规则识别的结果，需要根据敏感类型自动匹配算法 if (SensitiveRuleType.AI.equals(result.getSourceRuleType()) && result.getSensitiveType() != null) { return handleAiRecognitionResult(result.getSensitiveType()); } - // 对于传统规则，直接使用规则配置的脱敏算法ID return matchedRule.getMaskingAlgorithmId(); } - /** - * 处理AI识别结果的脱敏算法匹配 - */ private Long handleAiRecognitionResult(String sensitiveType) { - // 判断是否为默认敏感类型 if (DefaultSensitiveType.isDefaultType(sensitiveType)) { - // 通过DefaultSensitiveType获取算法名称，然后根据名称获取当前组织下的算法ID Optional algorithmNameOpt = DefaultSensitiveType.getAlgorithmNameBySensitiveType(sensitiveType); if (algorithmNameOpt.isPresent()) { try { MaskingAlgorithmService algorithmService = SpringContextUtil.getBean(MaskingAlgorithmService.class); Optional algorithmIdOpt = algorithmService.getAlgorithmIdByName(algorithmNameOpt.get(), - database.getOrganizationId()); + database.getOrganizationId()); if (algorithmIdOpt.isPresent()) { return algorithmIdOpt.get(); } @@ -242,19 +205,14 @@ private Long handleAiRecognitionResult(String sensitiveType) { } } - // 如果不是默认类型或获取失败，使用系统默认算法 return getSystemDefaultAlgorithmId(); } - /** - * 获取系统默认脱敏算法ID - */ private Long getSystemDefaultAlgorithmId() { try { MaskingAlgorithmService algorithmService = SpringContextUtil.getBean(MaskingAlgorithmService.class); return algorithmService.getDefaultAlgorithmIdByOrganizationId(database.getOrganizationId()); } catch (Exception e) { - // 记录错误日志，但不抛出异常，避免影响整个扫描流程 log.error("Failed to get default masking algorithm ID: {}", e.getMessage(), e); return null; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java index 63cca895ab..09102f78ca 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManager.java @@ -66,7 +66,7 @@ public class SensitiveColumnScanningTaskManager { private StatefulUuidStateIdGenerator statefulUuidStateIdGenerator; public SensitiveColumnScanningTaskInfo start(List databases, List rules, - ScanningModeType scanningMode, // 新增参数 + ScanningModeType scanningMode, ConnectionConfig connectionConfig, Map> databaseId2SensitiveColumns) { ConnectionSession session = new DefaultConnectSessionFactory(connectionConfig).generateSession(); try { diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java index 8130d77ef4..950fd57888 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnService.java @@ -410,7 +410,8 @@ public SensitiveColumnScanningTaskInfo startScanning(@NotNull Long projectId, PreConditions.notEmpty(rules, "sensitiveRules"); ConnectionConfig connectionConfig = databaseService.findDataSourceForConnectById(databases.get(0).getId()); Map> databaseId2SensitiveColumns = listExistSensitiveColumns(databaseIds); - return scanningTaskManager.start(databases, rules, req.getScanningMode(), connectionConfig, databaseId2SensitiveColumns); + return scanningTaskManager.start(databases, rules, req.getScanningMode(), connectionConfig, + databaseId2SensitiveColumns); } @Transactional(rollbackFor = Exception.class) @@ -429,8 +430,8 @@ public SensitiveColumnScanningTaskInfo getScanningResults(@NotNull Long projectI @Transactional(rollbackFor = Exception.class) @PreAuthenticate(hasAnyResourceRole = {"OWNER, DBA, SECURITY_ADMINISTRATOR"}, - actions = {"OWNER", "DBA", "SECURITY_ADMINISTRATOR"}, resourceType = "ODC_PROJECT", - indexOfIdParam = 0) + actions = {"OWNER", "DBA", "SECURITY_ADMINISTRATOR"}, resourceType = "ODC_PROJECT", + indexOfIdParam = 0) @StatefulRoute(stateName = StateName.UUID_STATEFUL_ID, stateIdExpression = "#taskId") public Boolean stopScanning(@NotNull Long projectId, @NotBlank String taskId) { SensitiveColumnScanningTaskInfo taskInfo = scanningTaskManager.get(taskId); @@ -553,88 +554,64 @@ private Map> getFilteringExistColumns(Long databaseI return filtered; } - /** - * 扫描单个表的敏感列 - */ @Transactional(rollbackFor = Exception.class) - @PreAuthenticate(hasAnyResourceRole = { "OWNER, DBA, SECURITY_ADMINISTRATOR" }, actions = { "OWNER", "DBA", - "SECURITY_ADMINISTRATOR" }, resourceType = "ODC_PROJECT", indexOfIdParam = 0) + @PreAuthenticate(hasAnyResourceRole = {"OWNER, DBA, SECURITY_ADMINISTRATOR"}, actions = {"OWNER", "DBA", + "SECURITY_ADMINISTRATOR"}, resourceType = "ODC_PROJECT", indexOfIdParam = 0) public String scanSingleTableAsync(@NotNull Long projectId, @NotNull @Valid SingleTableScanReq req) { - // 获取当前用户信息，用于在异步任务中设置认证上下文 final Long currentUserId = authenticationFacade.currentUserId(); final Long currentOrganizationId = authenticationFacade.currentOrganizationId(); final String currentUserAccountName = authenticationFacade.currentUserAccountName(); - - // 启动异步任务 String taskId = UUID.randomUUID().toString(); singleTableScanTaskManager.startTask(taskId, () -> { try { - // 在异步任务中设置用户认证上下文 com.oceanbase.odc.service.iam.util.SecurityContextUtils.setCurrentUser( - currentUserId, currentOrganizationId, currentUserAccountName); + currentUserId, currentOrganizationId, currentUserAccountName); List result = performSingleTableScan(projectId, req); singleTableScanTaskManager.setTaskResult(taskId, result); } catch (Exception e) { log.error("Single table scan failed for taskId: {}, projectId: {}, databaseId: {}, tableName: {}", - taskId, projectId, req.getDatabaseId(), req.getTableName(), e); + taskId, projectId, req.getDatabaseId(), req.getTableName(), e); String errorMessage = e.getMessage() != null ? e.getMessage() - : "扫描过程中发生未知错误: " + e.getClass().getSimpleName(); + : "An unknown error occurred during the scanning process: " + e.getClass().getSimpleName(); singleTableScanTaskManager.setTaskError(taskId, errorMessage); } }); return taskId; } - /** - * 获取单表扫描结果 - */ - @PreAuthenticate(hasAnyResourceRole = { "OWNER, DBA, SECURITY_ADMINISTRATOR" }, actions = { "OWNER", "DBA", - "SECURITY_ADMINISTRATOR" }, resourceType = "ODC_PROJECT", indexOfIdParam = 0) + @PreAuthenticate(hasAnyResourceRole = {"OWNER, DBA, SECURITY_ADMINISTRATOR"}, actions = {"OWNER", "DBA", + "SECURITY_ADMINISTRATOR"}, resourceType = "ODC_PROJECT", indexOfIdParam = 0) public SingleTableScanTaskManager.SingleTableScanTask getSingleTableScanResult(@NotNull Long projectId, - @NotBlank String taskId) { + @NotBlank String taskId) { return singleTableScanTaskManager.getTask(taskId); } - /** - * 执行单表扫描的具体逻辑 - */ private List performSingleTableScan(@NotNull Long projectId, - @NotNull @Valid SingleTableScanReq req) { - // 1. 获取数据库信息 + @NotNull @Valid SingleTableScanReq req) { Database database = databaseService.detail(req.getDatabaseId()); PreConditions.notNull(database, "database"); checkProjectDatabases(projectId, Collections.singletonList(req.getDatabaseId())); - - // 2. 获取连接配置 ConnectionConfig connectionConfig = connectionService - .getForConnectionSkipPermissionCheck(database.getDataSource().getId()); - - // 3. 获取表列信息 + .getForConnectionSkipPermissionCheck(database.getDataSource().getId()); List tableColumns = getTableColumns(connectionConfig, database.getName(), req.getTableName()); if (CollectionUtils.isEmpty(tableColumns)) { return Collections.emptyList(); } - - // 4. 获取扫描规则（使用预置的系统规则） List rules = getScanningRules(projectId, null); if (CollectionUtils.isEmpty(rules)) { return Collections.emptyList(); } - - // 5. 执行扫描 - 使用批量扫描以保持表级别上下文 ScanningStrategyFactory strategyFactory = new ScanningStrategyFactory(); SensitiveColumnScanner scanner = new SensitiveColumnScanner(rules, strategyFactory); - - // 使用批量扫描而不是逐个扫描，这样AI可以看到整个表的所有列 Map scanResults = scanner.scanBatch(tableColumns, req.getScanningMode()); List results = new ArrayList<>(); for (DBTableColumn column : tableColumns) { String columnKey = String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); ScanResult scanResult = scanResults.get(columnKey); if (scanResult != null) { @@ -650,15 +627,13 @@ private List performSingleTableScan(@NotNull Long projectId, sensitiveColumn.setEnabled(true); sensitiveColumn.setSensitiveRuleId(result.getMatchedRuleId()); sensitiveColumn.setLevel(result.getLevel()); - // 设置默认脱敏算法ID，可以从规则中获取 SensitiveRule matchedRule = rules.stream() - .filter(r -> r.getId().equals(result.getMatchedRuleId())) - .findFirst() - .orElse(null); + .filter(r -> r.getId().equals(result.getMatchedRuleId())) + .findFirst() + .orElse(null); if (matchedRule != null && matchedRule.getMaskingAlgorithmId() != null) { sensitiveColumn.setMaskingAlgorithmId(matchedRule.getMaskingAlgorithmId()); } else { - // 设置默认脱敏算法ID sensitiveColumn.setMaskingAlgorithmId(1L); } results.add(sensitiveColumn); @@ -669,11 +644,8 @@ private List performSingleTableScan(@NotNull Long projectId, return results; } - /** - * 获取指定表的列信息 - */ private List getTableColumns(ConnectionConfig connectionConfig, String databaseName, - String tableName) { + String tableName) { ConnectionSession session = new DefaultConnectSessionFactory(connectionConfig).generateSession(); try { DBSchemaAccessor accessor = DBSchemaAccessors.create(session); @@ -683,30 +655,23 @@ private List getTableColumns(ConnectionConfig connectionConfig, S } } - /** - * 获取单表扫描的预置规则 - */ private List getScanningRules(Long projectId, List sensitiveRuleIds) { - // 单表扫描使用预置的系统规则，不依赖用户配置的规则 SensitiveRule defaultRule = createDefaultScanningRule(); return Collections.singletonList(defaultRule); } - /** - * 创建单表扫描的默认规则 - */ private SensitiveRule createDefaultScanningRule() { SensitiveRule rule = new SensitiveRule(); - rule.setId(-1L); // 使用负数ID表示系统预置规则 + rule.setId(-1L); rule.setName("Single Table Scan Default Rule"); rule.setEnabled(true); - rule.setType(SensitiveRuleType.AI); // 使用AI类型 - rule.setAiSensitiveTypes(null); // 设置为null，对应提示词中的"No specified category" - rule.setAiCustomPrompt(null); // 设置为null，对应提示词中的"No supplementary rule" + rule.setType(SensitiveRuleType.AI); + rule.setAiSensitiveTypes(null); + rule.setAiCustomPrompt(null); Long organizationId = authenticationFacade.currentOrganizationId(); Long defaultAlgorithmId = algorithmService.getDefaultAlgorithmIdByOrganizationId(organizationId); rule.setMaskingAlgorithmId(defaultAlgorithmId); - rule.setLevel(SensitiveLevel.MEDIUM); // 设置默认敏感级别 + rule.setLevel(SensitiveLevel.MEDIUM); rule.setBuiltin(true); return rule; } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java index 4c1dc60f67..2bb5ef7cfa 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java @@ -16,8 +16,7 @@ import lombok.extern.slf4j.Slf4j; /** - * 单表扫描任务管理器 - * 提供轻量级的异步任务管理功能 + * Single-table scan task manager */ @Slf4j @Component @@ -29,14 +28,9 @@ public class SingleTableScanTaskManager { @Qualifier("scanSensitiveColumnExecutor") private ThreadPoolTaskExecutor executor; - /** - * 启动单表扫描任务 - */ public String startTask(String taskId, Runnable scanTask) { SingleTableScanTask task = new SingleTableScanTask(taskId); tasks.put(taskId, task); - - // 使用Spring的ThreadPoolTaskExecutor，它会自动传递Spring Security上下文 executor.submit(() -> { try { task.setStatus(TaskStatus.RUNNING); @@ -52,24 +46,15 @@ public String startTask(String taskId, Runnable scanTask) { return taskId; } - /** - * 启动单表扫描任务（自动生成taskId） - */ public String startTask(Runnable scanTask) { String taskId = UUID.randomUUID().toString(); return startTask(taskId, scanTask); } - /** - * 获取任务状态 - */ public SingleTableScanTask getTask(String taskId) { return tasks.get(taskId); } - /** - * 设置任务结果 - */ public void setTaskResult(String taskId, List result) { SingleTableScanTask task = tasks.get(taskId); if (task != null) { @@ -77,9 +62,6 @@ public void setTaskResult(String taskId, List result) { } } - /** - * 设置任务错误 - */ public void setTaskError(String taskId, String errorMessage) { SingleTableScanTask task = tasks.get(taskId); if (task != null) { @@ -88,28 +70,19 @@ public void setTaskError(String taskId, String errorMessage) { } } - /** - * 清理已完成的任务（可选的清理机制） - */ public void cleanupTask(String taskId) { tasks.remove(taskId); } - /** - * 任务状态枚举 - */ public enum TaskStatus { PENDING, RUNNING, COMPLETED, FAILED } - /** - * 单表扫描任务信息 - */ @Data public static class SingleTableScanTask { - private final String taskId; - private TaskStatus status = TaskStatus.PENDING; - private List result; - private String errorMessage; + private final String taskId; + private TaskStatus status = TaskStatus.PENDING; + private List result; + private String errorMessage; } } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java index b61ad5e640..a424901c18 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java @@ -34,7 +34,8 @@ import lombok.Data; /** - * AI功能配置类，支持从系统配置中动态读取配置 + * @author fenyf + * @date 2025/8/10 12:41 */ @Data @Component @@ -51,10 +52,6 @@ public class AIConfig { @Value("${odc.ai.model:gpt-3.5-turbo}") private String model; - // 硬编码超时和重试配置 - private static final int TIMEOUT_SECONDS = 30; - private static final int MAX_RETRIES = 3; - private Boolean enableThinking = AIParam.DEFAULT_ENABLE_THINKING; private Double temperature = AIParam.DEFAULT_TEMPERATURE; @@ -77,18 +74,16 @@ public Map loadAdditionalParams() { @ConditionalOnProperty(name = "odc.ai.enabled", havingValue = "true") public OpenAIClient openAIClient() { if (apiKey == null || apiKey.trim().isEmpty()) { - throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, new Object[]{"API key is not configured"}, "AI service is enabled but API key is not configured. Please set odc.ai.api-key configuration."); + throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, + new Object[] {"API key is not configured"}, + "AI service is enabled but API key is not configured. Please set odc.ai.api-key configuration."); } return OpenAIOkHttpClient.builder() - .apiKey(this.apiKey) - .baseUrl(this.baseUrl) - .build(); + .apiKey(this.apiKey) + .baseUrl(this.baseUrl) + .build(); } - /** - * 检查AI功能是否可用 - * @return true if AI功能启用且配置完整 - */ public boolean isAIAvailable() { return enabled && apiKey != null && !apiKey.trim().isEmpty(); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java index 02a1e0d6ed..04fe3ea82f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -25,66 +25,58 @@ import com.openai.models.chat.completions.ChatCompletion; import com.openai.models.chat.completions.ChatCompletionCreateParams; +/** + * @author fenyf + * @date 2025/8/10 12:41 + */ @Service public class AIInferenceService { private final AIConfig aiConfig; private final Optional openAIClient; - // 注入 AIConfig 和可选的 OpenAIClient Bean public AIInferenceService(AIConfig aiConfig, Optional openAIClient) { this.aiConfig = aiConfig; this.openAIClient = openAIClient; } - /** - * 检查AI功能是否可用 - * @throws IllegalStateException 如果AI功能不可用 - */ private void checkAIAvailability() { if (!aiConfig.isEnabled()) { - throw new BadRequestException(ErrorCodes.AIServiceNotAvailable, new Object[]{"AI service is not enabled"}, "AI service is not enabled. Please contact administrator to enable AI service."); + throw new BadRequestException(ErrorCodes.AIServiceNotAvailable, new Object[] {"AI service is not enabled"}, + "AI service is not enabled. Please contact administrator to enable AI service."); } if (!aiConfig.isAIAvailable()) { - throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, new Object[]{"AI configuration is incomplete"}, "AI configuration is incomplete. Please contact administrator to configure AI parameters."); + throw new BadRequestException(ErrorCodes.AIConfigurationIncomplete, + new Object[] {"AI configuration is incomplete"}, + "AI configuration is incomplete. Please contact administrator to configure AI parameters."); } if (!openAIClient.isPresent()) { - throw new BadRequestException(ErrorCodes.AIClientNotInitialized, new Object[]{"AI client is not initialized"}, "AI client is not initialized. Please check AI configuration and restart service."); + throw new BadRequestException(ErrorCodes.AIClientNotInitialized, + new Object[] {"AI client is not initialized"}, + "AI client is not initialized. Please check AI configuration and restart service."); } } - /** - * 使用系统提示词和用户提示词分别调用AI服务 - * - * @param systemPrompt 系统提示词 - * @param userPrompt 用户提示词 - * @return AI响应 - * @throws IllegalStateException 如果AI功能不可用 - */ public ChatCompletion chat(String systemPrompt, String userPrompt) { - // 检查AI功能可用性 checkAIAvailability(); try { ChatCompletionCreateParams params = ChatCompletionCreateParams.builder() - .addSystemMessage(systemPrompt) - .addUserMessage(userPrompt) - .model(aiConfig.getModel()) - .temperature(aiConfig.getTemperature()) - .topP(aiConfig.getTopP()) - .additionalBodyProperties(aiConfig.loadAdditionalParams()) - .build(); + .addSystemMessage(systemPrompt) + .addUserMessage(userPrompt) + .model(aiConfig.getModel()) + .temperature(aiConfig.getTemperature()) + .topP(aiConfig.getTopP()) + .additionalBodyProperties(aiConfig.loadAdditionalParams()) + .build(); return openAIClient.get().chat().completions().create(params); } catch (Exception e) { - throw new BadRequestException(ErrorCodes.AIInferenceServiceError, new Object[]{e.getMessage()}, "Failed to call AI inference service: " + e.getMessage(), e); + throw new BadRequestException(ErrorCodes.AIInferenceServiceError, new Object[] {e.getMessage()}, + "Failed to call AI inference service: " + e.getMessage(), e); } } - /** - * 检查AI功能是否可用（不抛出异常） - * @return true if AI功能可用 - */ public boolean isAIAvailable() { return aiConfig.isEnabled() && aiConfig.isAIAvailable() && openAIClient.isPresent(); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java index 3cd387fe7f..b5793e0664 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIParam.java @@ -15,6 +15,10 @@ */ package com.oceanbase.odc.service.datasecurity.ai; +/** + * @author fenyf + * @date 2025/8/10 12:41 + */ public class AIParam { /** @@ -27,4 +31,4 @@ public class AIParam { public static final Integer DEFAULT_MIN_P = 0; public static final Integer DEFAULT_BATCH_SIZE_IN_TABLE = 30; -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java index 1f59358382..266fefdd84 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java @@ -19,32 +19,18 @@ import lombok.Data; /** - * AI功能状态响应 - */ - @Data - public class AIStatusResponse { - /** - * AI功能是否启用 - */ - private boolean enabled; + * @author fenyf + * @date 2025/8/10 12:41 + */ +@Data +public class AIStatusResponse { + private boolean enabled; - /** - * AI功能是否可用（启用且配置完整） - */ - private boolean available; + private boolean available; - /** - * 使用的AI模型 - */ - private String model; + private String model; - /** - * API基础URL - */ - private String baseUrl; + private String baseUrl; - /** - * API密钥是否已配置 - */ - private boolean apiKeyConfigured; - } \ No newline at end of file + private boolean apiKeyConfigured; +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java index c0fae884da..6fa3d6d814 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java @@ -29,18 +29,14 @@ import lombok.var; /** - * AI 提示词 (Prompt) 模板加载器和构建器（重构版）。 - *

- * 该类负责加载结构化的 AI 提示词模板，并根据列元数据、指定的敏感类型和用户自定义提示来构建最终的提示词。 - *

+ * @author fenyf + * @date 2025/8/10 12:41 */ @Slf4j @Component public class PromptTemplateLoader { - - private static final String SYSTEM_TEMPLATE_PATH = "/ai-prompt-template/sensitive_column_recognize_system_prompt.txt"; - - // 定义占位符 + private static final String SYSTEM_TEMPLATE_PATH = + "/ai-prompt-template/sensitive_column_recognize_system_prompt.txt"; private static final String TYPES_PLACEHOLDER = "{sensitiveTypes}"; private static final String PROMPT_PLACEHOLDER = "{customPrompt}"; @@ -56,38 +52,26 @@ public void init() { this.systemTemplate = reader.lines().collect(Collectors.joining(System.lineSeparator())); } } catch (Exception e) { - // 在实际项目中，这里应该使用日志系统 - e.printStackTrace(); log.error("Failed to load AI system prompt template: {}", e.getMessage(), e); throw new IllegalStateException("Failed to load AI system prompt template", e); } } - /** - * 构建系统提示词 - * - * @param sensitiveTypes 用户指定的敏感类型列表 (例如 ["联系方式", "身份信息"]) - * @param customPrompt 用户为该规则自定义的补充说明提示 - * @return 填充了敏感类型和自定义提示的系统提示词字符串 - */ public String buildSystemPrompt(List sensitiveTypes, String customPrompt) { if (this.systemTemplate == null || this.systemTemplate.isEmpty()) { throw new IllegalStateException("System prompt template is not available. Check loading status."); } - // 1. 格式化敏感类型列表 String formattedTypes = (sensitiveTypes == null || sensitiveTypes.isEmpty()) - ? "No specified category." - : String.join(", ", sensitiveTypes); + ? "No specified category." + : String.join(", ", sensitiveTypes); - // 2. 格式化用户自定义提示 String formattedPrompt = (customPrompt == null || customPrompt.trim().isEmpty()) - ? "No supplementary rule." - : customPrompt; + ? "No supplementary rule." + : customPrompt; - // 3. 替换模板中的占位符 return this.systemTemplate - .replace(TYPES_PLACEHOLDER, formattedTypes) - .replace(PROMPT_PLACEHOLDER, formattedPrompt); + .replace(TYPES_PLACEHOLDER, formattedTypes) + .replace(PROMPT_PLACEHOLDER, formattedPrompt); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java index f2ab1db487..aab2f85985 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java @@ -31,11 +31,8 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** - * 扫描策略工厂类 - * 根据扫描模式类型返回对应的策略实例 - * - * @author Assistant - * @date 2025/1/27 + * @author fenyf + * @date 2025/8/10 12:41 */ @Component public class ScanningStrategyFactory { @@ -43,30 +40,19 @@ public class ScanningStrategyFactory { private final Map strategies = new HashMap<>(); public ScanningStrategyFactory() { - // 预创建所有策略实例 strategies.put(ScanningModeType.RULES_ONLY, new RulesOnlyStrategy()); strategies.put(ScanningModeType.JOINT_RECOGNITION, new JointRecognitionStrategy()); strategies.put(ScanningModeType.AI_ONLY, new AIOnlyStrategy()); } - /** - * 根据扫描模式获取对应的策略 - * - * @param mode 扫描模式 - * @return 对应的策略实例 - */ public ScanningStrategy getStrategy(ScanningModeType mode) { ScanningStrategy strategy = strategies.get(mode); if (strategy == null) { - // 返回默认的无操作策略 return new NoOpStrategy(); } return strategy; } - /** - * 无操作策略实现，用于处理未知或不支持的扫描模式 - */ private static class NoOpStrategy implements ScanningStrategy { @Override public ScanResult scan(DBTableColumn column, @@ -90,4 +76,4 @@ public Map scanBatch(java.util.List columns, return results; } } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java index f1b63279da..7b8733ebaf 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java @@ -15,82 +15,40 @@ */ package com.oceanbase.odc.service.datasecurity.model; -import java.util.Arrays; import java.util.Optional; /** - * 默认敏感类型枚举 - * 定义AI识别的13种默认敏感类型，每种类型对应一个同名的脱敏算法 - * 支持多语言匹配和模糊匹配，提高AI识别结果的容错率 - * - * @author AI Assistant - * @date 2024/01/01 + * 13 Default Sensitivity Types Identified by AI + * + * @author fenyf + * @date 2025/8/1 */ public enum DefaultSensitiveType { - /** - * 个人姓名（汉字类型） - */ PERSONAL_NAME_CHINESE("${com.oceanbase.odc.builtin-resource.masking-algorithm.personal-name-chinese.name}"), - /** - * 个人姓名（字母类型） - */ PERSONAL_NAME_ALPHABET("${com.oceanbase.odc.builtin-resource.masking-algorithm.personal-name-alphabet.name}"), - /** - * 昵称 - */ NICKNAME("${com.oceanbase.odc.builtin-resource.masking-algorithm.nickname.name}"), - /** - * 邮箱 - */ EMAIL("${com.oceanbase.odc.builtin-resource.masking-algorithm.email.name}"), - /** - * 地址 - */ ADDRESS("${com.oceanbase.odc.builtin-resource.masking-algorithm.address.name}"), - /** - * 手机号码 - */ PHONE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.phone-number.name}"), - /** - * 固定电话 - */ FIXED_LINE_PHONE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.fixed-line-phone-number.name}"), - /** - * 证件号码 - */ CERTIFICATE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.certificate-number.name}"), - /** - * 银行卡号 - */ BANK_CARD_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.bank-card-number.name}"), - /** - * 车牌号 - */ LICENSE_PLATE_NUMBER("${com.oceanbase.odc.builtin-resource.masking-algorithm.license-plate-number.name}"), - /** - * 设备唯一识别号 - */ DEVICE_ID("${com.oceanbase.odc.builtin-resource.masking-algorithm.device-id.name}"), - /** - * IP地址 - */ IP("${com.oceanbase.odc.builtin-resource.masking-algorithm.ip.name}"), - /** - * MAC地址 - */ MAC("${com.oceanbase.odc.builtin-resource.masking-algorithm.mac.name}"); private final String algorithmName; @@ -103,41 +61,23 @@ public String getAlgorithmName() { return algorithmName; } - /** - * 判断给定的敏感类型是否为默认类型（智能匹配） - * - * @param sensitiveType 敏感类型名称 - * @return 如果是默认类型返回true，否则返回false - */ public static boolean isDefaultType(String sensitiveType) { return findBestMatch(sensitiveType).isPresent(); } /** - * 根据敏感类型名称获取对应的脱敏算法名称（智能匹配） - * - * @param sensitiveType 敏感类型名称 - * @return 对应的脱敏算法名称，如果不是默认类型则返回空 + * Retrieve the corresponding algorithm name based on the name of the sensitive type. */ public static Optional getAlgorithmNameBySensitiveType(String sensitiveType) { return findBestMatch(sensitiveType).map(DefaultSensitiveType::getAlgorithmName); } - /** - * 根据敏感类型名称获取对应的枚举值（智能匹配） - * - * @param sensitiveType 敏感类型名称 - * @return 对应的枚举值，如果不是默认类型则返回空 - */ public static Optional getByDisplayName(String sensitiveType) { return findBestMatch(sensitiveType); } /** - * 智能匹配敏感类型（支持多种匹配策略） - * - * @param sensitiveType 敏感类型名称 - * @return 匹配到的枚举值 + * Match sensitive type */ public static Optional findBestMatch(String sensitiveType) { if (sensitiveType == null || sensitiveType.trim().isEmpty()) { @@ -146,14 +86,12 @@ public static Optional findBestMatch(String sensitiveType) String normalized = sensitiveType.toLowerCase().trim(); - // 精确匹配枚举名称（下划线格式） for (DefaultSensitiveType type : values()) { if (type.name().toLowerCase().equals(normalized)) { return Optional.of(type); } } - // 精确匹配连字符格式（AI返回的格式） for (DefaultSensitiveType type : values()) { String hyphenFormat = type.name().toLowerCase().replace("_", "-"); if (hyphenFormat.equals(normalized)) { @@ -161,12 +99,11 @@ public static Optional findBestMatch(String sensitiveType) } } - // 模糊匹配：检查是否包含关键词 for (DefaultSensitiveType type : values()) { String enumName = type.name().toLowerCase(); String hyphenFormat = enumName.replace("_", "-"); if (enumName.contains(normalized) || normalized.contains(enumName.replace("_", "")) || - hyphenFormat.contains(normalized) || normalized.contains(hyphenFormat.replace("-", ""))) { + hyphenFormat.contains(normalized) || normalized.contains(hyphenFormat.replace("-", ""))) { return Optional.of(type); } } @@ -174,4 +111,4 @@ public static Optional findBestMatch(String sensitiveType) return Optional.empty(); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java index f81ec26dda..d9064356ce 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java @@ -14,28 +14,23 @@ * limitations under the License. */ -package com.oceanbase.odc.service.datasecurity.model; // 建议放在 model 包下 +package com.oceanbase.odc.service.datasecurity.model; + -import java.util.ArrayList; -import java.util.List; import java.util.Optional; import lombok.AllArgsConstructor; import lombok.Getter; +/** + * @author fenyf + * @date 2025/7/18 17:52 + */ @Getter @AllArgsConstructor public class ScanResult { - // 基础规则的识别结果 (如果有) private final Optional basicRuleResult; - // AI 规则的识别结果 (如果有) private final Optional aiRuleResult; - /** - * 根据扫描模式获取最终的识别结果 - * - * @param scanningMode 扫描模式 - * @return 最终的识别结果 - */ public Optional getFinalResult(ScanningModeType scanningMode) { switch (scanningMode) { case RULES_ONLY: @@ -43,10 +38,9 @@ public Optional getFinalResult(ScanningModeType scanningMode) case AI_ONLY: return aiRuleResult; case JOINT_RECOGNITION: - // 对于联合识别，Scanner已经做过决策，直接返回存在的那个结果 return basicRuleResult.isPresent() ? basicRuleResult : aiRuleResult; default: return Optional.empty(); } } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java index f5a93092b7..09e638d4b1 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java @@ -15,7 +15,6 @@ */ package com.oceanbase.odc.service.datasecurity.model; -import java.util.List; import javax.validation.constraints.NotBlank; import javax.validation.constraints.NotNull; @@ -23,30 +22,21 @@ import lombok.Data; /** - * 单表敏感列扫描请求 - * - * @author Assistant - * @date 2025/1/27 + * Single-table sensitive column scan request + * + * @author fenyf + * @date 2025/8/18 17:52 */ @Data public class SingleTableScanReq { - /** - * 数据库ID - */ @NotNull private Long databaseId; - /** - * 表名 - */ @NotBlank private String tableName; - /** - * 扫描模式，默认为AI识别 - */ @NotNull private ScanningModeType scanningMode = ScanningModeType.JOINT_RECOGNITION; -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index be1cf87864..7864971943 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -45,16 +45,17 @@ import lombok.Data; /** - * AI Column Recognizer + * @author fenyf + * @date 2025/8/10 12:41 */ @Slf4j public class AIColumnRecognizer implements ColumnRecognizer { - private final SensitiveRule aiRule; // 直接保存整个规则对象 - private static final int BATCH_SIZE = AIParam.DEFAULT_BATCH_SIZE_IN_TABLE; // 单表内列数超过此值时进行分批处理 - private static final ObjectMapper objectMapper = new ObjectMapper(); // 用于解析JSON + private final SensitiveRule aiRule; + private static final int BATCH_SIZE = AIParam.DEFAULT_BATCH_SIZE_IN_TABLE; + private static final ObjectMapper objectMapper = new ObjectMapper(); private static final Pattern JSON_PATTERN = Pattern - .compile("(?s)```json\\s*([\\{\\[].*[\\}\\]])\\s*```|([\\{\\[].*[\\}\\]])"); + .compile("(?s)```json\\s*([\\{\\[].*[\\}\\]])\\s*```|([\\{\\[].*[\\}\\]])"); public AIColumnRecognizer(SensitiveRule rule) { this.aiRule = rule; @@ -62,44 +63,43 @@ public AIColumnRecognizer(SensitiveRule rule) { @Override public Optional recognize(DBTableColumn column) { - // 通过调用批量识别方法来实现单个识别 Map> batchResult = recognizeBatch(Collections.singletonList(column)); String columnKey = getColumnKey(column); return batchResult.getOrDefault(columnKey, Optional.empty()); } + /** + * If the data in the table columns is too large, scan them in batches. + * + * @param columns list of columns {@link DBTableColumn} + * @return + */ @Override public Map> recognizeBatch(List columns) { if (columns == null || columns.isEmpty()) { return Collections.emptyMap(); } - // 1. 获取依赖的服务 PromptTemplateLoader promptTemplateLoader = SpringContextUtil.getBean(PromptTemplateLoader.class); AIInferenceService aiService = SpringContextUtil.getBean(AIInferenceService.class); Map> finalAiResults = new HashMap<>(); - // 2. 如果列数超过批次大小，则分批处理；否则直接处理 if (columns.size() > BATCH_SIZE) { List> batches = Lists.partition(columns, BATCH_SIZE); try { - // 3. 遍历每一个小批次，分别调用 AI for (List batch : batches) { processBatch(batch, promptTemplateLoader, aiService, finalAiResults); } } catch (BadRequestException e) { - // 重新抛出BadRequestException，让调用方处理 throw e; } catch (Exception e) { log.error("Failed to process AI column recognition batch", e); return finalAiResults; } } else { - // 直接处理单批次 try { processBatch(columns, promptTemplateLoader, aiService, finalAiResults); } catch (BadRequestException e) { - // 重新抛出BadRequestException，让调用方处理 throw e; } catch (Exception e) { log.error("Failed to process AI column recognition", e); @@ -113,46 +113,34 @@ public Map> recognizeBatch(List batch, PromptTemplateLoader promptTemplateLoader, - AIInferenceService aiService, Map> finalAiResults) throws IOException { - // a. 构建系统提示词 + AIInferenceService aiService, Map> finalAiResults) throws IOException { String systemPrompt = promptTemplateLoader.buildSystemPrompt(aiRule.getAiSensitiveTypes(), - aiRule.getAiCustomPrompt()); - - // b. 构建用户提示词（列数据的JSON数组） + aiRule.getAiCustomPrompt()); String userPrompt = buildUserPrompt(batch); - - // c. 调用 AI ChatCompletion completion = aiService.chat(systemPrompt, userPrompt); String rawContent = completion.choices().get(0).message().content().orElse("[]"); - // c. 使用正则表达式从AI的返回结果中安全地提取JSON数组字符串 Matcher matcher = JSON_PATTERN.matcher(rawContent); String jsonArrayResponse = null; if (matcher.find()) { - // group(1) 对应被 ```json [...] ``` 包裹的内容, group(2) 对应裸露的 [...] - // 使用 Optional 来优雅地处理可能为null的捕获组 jsonArrayResponse = Optional.ofNullable(matcher.group(1)).orElse(matcher.group(2)); } if (jsonArrayResponse == null) { throw new BadRequestException(ErrorCodes.AIResponseFormatError, - new Object[] { "No valid JSON array found in AI response" }, - "AI response does not contain valid JSON format: " + rawContent); + new Object[] {"No valid JSON array found in AI response"}, + "AI response does not contain valid JSON format: " + rawContent); } - // d. 解析提取出的、更纯净的 JSON 数组 List batchResults; try { batchResults = objectMapper.readValue(jsonArrayResponse, - new TypeReference>() { - }); + new TypeReference>() {}); } catch (Exception e) { throw new BadRequestException(ErrorCodes.AIResponseFormatError, - new Object[] { "Failed to parse JSON: " + e.getMessage() }, - "Failed to parse AI response JSON: " + jsonArrayResponse, e); + new Object[] {"Failed to parse JSON: " + e.getMessage()}, + "Failed to parse AI response JSON: " + jsonArrayResponse, e); } - - // d. 将这批次的结果存入最终的 map，添加边界检查防止数组越界 int maxIndex = Math.min(batch.size(), batchResults.size()); for (int i = 0; i < maxIndex; i++) { DBTableColumn column = batch.get(i); @@ -161,22 +149,20 @@ private void processBatch(List batch, PromptTemplateLoader prompt if (dto.isSensitive()) { RecognitionResult result = RecognitionResult.builder() - .matched(true) - .matchedRuleId(this.aiRule.getId()) - .level(dto.getRiskLevel()) - .sourceRuleType(SensitiveRuleType.AI) - .sensitiveType(dto.getSensitiveCategory()) - .build(); + .matched(true) + .matchedRuleId(this.aiRule.getId()) + .level(dto.getRiskLevel()) + .sourceRuleType(SensitiveRuleType.AI) + .sensitiveType(dto.getSensitiveCategory()) + .build(); finalAiResults.put(columnKey, Optional.of(result)); } else { finalAiResults.put(columnKey, Optional.empty()); } } - - // Log warning if AI response count doesn't match input count if (batchResults.size() != batch.size()) { log.warn("AI response count ({}) does not match input column count ({})", - batchResults.size(), batch.size()); + batchResults.size(), batch.size()); } } @@ -184,7 +170,6 @@ private void processBatch(List batch, PromptTemplateLoader prompt * Build user prompt (JSON array of column data) */ private String buildUserPrompt(List batch) throws IOException { - // 将一批列的元数据转换为 JSON 数组字符串 List> columnMetadataList = batch.stream().map(c -> { Map meta = new HashMap<>(); meta.put("schemaName", c.getSchemaName()); @@ -204,4 +189,4 @@ private static class AiResponseDto { private SensitiveLevel riskLevel; private String sensitiveCategory; } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java index 381d08d119..fc3e64fb02 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/ColumnRecognizer.java @@ -25,8 +25,8 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** - * @author gaoda.xy - * @date 2023/5/30 10:32 + * @author fenyf + * @date 2025/8/10 12:41 */ public interface ColumnRecognizer { @@ -42,8 +42,7 @@ public interface ColumnRecognizer { * Batch recognizing the columns in database * * @param columns list of columns {@link DBTableColumn} - * @return map of recognizing results, key is column identifier, value is - * recognizing result + * @return map of recognizing results, key is column identifier, value is recognizing result */ default Map> recognizeBatch(List columns) { if (columns == null || columns.isEmpty()) { @@ -71,4 +70,4 @@ default String getColumnKey(DBTableColumn column) { column.getTableName() != null ? column.getTableName() : "unknown_table", column.getName() != null ? column.getName() : "unknown_column"); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java index e2b203abec..ceaa3f7e47 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizer.java @@ -41,7 +41,6 @@ public class GroovyColumnRecognizer implements ColumnRecognizer { private final Script script; private static final String COLUMN_KEYWORD = "column"; - // 【修改】构造函数接收 SensitiveRule 对象 public GroovyColumnRecognizer(SensitiveRule rule) { this.rule = rule; CompilerConfiguration config = new CompilerConfiguration(); @@ -58,7 +57,6 @@ public Optional recognize(DBTableColumn column) { binding.setVariable(COLUMN_KEYWORD, groovyColumnMeta); script.setBinding(binding); - // 【修改】获取脚本执行结果，并根据结果构建返回 boolean matched = (boolean) script.run(); if (matched) { RecognitionResult result = RecognitionResult.builder() diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java index 6fd489d518..430bb6f274 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizer.java @@ -39,7 +39,6 @@ public class PathColumnRecognizer implements ColumnRecognizer { private final List pathIncludeMatchers; private final List pathExcludeMatchers; - // 【修改】构造函数接收 SensitiveRule 对象 public PathColumnRecognizer(SensitiveRule rule) { this.rule = rule; pathIncludeMatchers = rule.getPathIncludes().stream().map(FieldPathMatcher::new).collect(Collectors.toList()); @@ -59,7 +58,6 @@ public Optional recognize(DBTableColumn column) { } for (FieldPathMatcher matcher : pathIncludeMatchers) { if (matcher.match(schemaName, tableName, columnName)) { - // 【修改】匹配成功，构建并返回 RecognitionResult RecognitionResult result = RecognitionResult.builder() .matched(true) .matchedRuleId(this.rule.getId()) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java index f3017e4cd0..117d4b04cf 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizer.java @@ -32,7 +32,6 @@ */ public class RegexColumnRecognizer implements ColumnRecognizer { - // 【修改】直接保存整个规则对象，以便获取 ID 和 Level private final SensitiveRule rule; private final Pattern databasePattern; private final Pattern tablePattern; @@ -41,13 +40,20 @@ public class RegexColumnRecognizer implements ColumnRecognizer { private static final long MATCH_TIMEOUT_MILLIS = 100L; - // 【修改】构造函数接收 SensitiveRule 对象 public RegexColumnRecognizer(SensitiveRule rule) { this.rule = rule; - databasePattern = StringUtils.isNotBlank(rule.getDatabaseRegexExpression()) ? Pattern.compile(rule.getDatabaseRegexExpression()) : null; - tablePattern = StringUtils.isNotBlank(rule.getTableRegexExpression()) ? Pattern.compile(rule.getTableRegexExpression()) : null; - columnPattern = StringUtils.isNotBlank(rule.getColumnRegexExpression()) ? Pattern.compile(rule.getColumnRegexExpression()) : null; - columnCommentPattern = StringUtils.isNotBlank(rule.getColumnCommentRegexExpression()) ? Pattern.compile(rule.getColumnCommentRegexExpression()) : null; + databasePattern = StringUtils.isNotBlank(rule.getDatabaseRegexExpression()) + ? Pattern.compile(rule.getDatabaseRegexExpression()) + : null; + tablePattern = + StringUtils.isNotBlank(rule.getTableRegexExpression()) ? Pattern.compile(rule.getTableRegexExpression()) + : null; + columnPattern = StringUtils.isNotBlank(rule.getColumnRegexExpression()) + ? Pattern.compile(rule.getColumnRegexExpression()) + : null; + columnCommentPattern = StringUtils.isNotBlank(rule.getColumnCommentRegexExpression()) + ? Pattern.compile(rule.getColumnCommentRegexExpression()) + : null; } @Override @@ -69,7 +75,6 @@ public Optional recognize(DBTableColumn column) { .matcher(new TimeoutCharSequence(column.getComment(), getTimeoutMillis())).matches()) { return Optional.empty(); } - // 【修改】如果所有条件都通过，说明匹配成功，构建并返回 RecognitionResult RecognitionResult result = RecognitionResult.builder() .matched(true) .matchedRuleId(this.rule.getId()) diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java index c69207be8a..5625d59f03 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java @@ -27,11 +27,8 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; /** - * 仅AI扫描策略实现 - * 只使用AI识别器进行识别，忽略基础规则 - * - * @author Assistant - * @date 2025/1/27 + * @author fenyf + * @date 2025/8/10 12:41 */ public class AIOnlyStrategy extends AbstractScanningStrategy { @@ -56,4 +53,4 @@ public Map scanBatch(List columns, List findFirstMatch(List recognizers, DBTableColumn column) { for (ColumnRecognizer recognizer : recognizers) { Optional result = recognizer.recognize(column); @@ -50,26 +40,16 @@ protected Optional findFirstMatch(List reco return Optional.empty(); } - /** - * 批量查找所有列的第一个匹配结果 - * - * @param recognizers 识别器列表 - * @param columns 待识别的列列表 - * @return 列标识符到识别结果的映射 - */ protected Map> findAllFirstMatches(List recognizers, List columns) { if (recognizers.isEmpty() || columns.isEmpty()) { return createEmptyResultMap(columns); } - - // 尝试使用批量识别（优先用于AI识别器） if (recognizers.size() == 1) { ColumnRecognizer recognizer = recognizers.get(0); return recognizer.recognizeBatch(columns); } - // 多个识别器时，逐个处理以保证优先级 Map> results = new HashMap<>(); for (DBTableColumn column : columns) { String columnKey = getColumnKey(column); @@ -79,12 +59,6 @@ protected Map> findAllFirstMatches(List> createEmptyResultMap(List columns) { Map> results = new HashMap<>(); for (DBTableColumn column : columns) { @@ -94,16 +68,10 @@ protected Map> createEmptyResultMap(List basicRecognizers, List aiRecognizers) { Optional basicResult = findFirstMatch(basicRecognizers, column); - - // 如果基础规则已经识别出来，就以此为准，不再调用AI if (basicResult.isPresent()) { return new ScanResult(basicResult, Optional.empty()); } - - // 否则，调用AI作为补充 Optional aiResult = findFirstMatch(aiRecognizers, column); return new ScanResult(Optional.empty(), aiResult); } @@ -57,7 +49,6 @@ public Map scanBatch(List columns, List aiRecognizers) { Map> basicResults = findAllFirstMatches(basicRecognizers, columns); - // 收集没有被基础规则识别的列 List remainingColumns = new ArrayList<>(); for (DBTableColumn column : columns) { String columnKey = getColumnKey(column); @@ -66,11 +57,8 @@ public Map scanBatch(List columns, List> aiResults = findAllFirstMatches(aiRecognizers, remainingColumns); - // 合并结果 Map results = new HashMap<>(); for (DBTableColumn column : columns) { String columnKey = getColumnKey(column); @@ -86,4 +74,4 @@ public Map scanBatch(List columns, List scanBatch(List columns, List basicRecognizers, List aiRecognizers); - /** - * 执行批量列的扫描 - * - * @param columns 待扫描的列列表 - * @param basicRecognizers 基础规则识别器列表 - * @param aiRecognizers AI识别器列表 - * @return 扫描结果映射，key为列标识符，value为扫描结果 - */ Map scanBatch(List columns, List basicRecognizers, List aiRecognizers); -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java index 257f88ebb4..b82d68aefe 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java @@ -225,4 +225,4 @@ public void test_gettersAndSetters_workCorrectly() { aiConfig.setMinP(20); Assert.assertEquals(Integer.valueOf(20), aiConfig.getMinP()); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java index b0a7810f77..e720f3f882 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java @@ -54,7 +54,6 @@ public class AIInferenceServiceTest { @Before public void setUp() { // Setup mock for openAIClient.chat().completions().create() call chain - // We'll mock this in individual test methods as needed } @Test @@ -69,7 +68,6 @@ public void test_chat_allConditionsMet_callsConfigMethods() { try { aiInferenceService.chat(systemPrompt, userPrompt); } catch (Exception e) { - // Expected - we can't mock the OpenAI client chain easily } // Then - Verify that config methods were called @@ -179,7 +177,8 @@ public void test_isAIAvailable_aiNotEnabled_returnsFalse() { // Given aiInferenceService = new AIInferenceService(aiConfig, Optional.of(openAIClient)); Mockito.when(aiConfig.isEnabled()).thenReturn(false); - // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit evaluation + // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit + // evaluation // When boolean result = aiInferenceService.isAIAvailable(); @@ -221,7 +220,8 @@ public void test_isAIAvailable_noConditionsMet_returnsFalse() { // Given aiInferenceService = new AIInferenceService(aiConfig, Optional.empty()); Mockito.when(aiConfig.isEnabled()).thenReturn(false); - // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit evaluation + // Note: aiConfig.isAIAvailable() stubbing removed as it's not called due to short-circuit + // evaluation // When boolean result = aiInferenceService.isAIAvailable(); @@ -256,4 +256,4 @@ private void setupValidAIConfig() { Mockito.when(aiConfig.getTopP()).thenReturn(0.9); Mockito.when(aiConfig.loadAdditionalParams()).thenReturn(new HashMap<>()); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java index 9af485d4e6..6752fd8fd9 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java @@ -51,9 +51,9 @@ public void test_buildSystemPrompt_withValidSensitiveTypesAndCustomPrompt_return // Then Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain formatted sensitive types", - result.contains("contact_info, identity_info")); + result.contains("contact_info, identity_info")); Assert.assertTrue("Should contain custom prompt", - result.contains("Additional rules for identification")); + result.contains("Additional rules for identification")); } @Test @@ -68,7 +68,7 @@ public void test_buildSystemPrompt_withEmptySensitiveTypes_returnsDefaultMessage // Then Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain default message for empty types", - result.contains("No specified category.")); + result.contains("No specified category.")); Assert.assertTrue("Should contain custom prompt", result.contains("Custom rules")); } @@ -84,7 +84,7 @@ public void test_buildSystemPrompt_withNullSensitiveTypes_returnsDefaultMessage( // Then Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain default message for null types", - result.contains("No specified category.")); + result.contains("No specified category.")); Assert.assertTrue("Should contain custom prompt", result.contains("Custom rules")); } @@ -100,9 +100,9 @@ public void test_buildSystemPrompt_withEmptyCustomPrompt_returnsDefaultMessage() // Then Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain formatted sensitive types", - result.contains("email, phone")); + result.contains("email, phone")); Assert.assertTrue("Should contain default message for empty prompt", - result.contains("No supplementary rule.")); + result.contains("No supplementary rule.")); } @Test @@ -118,7 +118,7 @@ public void test_buildSystemPrompt_withNullCustomPrompt_returnsDefaultMessage() Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain formatted sensitive types", result.contains("address")); Assert.assertTrue("Should contain default message for null prompt", - result.contains("No supplementary rule.")); + result.contains("No supplementary rule.")); } @Test @@ -134,7 +134,7 @@ public void test_buildSystemPrompt_withWhitespaceCustomPrompt_returnsDefaultMess Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain formatted sensitive types", result.contains("name")); Assert.assertTrue("Should contain default message for whitespace prompt", - result.contains("No supplementary rule.")); + result.contains("No supplementary rule.")); } @Test @@ -150,9 +150,9 @@ public void test_buildSystemPrompt_withSingleSensitiveType_returnsCorrectFormat( Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain single sensitive type", result.contains("credit_card")); Assert.assertFalse("Should not contain comma for single type", - result.contains("credit_card,")); + result.contains("credit_card,")); Assert.assertTrue("Should contain custom prompt", - result.contains("Strict validation required")); + result.contains("Strict validation required")); } @Test(expected = IllegalStateException.class) @@ -193,9 +193,9 @@ public void test_buildSystemPrompt_withMultipleSensitiveTypes_returnsCommaSepara // Then Assert.assertNotNull("Result should not be null", result); Assert.assertTrue("Should contain all types comma-separated", - result.contains("email, phone, address, name")); + result.contains("email, phone, address, name")); Assert.assertTrue("Should contain custom prompt", - result.contains("Multiple type validation")); + result.contains("Multiple type validation")); } @Test @@ -218,4 +218,4 @@ public void test_buildSystemPrompt_preservesOriginalTemplate_afterMultipleCalls( Assert.assertFalse("First result should not contain type2", result1.contains("type2")); Assert.assertFalse("Second result should not contain type1", result2.contains("type1")); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java index ea10732fc7..1a80b04fdd 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java @@ -119,21 +119,23 @@ public void test_getStrategy_returnsDifferentInstancesForDifferentModes() { @Test public void test_allStrategies_implementScanningStrategy() { // Given - ScanningModeType[] allModes = {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; + ScanningModeType[] allModes = + {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; // When & Then for (ScanningModeType mode : allModes) { ScanningStrategy strategy = factory.getStrategy(mode); Assert.assertNotNull("Strategy should not be null for mode: " + mode, strategy); Assert.assertTrue("Strategy should implement ScanningStrategy for mode: " + mode, - strategy instanceof ScanningStrategy); + strategy instanceof ScanningStrategy); } } @Test public void test_allStrategies_canHandleBasicOperations() { // Given - ScanningModeType[] allModes = {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; + ScanningModeType[] allModes = + {ScanningModeType.RULES_ONLY, ScanningModeType.AI_ONLY, ScanningModeType.JOINT_RECOGNITION}; List columns = Arrays.asList(testColumn); // When & Then @@ -148,7 +150,7 @@ public void test_allStrategies_canHandleBasicOperations() { Map batchResults = strategy.scanBatch(columns, emptyRecognizers, emptyRecognizers); Assert.assertNotNull("Batch scan results should not be null for mode: " + mode, batchResults); Assert.assertEquals("Batch scan should return result for each column for mode: " + mode, - 1, batchResults.size()); + 1, batchResults.size()); } } @@ -188,7 +190,7 @@ public void test_noOpStrategy_handlesNullColumnsGracefully() { // Then Assert.assertEquals("Should return results for all columns", 1, results.size()); Assert.assertTrue("Should contain key for unknown column", - results.containsKey("unknown_schema.unknown_table.unknown_column")); + results.containsKey("unknown_schema.unknown_table.unknown_column")); } private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { @@ -200,4 +202,4 @@ private DBTableColumn createTestColumn(String columnName, String typeName, Strin column.setTableName("test_table"); return column; } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java index 50775763c3..42cc6723ce 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java @@ -70,7 +70,8 @@ public void test_recognize_singleColumn_returnsSensitive() { // Given DBTableColumn column = createTestColumn("user_phone", "varchar", "user phone number"); String systemPrompt = "System prompt for AI"; - String aiResponse = "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]\n```"; + String aiResponse = + "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]\n```"; try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { setupMocks(mockedSpringContext, systemPrompt, aiResponse); @@ -109,16 +110,15 @@ public void test_recognize_singleColumn_returnsNotSensitive() { public void test_recognizeBatch_multipleColumns_returnsMixedResults() { // Given List columns = Arrays.asList( - createTestColumn("user_phone", "varchar", "user phone number"), - createTestColumn("id", "bigint", "primary key"), - createTestColumn("email", "varchar", "user email address") - ); + createTestColumn("user_phone", "varchar", "user phone number"), + createTestColumn("id", "bigint", "primary key"), + createTestColumn("email", "varchar", "user email address")); String systemPrompt = "System prompt for AI"; String aiResponse = "```json\n[" + - "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + - "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}," + - "{\"sensitive\": true, \"riskLevel\": \"MEDIUM\", \"sensitiveCategory\": \"contact_info\"}" + - "]```"; + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + + "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}," + + "{\"sensitive\": true, \"riskLevel\": \"MEDIUM\", \"sensitiveCategory\": \"contact_info\"}" + + "]```"; try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { setupMocks(mockedSpringContext, systemPrompt, aiResponse); @@ -171,14 +171,14 @@ public void test_recognizeBatch_aiServiceThrowsException_returnsEmptyResults() { try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { mockedSpringContext.when(() -> SpringContextUtil.getBean(PromptTemplateLoader.class)) - .thenReturn(promptTemplateLoader); + .thenReturn(promptTemplateLoader); mockedSpringContext.when(() -> SpringContextUtil.getBean(AIInferenceService.class)) - .thenReturn(aiInferenceService); + .thenReturn(aiInferenceService); Mockito.when(promptTemplateLoader.buildSystemPrompt(Mockito.anyList(), Mockito.anyString())) - .thenReturn(systemPrompt); + .thenReturn(systemPrompt); Mockito.when(aiInferenceService.chat(Mockito.anyString(), Mockito.anyString())) - .thenThrow(new RuntimeException("AI service error")); + .thenThrow(new RuntimeException("AI service error")); // When Map> results = recognizer.recognizeBatch(columns); @@ -211,7 +211,10 @@ public void test_recognizeBatch_malformedJsonResponse_throwsException() { // Given List columns = Arrays.asList(createTestColumn("test", "varchar", "test")); String systemPrompt = "System prompt for AI"; - String malformedResponse = "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\" missing_comma \"field\": \"value\"}]```"; // Invalid JSON syntax + String malformedResponse = + "```json\n[{\"sensitive\": true, \"riskLevel\": \"HIGH\" missing_comma \"field\": \"value\"}]```"; // Invalid + // JSON + // syntax try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { setupMocks(mockedSpringContext, systemPrompt, malformedResponse); @@ -229,7 +232,9 @@ public void test_recognizeBatch_responseWithoutJsonWrapper_parsesCorrectly() { // Given DBTableColumn column = createTestColumn("user_phone", "varchar", "user phone number"); String systemPrompt = "System prompt for AI"; - String aiResponse = "[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]"; // No ```json wrapper + String aiResponse = "[{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}]"; // No + // ```json + // wrapper try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { setupMocks(mockedSpringContext, systemPrompt, aiResponse); @@ -248,16 +253,15 @@ public void test_recognizeBatch_responseWithoutJsonWrapper_parsesCorrectly() { public void test_recognizeBatch_mismatchedResponseCount_handlesGracefully() { // Given List columns = Arrays.asList( - createTestColumn("col1", "varchar", "test1"), - createTestColumn("col2", "varchar", "test2"), - createTestColumn("col3", "varchar", "test3") - ); + createTestColumn("col1", "varchar", "test1"), + createTestColumn("col2", "varchar", "test2"), + createTestColumn("col3", "varchar", "test3")); String systemPrompt = "System prompt for AI"; // AI returns only 2 results for 3 columns String aiResponse = "```json\n[" + - "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + - "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}" + - "]```"; + "{\"sensitive\": true, \"riskLevel\": \"HIGH\", \"sensitiveCategory\": \"contact_info\"}," + + "{\"sensitive\": false, \"riskLevel\": null, \"sensitiveCategory\": null}" + + "]```"; try (MockedStatic mockedSpringContext = Mockito.mockStatic(SpringContextUtil.class)) { setupMocks(mockedSpringContext, systemPrompt, aiResponse); @@ -279,22 +283,23 @@ public void test_recognizeBatch_mismatchedResponseCount_handlesGracefully() { } // Helper methods - private void setupMocks(MockedStatic mockedSpringContext, String systemPrompt, String aiResponse) { + private void setupMocks(MockedStatic mockedSpringContext, String systemPrompt, + String aiResponse) { mockedSpringContext.when(() -> SpringContextUtil.getBean(PromptTemplateLoader.class)) - .thenReturn(promptTemplateLoader); + .thenReturn(promptTemplateLoader); mockedSpringContext.when(() -> SpringContextUtil.getBean(AIInferenceService.class)) - .thenReturn(aiInferenceService); + .thenReturn(aiInferenceService); Mockito.when(promptTemplateLoader.buildSystemPrompt(Mockito.anyList(), Mockito.anyString())) - .thenReturn(systemPrompt); + .thenReturn(systemPrompt); // Mock the chain: completion.choices().get(0).message().content().orElse("[]") // Use Mockito's deep stubbing with RETURNS_DEEP_STUBS ChatCompletion mockCompletion = Mockito.mock(ChatCompletion.class, Mockito.RETURNS_DEEP_STUBS); Mockito.when(aiInferenceService.chat(Mockito.anyString(), Mockito.anyString())) - .thenReturn(mockCompletion); + .thenReturn(mockCompletion); Mockito.when(mockCompletion.choices().get(0).message().content()) - .thenReturn(Optional.of(aiResponse)); + .thenReturn(Optional.of(aiResponse)); } private DBTableColumn createTestColumn(String columnName, String typeName, String comment) { @@ -321,4 +326,4 @@ private SensitiveRule createAIRule() { private String getColumnKey(DBTableColumn column) { return column.getSchemaName() + "." + column.getTableName() + "." + column.getName(); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java index 4eb1c2ebbb..ab575b7a5c 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java @@ -42,10 +42,11 @@ public void test_recognize_true() { ColumnRecognizer recognizer = new GroovyColumnRecognizer(rule); DBTableColumn dbTableColumn = createTestColumn(); Optional resultOpt = recognizer.recognize(dbTableColumn); - Assert.assertTrue("脚本匹配成功，应返回有值的 Optional", resultOpt.isPresent()); + Assert.assertTrue("Script matching is successful. An Optional with a value should be returned.", + resultOpt.isPresent()); RecognitionResult result = resultOpt.get(); - Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), result.getMatchedRuleId()); - Assert.assertEquals("规则类型应为 GROOVY", SensitiveRuleType.GROOVY, result.getSourceRuleType()); + Assert.assertEquals("The matching rule ID should be 1.", rule.getId(), result.getMatchedRuleId()); + Assert.assertEquals("The rule type should be GROOVY", SensitiveRuleType.GROOVY, result.getSourceRuleType()); } @Test @@ -55,7 +56,7 @@ public void test_recognize_false() { DBTableColumn dbTableColumn = createTestColumn(); dbTableColumn.setTableName("unmatched_table"); Optional resultOpt = recognizer.recognize(dbTableColumn); - Assert.assertFalse("脚本匹配失败，应返回空的 Optional", resultOpt.isPresent()); + Assert.assertFalse("Script matching failed. It should return an empty Optional.", resultOpt.isPresent()); } @Test @@ -65,7 +66,8 @@ public void test_recognize_nullColumnName() { DBTableColumn dbTableColumn = createTestColumn(); dbTableColumn.setName(null); Optional resultOpt = recognizer.recognize(dbTableColumn); - Assert.assertFalse("脚本执行异常，应返回空的 Optional", resultOpt.isPresent()); + Assert.assertFalse("The script execution has failed. It should return an empty Optional.", + resultOpt.isPresent()); } @Test @@ -81,8 +83,8 @@ public void test_securityInterceptor_forLoop() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("ForStatements are not allowed"); String script = "for (int i = 0; i < 1; i++) {\n" - + " i = 0;\n" - + "}"; + + " i = 0;\n" + + "}"; new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @@ -91,8 +93,8 @@ public void test_securityInterceptor_whileLoop() { thrown.expect(MultipleCompilationErrorsException.class); thrown.expectMessage("WhileStatements are not allowed"); String script = "while(true) {\n" - + " int i = 0;\n" - + "}"; + + " int i = 0;\n" + + "}"; new GroovyColumnRecognizer(createGroovyRule(1L, script)); } @@ -116,17 +118,17 @@ public void test_securityInterceptor_importPackage() { private String buildDefaultGroovyScript() { return "if (column.name.equals(\"column\")) {\n" - + " if (column.table.equalsIgnoreCase(\"iam_user\")) {\n" - + " if (column.schema.length() > 0) {\n" - + " if (column.comment.indexOf(\"user\") > 0) {\n" - + " if (column.type.toLowerCase().equals(\"varchar\")) {\n" - + " return true;\n" - + " }\n" - + " }\n" - + " }\n" - + " }\n" - + "}\n" - + "return false;"; + + " if (column.table.equalsIgnoreCase(\"iam_user\")) {\n" + + " if (column.schema.length() > 0) {\n" + + " if (column.comment.indexOf(\"user\") > 0) {\n" + + " if (column.type.toLowerCase().equals(\"varchar\")) {\n" + + " return true;\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + "}\n" + + "return false;"; } private DBTableColumn createTestColumn() { @@ -139,7 +141,6 @@ private DBTableColumn createTestColumn() { return dbTableColumn; } - // 【新增】辅助方法，用于快速创建测试用的 Groovy 规则 private SensitiveRule createGroovyRule(Long id, String script) { SensitiveRule rule = new SensitiveRule(); rule.setId(id); diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java index 877a6da112..da12ab3a42 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/PathColumnRecognizerTest.java @@ -36,14 +36,14 @@ public void test_recognize_true() { ColumnRecognizer recognizer = new PathColumnRecognizer(rule); Optional result1 = recognizer.recognize(createDBTableColumn("a", "b12", "c")); - Assert.assertTrue("路径 'a.b12.c' 应匹配成功", result1.isPresent()); + Assert.assertTrue("Path 'a.b12.c' should match successfully", result1.isPresent()); Assert.assertEquals(rule.getId(), result1.get().getMatchedRuleId()); Optional result2 = recognizer.recognize(createDBTableColumn("a12", "34b56", "c")); - Assert.assertTrue("路径 'a12.34b56.c' 应匹配成功", result2.isPresent()); + Assert.assertTrue("Path 'a12.34b56.c' should match successfully", result2.isPresent()); Optional result3 = recognizer.recognize(createDBTableColumn("a12", "34b", "c")); - Assert.assertTrue("路径 'a12.34b.c' 应匹配成功", result3.isPresent()); + Assert.assertTrue("Path 'a12.34b.c' should match successfully", result3.isPresent()); } @Test @@ -51,13 +51,13 @@ public void test_recognize_false() { SensitiveRule rule = createPathRule(1L, Arrays.asList("*.*b*.c"), Arrays.asList("a.b.*")); ColumnRecognizer recognizer = new PathColumnRecognizer(rule); - Assert.assertFalse("路径 'a.b.c' 应被排除，匹配失败", + Assert.assertFalse("The path 'a.b.c' should be excluded as the match failed.", recognizer.recognize(createDBTableColumn("a", "b", "c")).isPresent()); - Assert.assertFalse("路径 'a12.b34.c56' 不应匹配，匹配失败", + Assert.assertFalse("Path 'a12.b34.c56' should not match; the match failed.", recognizer.recognize(createDBTableColumn("a12", "b34", "c56")).isPresent()); - Assert.assertFalse("路径 'a12.b34.null' 不应匹配，匹配失败", + Assert.assertFalse("Path 'a12.b34.null' should not match; the match failed.", recognizer.recognize(createDBTableColumn("a12", "b34", null)).isPresent()); } @@ -69,6 +69,7 @@ private DBTableColumn createDBTableColumn(String schemaName, String tableName, S column.setName(columnName); return column; } + private SensitiveRule createPathRule(Long id, List includes, List excludes) { SensitiveRule rule = new SensitiveRule(); rule.setId(id); diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java index 3a2cdad4d2..92965b4f7c 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java @@ -37,8 +37,8 @@ public void recognize_returnTrue() { Optional resultOpt = recognizer.recognize(column); - Assert.assertTrue("正则表达式应匹配成功", resultOpt.isPresent()); - Assert.assertEquals("匹配的规则ID应为 1", rule.getId(), resultOpt.get().getMatchedRuleId()); + Assert.assertTrue("The regular expression should match successfully.", resultOpt.isPresent()); + Assert.assertEquals("The matching rule ID should be 1.", rule.getId(), resultOpt.get().getMatchedRuleId()); } @Test @@ -46,13 +46,13 @@ public void recognize_returnFalse() { SensitiveRule rule = createTestRegexRule(1L); ColumnRecognizer recognizer = new RegexColumnRecognizer(rule); - Assert.assertFalse("Comment 为 null 时不应匹配", + Assert.assertFalse("When Comment is null, it should not match.", recognizer.recognize(createDBTableColumn("xxx", "xxx", "user_email", null)).isPresent()); - Assert.assertFalse("SchemaName 不匹配时应失败", + Assert.assertFalse("It should fail when the SchemaName does not match.", recognizer.recognize(createDBTableColumn(" ", "xxx", "user_email", "email of user")).isPresent()); - Assert.assertFalse("ColumnName 和 Comment 都不匹配时应失败", + Assert.assertFalse("When both ColumnName and Comment do not match, it should fail.", recognizer.recognize(createDBTableColumn("xxx", "xxx", "user", "some info")).isPresent()); } diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java index 3a27836924..a64f49ad69 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java @@ -182,9 +182,8 @@ public void test_scanBatch_withSingleAiRecognizer_usesBatchRecognition() { // Given List columns = Arrays.asList(testColumn); Map> batchResults = Collections.singletonMap( - getColumnKey(testColumn), - Optional.of(createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.AI)) - ); + getColumnKey(testColumn), + Optional.of(createRecognitionResult(1L, SensitiveLevel.HIGH, SensitiveRuleType.AI))); Mockito.when(mockAiRecognizer.recognizeBatch(columns)).thenReturn(batchResults); @@ -213,17 +212,17 @@ private DBTableColumn createTestColumn(String columnName, String typeName, Strin private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { return RecognitionResult.builder() - .matched(true) - .matchedRuleId(ruleId) - .level(level) - .sourceRuleType(ruleType) - .build(); + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); } private String getColumnKey(DBTableColumn column) { return String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java index 69e9fc3062..19ea6a5591 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java @@ -120,7 +120,8 @@ public void test_findAllFirstMatches_withMultipleColumns_returnsCorrectMapping() // Then Assert.assertEquals("Should return results for all columns", 2, results.size()); Assert.assertTrue("Should find result for column1", results.get(strategy.getColumnKey(column1)).isPresent()); - Assert.assertFalse("Should not find result for column2", results.get(strategy.getColumnKey(column2)).isPresent()); + Assert.assertFalse("Should not find result for column2", + results.get(strategy.getColumnKey(column2)).isPresent()); } @Test @@ -172,7 +173,8 @@ public void test_findAllFirstMatches_withMultipleRecognizers_usesIndividualRecog // Then Assert.assertEquals("Should return results for all columns", 2, results.size()); Assert.assertTrue("Should find result for column1", results.get(strategy.getColumnKey(column1)).isPresent()); - Assert.assertFalse("Should not find result for column2", results.get(strategy.getColumnKey(column2)).isPresent()); + Assert.assertFalse("Should not find result for column2", + results.get(strategy.getColumnKey(column2)).isPresent()); // Verify that recognizeBatch was not called for multiple recognizers Mockito.verify(mockRecognizer1, Mockito.never()).recognizeBatch(Mockito.any()); @@ -205,28 +207,28 @@ private DBTableColumn createTestColumn(String columnName, String typeName, Strin private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level) { return RecognitionResult.builder() - .matched(true) - .matchedRuleId(ruleId) - .level(level) - .sourceRuleType(SensitiveRuleType.REGEX) - .build(); + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(SensitiveRuleType.REGEX) + .build(); } // Testable implementation of AbstractScanningStrategy for testing protected methods private static class TestableAbstractScanningStrategy extends AbstractScanningStrategy { @Override public com.oceanbase.odc.service.datasecurity.model.ScanResult scan( - DBTableColumn column, - List basicRecognizers, - List aiRecognizers) { + DBTableColumn column, + List basicRecognizers, + List aiRecognizers) { return null; // Not used in these tests } @Override public Map scanBatch( - List columns, - List basicRecognizers, - List aiRecognizers) { + List columns, + List basicRecognizers, + List aiRecognizers) { return null; // Not used in these tests } @@ -237,7 +239,8 @@ public Optional findFirstMatch(List recogni } @Override - public Map> findAllFirstMatches(List recognizers, List columns) { + public Map> findAllFirstMatches(List recognizers, + List columns) { return super.findAllFirstMatches(recognizers, columns); } @@ -246,4 +249,4 @@ public String getColumnKey(DBTableColumn column) { return super.getColumnKey(column); } } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java index 1e1f492937..c6f08b72fb 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java @@ -279,17 +279,17 @@ private DBTableColumn createTestColumn(String columnName, String typeName, Strin private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { return RecognitionResult.builder() - .matched(true) - .matchedRuleId(ruleId) - .level(level) - .sourceRuleType(ruleType) - .build(); + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); } private String getColumnKey(DBTableColumn column) { return String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } -} \ No newline at end of file +} diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java index e9a20e2120..fe731415d8 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java @@ -190,17 +190,17 @@ private DBTableColumn createTestColumn(String columnName, String typeName, Strin private RecognitionResult createRecognitionResult(Long ruleId, SensitiveLevel level, SensitiveRuleType ruleType) { return RecognitionResult.builder() - .matched(true) - .matchedRuleId(ruleId) - .level(level) - .sourceRuleType(ruleType) - .build(); + .matched(true) + .matchedRuleId(ruleId) + .level(level) + .sourceRuleType(ruleType) + .build(); } private String getColumnKey(DBTableColumn column) { return String.format("%s.%s.%s", - column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", - column.getTableName() != null ? column.getTableName() : "unknown_table", - column.getName() != null ? column.getName() : "unknown_column"); + column.getSchemaName() != null ? column.getSchemaName() : "unknown_schema", + column.getTableName() != null ? column.getTableName() : "unknown_table", + column.getName() != null ? column.getName() : "unknown_column"); } -} \ No newline at end of file +} From b6c6ce8843b76f773a019636cacbd74f32e0389c Mon Sep 17 00:00:00 2001 From: fenyf Date: Thu, 18 Sep 2025 20:59:24 +0800 Subject: [PATCH 08/10] =?UTF-8?q?fix(ai=5Frecognition):=E4=BF=AE=E5=A4=8D?= =?UTF-8?q?=E9=9B=86=E6=88=90=E6=B5=8B=E8=AF=95bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../SensitiveColumnScanningTaskManagerTest.java | 4 ++-- .../src/main/resources/i18n/ErrorMessages.properties | 8 +++++++- .../main/resources/i18n/ErrorMessages_zh_CN.properties | 9 ++++++++- .../datasecurity/recognizer/AIColumnRecognizer.java | 6 +++--- 4 files changed, 20 insertions(+), 7 deletions(-) diff --git a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java index a75b8325e6..3f4d414025 100644 --- a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java +++ b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java @@ -108,7 +108,7 @@ public static void tearDown() { public void test_start_groovyRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createGroovySensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -128,7 +128,7 @@ public void test_start_groovyRule_OBOracle() { public void test_start_pathRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createPathSensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); diff --git a/server/odc-core/src/main/resources/i18n/ErrorMessages.properties b/server/odc-core/src/main/resources/i18n/ErrorMessages.properties index 89ccb301f4..808054c177 100644 --- a/server/odc-core/src/main/resources/i18n/ErrorMessages.properties +++ b/server/odc-core/src/main/resources/i18n/ErrorMessages.properties @@ -212,4 +212,10 @@ com.oceanbase.odc.ErrorCodes.UpdateNotAllowed=Editing is not allowed in the curr com.oceanbase.odc.ErrorCodes.PauseNotAllowed=Disabling is not allowed in the current state, please check if there are any records in execution. com.oceanbase.odc.ErrorCodes.DeleteNotAllowed=Deletion is not allowed in the current state. com.oceanbase.odc.ErrorCodes.ExtractFileFailed=Failed to extract the file. Please check whether the file is correct. Details: {0} -com.oceanbase.odc.ErrorCodes.InvalidSignature=File verification failed. Do not modify exported files or check if the correct key was used. \ No newline at end of file +com.oceanbase.odc.ErrorCodes.InvalidSignature=File verification failed. Do not modify exported files or check if the correct key was used. +com.oceanbase.odc.ErrorCodes.AIServiceNotAvailable=AI service is not available. Please contact administrator to enable AI service. +com.oceanbase.odc.ErrorCodes.AIConfigurationIncomplete=AI configuration is incomplete. Please contact administrator to configure AI parameters. +com.oceanbase.odc.ErrorCodes.AIClientNotInitialized=AI client is not initialized. Please check AI configuration and restart service. +com.oceanbase.odc.ErrorCodes.AIInferenceServiceError=Failed to call AI inference service. Details: {0} +com.oceanbase.odc.ErrorCodes.AIResponseFormatError=AI response format is invalid. Details: {0} +com.oceanbase.odc.ErrorCodes.AIResponseCountMismatch=AI response count does not match expected count. Expected: {0}, Actual: {1} \ No newline at end of file diff --git a/server/odc-core/src/main/resources/i18n/ErrorMessages_zh_CN.properties b/server/odc-core/src/main/resources/i18n/ErrorMessages_zh_CN.properties index 1a84987660..7bb054aa71 100644 --- a/server/odc-core/src/main/resources/i18n/ErrorMessages_zh_CN.properties +++ b/server/odc-core/src/main/resources/i18n/ErrorMessages_zh_CN.properties @@ -214,4 +214,11 @@ com.oceanbase.odc.ErrorCodes.UnsupportedSyncTableStructure=结构同步暂不支 com.oceanbase.odc.ErrorCodes.ScheduleIntervalTooShort=执行间隔配置过短，请重新配置。最小间隔为：{0} 秒 com.oceanbase.odc.ErrorCodes.ExtractFileFailed=提取文件失败，请确认文件是否正确，错误详情 {0} -com.oceanbase.odc.ErrorCodes.InvalidSignature=文件验签不通过，请勿修改导出文件或检查密钥是否正确 \ No newline at end of file +com.oceanbase.odc.ErrorCodes.InvalidSignature=文件验签不通过，请勿修改导出文件或检查密钥是否正确 + +com.oceanbase.odc.ErrorCodes.AIServiceNotAvailable=AI 服务不可用，请联系管理员启用 AI 服务 +com.oceanbase.odc.ErrorCodes.AIConfigurationIncomplete=AI 配置不完整，请联系管理员配置 AI 参数 +com.oceanbase.odc.ErrorCodes.AIClientNotInitialized=AI 客户端未初始化，请检查 AI 配置并重启服务 +com.oceanbase.odc.ErrorCodes.AIInferenceServiceError=调用 AI 推理服务失败，错误详情：{0} +com.oceanbase.odc.ErrorCodes.AIResponseFormatError=AI 响应格式无效，错误详情：{0} +com.oceanbase.odc.ErrorCodes.AIResponseCountMismatch=AI 响应数量与预期不符，预期：{0}，实际：{1} \ No newline at end of file diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java index 7864971943..9870d29d00 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizer.java @@ -28,11 +28,9 @@ import com.fasterxml.jackson.core.type.TypeReference; import com.fasterxml.jackson.databind.ObjectMapper; import com.google.common.collect.Lists; -import com.oceanbase.odc.service.common.util.SpringContextUtil; - -import lombok.extern.slf4j.Slf4j; import com.oceanbase.odc.core.shared.constant.ErrorCodes; import com.oceanbase.odc.core.shared.exception.BadRequestException; +import com.oceanbase.odc.service.common.util.SpringContextUtil; import com.oceanbase.odc.service.datasecurity.ai.AIInferenceService; import com.oceanbase.odc.service.datasecurity.ai.AIParam; import com.oceanbase.odc.service.datasecurity.ai.PromptTemplateLoader; @@ -42,7 +40,9 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import com.openai.models.chat.completions.ChatCompletion; + import lombok.Data; +import lombok.extern.slf4j.Slf4j; /** * @author fenyf From bee19d47f7d825b11da4a48e122f525b26c782b3 Mon Sep 17 00:00:00 2001 From: fenyf Date: Fri, 19 Sep 2025 22:43:50 +0800 Subject: [PATCH 09/10] =?UTF-8?q?fix(ai=5Frecognition):=E4=BB=A3=E7=A0=81?= =?UTF-8?q?=E6=A0=BC=E5=BC=8F=E5=8C=96?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit fix(ai_recognition):代码格式化和集成测试 --- .gitattributes | 1 + ...ensitiveColumnScanningTaskManagerTest.java | 20 ++++++++++++------- .../web/controller/v2/AIController.java | 5 ++--- .../v2/SensitiveColumnController.java | 10 ++++++---- .../supervisor/SupervisorApplicationTest.java | 8 -------- .../datasecurity/MaskingAlgorithmService.java | 4 ++-- .../datasecurity/SensitiveColumnScanner.java | 3 +-- .../SensitiveColumnScanningTask.java | 4 ++-- .../SingleTableScanTaskManager.java | 15 ++++++++++++++ .../odc/service/datasecurity/ai/AIConfig.java | 2 +- .../datasecurity/ai/AIInferenceService.java | 2 +- .../datasecurity/ai/AIStatusResponse.java | 3 +-- .../datasecurity/ai/PromptTemplateLoader.java | 6 ++---- .../factory/ColumnRecognizerFactory.java | 2 +- .../factory/ScanningStrategyFactory.java | 3 +-- .../model/DefaultSensitiveType.java | 2 +- .../datasecurity/model/RecognitionResult.java | 6 ++---- .../datasecurity/model/ScanResult.java | 5 ++--- .../SensitiveColumnScanningTaskInfo.java | 2 +- .../model/SingleTableScanReq.java | 3 +-- .../datasecurity/strategy/AIOnlyStrategy.java | 3 +-- .../strategy/AbstractScanningStrategy.java | 3 +-- .../strategy/JointRecognitionStrategy.java | 3 +-- .../strategy/RulesOnlyStrategy.java | 3 +-- .../strategy/ScanningStrategy.java | 3 +-- .../service/datasecurity/ai/AIConfigTest.java | 2 +- .../ai/AIInferenceServiceTest.java | 2 +- .../ai/PromptTemplateLoaderTest.java | 2 +- .../factory/ScanningStrategyFactoryTest.java | 2 +- .../recognizer/AIColumnRecognizerTest.java | 3 +-- .../GroovyColumnRecognizerTest.java | 1 - .../recognizer/RegexColumnRecognizerTest.java | 1 - .../strategy/AIOnlyStrategyTest.java | 2 +- .../AbstractScanningStrategyTest.java | 2 +- .../JointRecognitionStrategyTest.java | 2 +- .../strategy/RulesOnlyStrategyTest.java | 3 +-- 36 files changed, 70 insertions(+), 73 deletions(-) create mode 100644 .gitattributes diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000000..94f480de94 --- /dev/null +++ b/.gitattributes @@ -0,0 +1 @@ +* text=auto eol=lf \ No newline at end of file diff --git a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java index 3f4d414025..0a30b3cba7 100644 --- a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java +++ b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java @@ -108,17 +108,19 @@ public static void tearDown() { public void test_start_groovyRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createGroovySensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); - Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); + Assert.assertEquals(3, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); } @Test public void test_start_groovyRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createGroovySensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -128,7 +130,8 @@ public void test_start_groovyRule_OBOracle() { public void test_start_pathRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createPathSensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -138,7 +141,8 @@ public void test_start_pathRule_OBMySQL() { public void test_start_pathRule_OBMOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createPathSensitiveRules()); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -148,7 +152,8 @@ public void test_start_pathRule_OBMOracle() { public void test_start_RegexRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_MYSQL)); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -158,7 +163,8 @@ public void test_start_RegexRule_OBMySQL() { public void test_start_RegexRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_ORACLE)); - SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + SensitiveColumnScanningTaskInfo taskInfo = + manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java index f8d702ae6f..57a2dd3639 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.server.web.controller.v2; import org.springframework.beans.factory.annotation.Autowired; @@ -48,7 +47,7 @@ public class AIController { @ApiOperation(value = "Query the status of the AI function", - notes = "Return the status of whether the AI function is enabled and its configuration status") + notes = "Return the status of whether the AI function is enabled and its configuration status") @SkipAuthorize("AI status is safe to query for authenticated users") @GetMapping("/status") public SuccessResponse getAIStatus() { diff --git a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java index d48cbc4654..dd7370ba16 100644 --- a/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java +++ b/server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/SensitiveColumnController.java @@ -148,20 +148,22 @@ public SuccessResponse getScanningResults(@Path @ApiOperation(value = "stopScanning", notes = "Stop a sensitive column scanning task") @RequestMapping(value = "/stopScanning", method = RequestMethod.POST) public SuccessResponse stopScanning(@PathVariable Long projectId, - @RequestParam String taskId) { + @RequestParam String taskId) { return Responses.success(service.stopScanning(projectId, taskId)); } @ApiOperation(value = "getSingleTableScanResult", notes = "Get single table scan result") @RequestMapping(value = "/singleTableScan/{taskId}/result", method = RequestMethod.GET) - public SuccessResponse getSingleTableScanResult(@PathVariable Long projectId, - @PathVariable String taskId) { + public SuccessResponse getSingleTableScanResult( + @PathVariable Long projectId, + @PathVariable String taskId) { return Responses.success(service.getSingleTableScanResult(projectId, taskId)); } + @ApiOperation(value = "scanSingleTableAsync", notes = "Start an asynchronous single table scan") @RequestMapping(value = "/scanSingleTableAsync", method = RequestMethod.POST) public SuccessResponse scanSingleTableAsync(@PathVariable Long projectId, - @RequestBody SingleTableScanReq req) { + @RequestBody SingleTableScanReq req) { return Responses.success(service.scanSingleTableAsync(projectId, req)); } } diff --git a/server/odc-server/src/test/java/com/oceanbase/odc/supervisor/SupervisorApplicationTest.java b/server/odc-server/src/test/java/com/oceanbase/odc/supervisor/SupervisorApplicationTest.java index 5fb40d6311..6946197185 100644 --- a/server/odc-server/src/test/java/com/oceanbase/odc/supervisor/SupervisorApplicationTest.java +++ b/server/odc-server/src/test/java/com/oceanbase/odc/supervisor/SupervisorApplicationTest.java @@ -15,14 +15,6 @@ */ package com.oceanbase.odc.supervisor; -/** - * @author longpeng.zlp - * @date 2024/12/9 15:59 - */ -/** - * @author longpeng.zlp - * @date 2024/12/9 15:59 - */ import java.io.IOException; import java.net.InetSocketAddress; import java.net.Socket; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java index 479ed36dca..744ed3e3e3 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/MaskingAlgorithmService.java @@ -210,14 +210,14 @@ public Long getDefaultAlgorithmIdByOrganizationId(@NonNull Long organizationId) @SkipAuthorize("odc internal usages") public Optional getAlgorithmIdByName(@NonNull String algorithmName, @NonNull Long organizationId) { List entities = - algorithmRepository.findByNameAndOrganizationId(algorithmName, organizationId); + algorithmRepository.findByNameAndOrganizationId(algorithmName, organizationId); if (entities.isEmpty()) { log.warn("No masking algorithm found with name: {} for organization: {}", algorithmName, organizationId); return Optional.empty(); } if (entities.size() > 1) { log.warn("Multiple masking algorithms found with name: {} for organization: {}, using the first one", - algorithmName, organizationId); + algorithmName, organizationId); } return Optional.of(entities.get(0).getId()); } diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java index 4788d4c7cc..236c3550f8 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanner.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity; import java.util.List; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java index 25b1551c2d..2dd234a21c 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTask.java @@ -28,8 +28,10 @@ import java.util.stream.Collectors; import com.oceanbase.odc.core.shared.constant.ErrorCodes; +import com.oceanbase.odc.service.common.util.SpringContextUtil; import com.oceanbase.odc.service.connection.database.model.Database; import com.oceanbase.odc.service.datasecurity.factory.ScanningStrategyFactory; +import com.oceanbase.odc.service.datasecurity.model.DefaultSensitiveType; import com.oceanbase.odc.service.datasecurity.model.RecognitionResult; import com.oceanbase.odc.service.datasecurity.model.ScanResult; import com.oceanbase.odc.service.datasecurity.model.ScanningModeType; @@ -40,9 +42,7 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveColumnType; import com.oceanbase.odc.service.datasecurity.model.SensitiveRule; import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; -import com.oceanbase.odc.service.datasecurity.model.DefaultSensitiveType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; -import com.oceanbase.odc.service.common.util.SpringContextUtil; import lombok.extern.slf4j.Slf4j; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java index 2bb5ef7cfa..2472426bc6 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/SingleTableScanTaskManager.java @@ -1,3 +1,18 @@ +/* + * Copyright (c) 2023 OceanBase. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package com.oceanbase.odc.service.datasecurity; import java.util.List; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java index a424901c18..a8a81a27b4 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIConfig.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java index 04fe3ea82f..a1bf62e906 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceService.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java index 266fefdd84..d5ed6d977b 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/AIStatusResponse.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.ai; import lombok.Data; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java index 6fa3d6d814..e36d837239 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoader.java @@ -1,11 +1,11 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. * You may obtain a copy of the License at * - * http://www.apache.org/licenses/LICENSE-2.0 + * http://www.apache.org/licenses/LICENSE-2.0 * * Unless required by applicable law or agreed to in writing, software * distributed under the License is distributed on an "AS IS" BASIS, @@ -23,8 +23,6 @@ import org.springframework.stereotype.Component; -import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - import lombok.extern.slf4j.Slf4j; import lombok.var; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java index 8f67a4a4d3..c38ca4034e 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ColumnRecognizerFactory.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java index aab2f85985..46bbba6109 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactory.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.factory; import java.util.HashMap; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java index 7b8733ebaf..7fd990f5df 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/DefaultSensitiveType.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java index 4d5254a00e..e3c004a132 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/RecognitionResult.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,10 +13,8 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.model; -import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import lombok.Builder; import lombok.Data; @@ -31,4 +29,4 @@ public class RecognitionResult { // AI 规则 private String sensitiveType; // AI 判断出的具体敏感类型 -} \ No newline at end of file +} diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java index d9064356ce..9fe6da5253 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/ScanResult.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,11 +13,10 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.model; - import java.util.Optional; + import lombok.AllArgsConstructor; import lombok.Getter; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java index 03ac7eef62..68cfb2f262 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SensitiveColumnScanningTaskInfo.java @@ -44,7 +44,7 @@ public class SensitiveColumnScanningTaskInfo { private volatile boolean cancelled = false; public SensitiveColumnScanningTaskInfo(@NonNull String taskId, @NonNull Long projectId, - @NonNull Integer allTableCount) { + @NonNull Integer allTableCount) { this.taskId = taskId; this.projectId = projectId; this.status = ScanningTaskStatus.CREATED; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java index 09e638d4b1..04dbf36e76 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/model/SingleTableScanReq.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -15,7 +15,6 @@ */ package com.oceanbase.odc.service.datasecurity.model; - import javax.validation.constraints.NotBlank; import javax.validation.constraints.NotNull; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java index 5625d59f03..78d6bb0f61 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.strategy; import java.util.HashMap; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java index f6c562ecb7..472e2c7651 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.strategy; import java.util.HashMap; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java index 5b3f1ffd75..6156f6531b 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.strategy; import java.util.ArrayList; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java index 3a23302dd4..e209e31ec8 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.strategy; import java.util.HashMap; diff --git a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java index 917a98379f..b156e92c4f 100644 --- a/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java +++ b/server/odc-service/src/main/java/com/oceanbase/odc/service/datasecurity/strategy/ScanningStrategy.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -13,7 +13,6 @@ * See the License for the specific language governing permissions and * limitations under the License. */ - package com.oceanbase.odc.service.datasecurity.strategy; import java.util.List; diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java index b82d68aefe..7788d6cc77 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIConfigTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java index e720f3f882..977827bc26 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/AIInferenceServiceTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java index 6752fd8fd9..e663f3f679 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/ai/PromptTemplateLoaderTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java index 1a80b04fdd..3a2c3bdc3d 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/factory/ScanningStrategyFactoryTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java index 42cc6723ce..b72abe3592 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/AIColumnRecognizerTest.java @@ -28,8 +28,8 @@ import org.junit.rules.ExpectedException; import org.junit.runner.RunWith; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.MockedStatic; +import org.mockito.Mockito; import org.mockito.junit.MockitoJUnitRunner; import com.oceanbase.odc.core.shared.exception.BadRequestException; @@ -43,7 +43,6 @@ import com.oceanbase.tools.dbbrowser.model.DBTableColumn; import com.openai.models.chat.completions.ChatCompletion; - @RunWith(MockitoJUnitRunner.class) public class AIColumnRecognizerTest { diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java index ab575b7a5c..c885dc682e 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/GroovyColumnRecognizerTest.java @@ -29,7 +29,6 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - public class GroovyColumnRecognizerTest { diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java index 92965b4f7c..93f312d1a2 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/recognizer/RegexColumnRecognizerTest.java @@ -26,7 +26,6 @@ import com.oceanbase.odc.service.datasecurity.model.SensitiveRuleType; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - public class RegexColumnRecognizerTest { @Test diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java index a64f49ad69..1f2d94ebad 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AIOnlyStrategyTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java index 19ea6a5591..2e1317778c 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/AbstractScanningStrategyTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java index c6f08b72fb..bb94030861 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/JointRecognitionStrategyTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. diff --git a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java index fe731415d8..29940e30a5 100644 --- a/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java +++ b/server/odc-service/src/test/java/com/oceanbase/odc/service/datasecurity/strategy/RulesOnlyStrategyTest.java @@ -1,5 +1,5 @@ /* - * Copyright (c) 2025 OceanBase. + * Copyright (c) 2023 OceanBase. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -37,7 +37,6 @@ import com.oceanbase.odc.service.datasecurity.recognizer.ColumnRecognizer; import com.oceanbase.tools.dbbrowser.model.DBTableColumn; - @RunWith(MockitoJUnitRunner.class) public class RulesOnlyStrategyTest { From ecb352e956b4445bc876679a5521f80f4532a6c0 Mon Sep 17 00:00:00 2001 From: fenyf Date: Sat, 20 Sep 2025 01:08:55 +0800 Subject: [PATCH 10/10] =?UTF-8?q?fix(ai=5Frecognition):=E9=9B=86=E6=88=90?= =?UTF-8?q?=E6=B5=8B=E8=AF=95bug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../SensitiveColumnScanningTaskManagerTest.java | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java index 0a30b3cba7..d1123b82c7 100644 --- a/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java +++ b/server/integration-test/src/test/java/com/oceanbase/odc/service/datasecurity/SensitiveColumnScanningTaskManagerTest.java @@ -109,10 +109,10 @@ public void test_start_groovyRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createGroovySensitiveRules()); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); - Assert.assertEquals(3, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); + Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); } @Test @@ -120,7 +120,7 @@ public void test_start_groovyRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createGroovySensitiveRules()); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -131,7 +131,7 @@ public void test_start_pathRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createPathSensitiveRules()); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -142,7 +142,7 @@ public void test_start_pathRule_OBMOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createPathSensitiveRules()); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(20, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -153,7 +153,7 @@ public void test_start_RegexRule_OBMySQL() { List databases = createDatabases(ConnectType.OB_MYSQL); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_MYSQL)); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, mysqlConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, mysqlConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size()); @@ -164,7 +164,7 @@ public void test_start_RegexRule_OBOracle() { List databases = createDatabases(ConnectType.OB_ORACLE); List rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_ORACLE)); SensitiveColumnScanningTaskInfo taskInfo = - manager.start(databases, rules, ScanningModeType.JOINT_RECOGNITION, oracleConnectionConfig, null); + manager.start(databases, rules, ScanningModeType.RULES_ONLY, oracleConnectionConfig, null); await().atMost(20, SECONDS) .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS); Assert.assertEquals(6, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size());