From 8299084c95fcabfba2533653c7e43b576b5d7ddb Mon Sep 17 00:00:00 2001
From: lxwcodemonkey <jolunoluo@tencent.com>
Date: Sat, 30 Nov 2024 21:16:46 +0800
Subject: [PATCH] [improvement][Headless] Add get database by type

---
 .../headless/server/service/DatabaseService.java       |  3 +++
 .../server/service/impl/DatabaseServiceImpl.java       | 10 ++++++++++
 2 files changed, 13 insertions(+)

diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/DatabaseService.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/DatabaseService.java
index 8d6811a6c..a556d4bcc 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/DatabaseService.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/DatabaseService.java
@@ -2,6 +2,7 @@ package com.tencent.supersonic.headless.server.service;
 
 import com.tencent.supersonic.common.pojo.User;
 import com.tencent.supersonic.headless.api.pojo.DBColumn;
+import com.tencent.supersonic.headless.api.pojo.enums.DataType;
 import com.tencent.supersonic.headless.api.pojo.request.DatabaseReq;
 import com.tencent.supersonic.headless.api.pojo.request.ModelBuildReq;
 import com.tencent.supersonic.headless.api.pojo.request.SqlExecuteReq;
@@ -17,6 +18,8 @@ public interface DatabaseService {
 
     SemanticQueryResp executeSql(String sql, DatabaseResp databaseResp);
 
+    List<DatabaseResp> getDatabaseByType(DataType dataType);
+
     SemanticQueryResp executeSql(SqlExecuteReq sqlExecuteReq, Long id, User user);
 
     DatabaseResp getDatabase(Long id, User user);
diff --git a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
index 975847918..61a1d0906 100644
--- a/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
+++ b/headless/server/src/main/java/com/tencent/supersonic/headless/server/service/impl/DatabaseServiceImpl.java
@@ -1,11 +1,13 @@
 package com.tencent.supersonic.headless.server.service.impl;
 
+import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
 import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
 import com.google.common.collect.Lists;
 import com.tencent.supersonic.common.pojo.QueryColumn;
 import com.tencent.supersonic.common.pojo.User;
 import com.tencent.supersonic.common.pojo.enums.EngineType;
 import com.tencent.supersonic.headless.api.pojo.DBColumn;
+import com.tencent.supersonic.headless.api.pojo.enums.DataType;
 import com.tencent.supersonic.headless.api.pojo.request.DatabaseReq;
 import com.tencent.supersonic.headless.api.pojo.request.ModelBuildReq;
 import com.tencent.supersonic.headless.api.pojo.request.SqlExecuteReq;
@@ -131,6 +133,14 @@ public class DatabaseServiceImpl extends ServiceImpl<DatabaseDOMapper, DatabaseD
         return databaseResp;
     }
 
+    @Override
+    public List<DatabaseResp> getDatabaseByType(DataType dataType) {
+        QueryWrapper<DatabaseDO> queryWrapper = new QueryWrapper<>();
+        queryWrapper.lambda().eq(DatabaseDO::getType, dataType.getFeature());
+        List<DatabaseDO> list = list(queryWrapper);
+        return list.stream().map(DatabaseConverter::convertWithPassword).collect(Collectors.toList());
+    }
+
     @Override
     public SemanticQueryResp executeSql(SqlExecuteReq sqlExecuteReq, Long id, User user) {
         DatabaseResp databaseResp = getDatabase(id);