Implement AI-based automatic identification of sensitive columns in databases

.gitattributes

@@ -0,0 +1 @@
+* text=auto eol=lf

pom.xml

@@ -196,6 +196,17 @@
                 <version>${springboot.version}</version>
             </dependency>
 
+            <!-- OpenAI SDK -->
+            <dependency>
+                <groupId>com.openai</groupId>
+                <artifactId>openai-java</artifactId>
+                <version>2.16.0</version>
+            </dependency>
+            <dependency>
+                <groupId>org.jetbrains.kotlin</groupId>
+                <artifactId>kotlin-stdlib-jdk7</artifactId>
+                <version>1.9.24</version>
+            </dependency>
 
             <dependency>
                 <groupId>org.codehaus.groovy</groupId>

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,8 @@ public static void tearDown() {
     public void test_start_groovyRule_OBMySQL() {
         List<Database> databases = createDatabases(ConnectType.OB_MYSQL);
         List<SensitiveRule> rules = Arrays.asList(createGroovySensitiveRules());
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                manager.start(databases, rules, ScanningModeType.RULES_ONLY, mysqlConnectionConfig, null);
         await().atMost(20, SECONDS)
                 .until(() -> manager.get(taskInfo.getTaskId()).getStatus() == ScanningTaskStatus.SUCCESS);
         Assert.assertEquals(2, manager.get(taskInfo.getTaskId()).getSensitiveColumns().size());
@@ -117,7 +119,8 @@ public void test_start_groovyRule_OBMySQL() {
     public void test_start_groovyRule_OBOracle() {
         List<Database> databases = createDatabases(ConnectType.OB_ORACLE);
         List<SensitiveRule> rules = Arrays.asList(createGroovySensitiveRules());
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                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());
@@ -127,7 +130,8 @@ public void test_start_groovyRule_OBOracle() {
     public void test_start_pathRule_OBMySQL() {
         List<Database> databases = createDatabases(ConnectType.OB_MYSQL);
         List<SensitiveRule> rules = Arrays.asList(createPathSensitiveRules());
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                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());
@@ -137,7 +141,8 @@ public void test_start_pathRule_OBMySQL() {
     public void test_start_pathRule_OBMOracle() {
         List<Database> databases = createDatabases(ConnectType.OB_ORACLE);
         List<SensitiveRule> rules = Arrays.asList(createPathSensitiveRules());
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                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());
@@ -147,7 +152,8 @@ public void test_start_pathRule_OBMOracle() {
     public void test_start_RegexRule_OBMySQL() {
         List<Database> databases = createDatabases(ConnectType.OB_MYSQL);
         List<SensitiveRule> rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_MYSQL));
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, mysqlConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                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());
@@ -157,7 +163,8 @@ public void test_start_RegexRule_OBMySQL() {
     public void test_start_RegexRule_OBOracle() {
         List<Database> databases = createDatabases(ConnectType.OB_ORACLE);
         List<SensitiveRule> rules = Arrays.asList(createRegexSensitiveRules(ConnectType.OB_ORACLE));
-        SensitiveColumnScanningTaskInfo taskInfo = manager.start(databases, rules, oracleConnectionConfig, null);
+        SensitiveColumnScanningTaskInfo taskInfo =
+                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());

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,
+
 
     ;
 

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.
+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}

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=文件验签不通过，请勿修改导出文件或检查密钥是否正确
+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}

server/odc-migrate/src/main/resources/migrate/common/V_4_3_4_20__alter_sensitive_rule.sql

@@ -0,0 +1,7 @@
+-- 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_custom_prompt` text default null comment 'User-defined custom prompt for AI rules.';
+

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`;

server/odc-server/src/main/java/com/oceanbase/odc/server/web/controller/v2/AIController.java

@@ -0,0 +1,63 @@
+/*
+ * 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.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.AIInferenceService;
+import com.oceanbase.odc.service.datasecurity.ai.AIStatusResponse;
+
+import io.swagger.annotations.Api;
+import io.swagger.annotations.ApiOperation;
+
+/**
+ * @author fenyf
+ * @date 2025/8/10 12:41
+ */
+@Api(tags = "AI")
+@RestController
+@RequestMapping("/api/v2/ai")
+public class AIController {
+
+    @Autowired
+    private AIConfig aiConfig;
+
+    @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<AIStatusResponse> getAIStatus() {
+        AIStatusResponse response = new AIStatusResponse();
+        response.setEnabled(aiConfig.isEnabled());
+        response.setAvailable(aiInferenceService.isAIAvailable());
+        response.setModel(aiConfig.getModel());
+        response.setBaseUrl(aiConfig.getBaseUrl());
+        response.setApiKeyConfigured(aiConfig.getApiKey() != null && !aiConfig.getApiKey().trim().isEmpty());
+
+        return Responses.success(response);
+    }
+}

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;
@@ -143,4 +145,25 @@ public SuccessResponse<SensitiveColumnScanningTaskInfo> 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<Boolean> stopScanning(@PathVariable Long projectId,
+            @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<SingleTableScanTaskManager.SingleTableScanTask> 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<String> scanSingleTableAsync(@PathVariable Long projectId,
+            @RequestBody SingleTableScanReq req) {
+        return Responses.success(service.scanSingleTableAsync(projectId, req));
+    }
 }

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;

server/odc-service/pom.xml

@@ -75,6 +75,14 @@
             <groupId>commons-beanutils</groupId>
             <artifactId>commons-beanutils</artifactId>
         </dependency>
+        <dependency>
+            <groupId>com.openai</groupId>
+            <artifactId>openai-java</artifactId>
+        </dependency>
+        <dependency>
+            <groupId>org.jetbrains.kotlin</groupId>
+            <artifactId>kotlin-stdlib-jdk7</artifactId>
+        </dependency>
         <dependency>
             <groupId>org.apache.commons</groupId>
             <artifactId>commons-compress</artifactId>

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[] {
             "/",

server/odc-service/src/main/java/com/oceanbase/odc/metadb/datasecurity/SensitiveRuleEntity.java

@@ -113,4 +113,10 @@ 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<String> aiSensitiveTypes;
+
+    @Column(name = "ai_custom_prompt")
+    private String aiCustomPrompt;
 }

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<Long> getAlgorithmIdByName(@NonNull String algorithmName, @NonNull Long organizationId) {
+        List<MaskingAlgorithmEntity> 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<MaskingAlgorithm> getMaskingAlgorithmsByOrganizationId(@NonNull Long organizationId) {
         return organizationId2Algorithms.get(organizationId);