Auto Upgrade

chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/model/DataSource.java

@@ -1,6 +1,7 @@
 package ai.chat2db.server.domain.api.model;
 
 import java.time.LocalDateTime;
+import java.util.LinkedHashMap;
 import java.util.List;
 
 import ai.chat2db.spi.config.DriverConfig;
@@ -8,6 +9,7 @@ import ai.chat2db.spi.model.KeyValue;
 import ai.chat2db.spi.model.SSHInfo;
 import ai.chat2db.spi.model.SSLInfo;
 import lombok.Data;
+import org.springframework.util.ObjectUtils;
 
 /**
  * @author moji
@@ -107,4 +109,16 @@ public class DataSource {
      * 驱动配置
      */
     private DriverConfig driverConfig;
+
+
+    public LinkedHashMap<String,Object> getExtendMap() {
+        if (ObjectUtils.isEmpty(extendInfo)) {
+            return new LinkedHashMap<>();
+        }
+        LinkedHashMap<String,Object> map = new LinkedHashMap<>();
+        for (KeyValue keyValue : extendInfo) {
+            map.put(keyValue.getKey(),keyValue.getValue());
+        }
+        return map;
+    }
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/DatabaseQueryAllParam.java

@@ -1,7 +1,8 @@
 package ai.chat2db.server.domain.api.param;
 
-import jakarta.validation.constraints.NotNull;
+import java.sql.Connection;
 
+import jakarta.validation.constraints.NotNull;
 import lombok.AllArgsConstructor;
 import lombok.Data;
 import lombok.NoArgsConstructor;
@@ -27,4 +28,14 @@ public class DatabaseQueryAllParam {
      * if true, refresh the cache
      */
     private boolean refresh;
+
+    /**
+     * Can be null, if null, use the default connection
+     */
+    private Connection connection;
+
+    /**
+     * Can be null, if null, use the default dbType
+     */
+    private String dbType;
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-api/src/main/java/ai/chat2db/server/domain/api/param/SchemaQueryParam.java

@@ -1,6 +1,8 @@
 
 package ai.chat2db.server.domain.api.param;
 
+import java.sql.Connection;
+
 import jakarta.validation.constraints.NotNull;
 
 import lombok.AllArgsConstructor;
@@ -29,4 +31,9 @@ public class SchemaQueryParam {
      * if true, refresh the cache
      */
     private boolean refresh;
+
+    /**
+     * Can be null, if null, use the default connection
+     */
+    private Connection connection;
 }

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DataSourceServiceImpl.java

@@ -1,5 +1,7 @@
 package ai.chat2db.server.domain.core.impl;
 
+import java.sql.Connection;
+import java.sql.SQLException;
 import java.time.LocalDateTime;
 import java.util.List;
 
@@ -13,6 +15,7 @@ import ai.chat2db.server.domain.api.param.DataSourceTestParam;
 import ai.chat2db.server.domain.api.param.DataSourceUpdateParam;
 import ai.chat2db.server.domain.api.param.DatabaseQueryAllParam;
 import ai.chat2db.server.domain.api.service.DataSourceService;
+import ai.chat2db.server.domain.api.service.DatabaseService;
 import ai.chat2db.server.domain.core.converter.DataSourceConverter;
 import ai.chat2db.server.domain.repository.entity.DataSourceDO;
 import ai.chat2db.server.domain.repository.mapper.DataSourceMapper;
@@ -25,11 +28,13 @@ import ai.chat2db.spi.model.DataSourceConnect;
 import ai.chat2db.spi.model.Database;
 import ai.chat2db.spi.model.KeyValue;
 import ai.chat2db.spi.sql.Chat2DBContext;
+import ai.chat2db.spi.sql.IDriverManager;
 import ai.chat2db.spi.sql.SQLExecutor;
 import ai.chat2db.spi.util.JdbcUtils;
 import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
 import com.baomidou.mybatisplus.core.metadata.IPage;
 import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
+import lombok.extern.slf4j.Slf4j;
 import org.apache.commons.lang3.BooleanUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -40,6 +45,7 @@ import org.springframework.stereotype.Service;
  * @version DataSourceCoreServiceImpl.java, v 0.1 2022年09月23日 15:51 moji Exp $
  * @date 2022/09/23
  */
+@Slf4j
 @Service
 public class DataSourceServiceImpl implements DataSourceService {
 
@@ -49,15 +55,37 @@ public class DataSourceServiceImpl implements DataSourceService {
     @Autowired
     private DataSourceConverter dataSourceConverter;
 
+    @Autowired
+    private DatabaseService databaseService;
+
     @Override
     public DataResult<Long> create(DataSourceCreateParam param) {
         DataSourceDO dataSourceDO = dataSourceConverter.param2do(param);
         dataSourceDO.setGmtCreate(LocalDateTime.now());
         dataSourceDO.setGmtModified(LocalDateTime.now());
         dataSourceMapper.insert(dataSourceDO);
+        preWarmingData(dataSourceDO.getId());
         return DataResult.of(dataSourceDO.getId());
     }
 
+    private void preWarmingData(Long dataSourceId) {
+        DataResult<DataSource> dataResult = queryById(dataSourceId);
+        if (dataResult.success() && dataResult.getData() != null) {
+            DataSource dataSource = dataResult.getData();
+            try (Connection connection = IDriverManager.getConnection(dataSource.getUrl(), dataSource.getUserName(),
+                dataSource.getPassword(), dataSource.getDriverConfig(), dataSource.getExtendMap())) {
+                DatabaseQueryAllParam databaseQueryAllParam = new DatabaseQueryAllParam();
+                databaseQueryAllParam.setDataSourceId(dataSourceId);
+                databaseQueryAllParam.setConnection(connection);
+                databaseQueryAllParam.setDbType(dataSource.getType());
+                databaseQueryAllParam.setRefresh(true);
+                databaseService.queryAll(databaseQueryAllParam);
+            } catch (SQLException e) {
+                log.error("preWarmingData error", e);
+            }
+        }
+    }
+
     @Override
     public ActionResult update(DataSourceUpdateParam param) {
         DataSourceDO dataSourceDO = dataSourceConverter.param2do(param);

chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DatabaseServiceImpl.java

@@ -41,8 +41,9 @@ public class DatabaseServiceImpl implements DatabaseService {
     public ListResult<Database> queryAll(DatabaseQueryAllParam param) {
         List<Database> databases = CacheManage.getList(getDataBasesKey(param.getDataSourceId()), Database.class,
             (key) -> param.isRefresh(), (key) -> {
-                Connection connection = Chat2DBContext.getConnection();
+                Connection connection = param.getConnection() == null ? Chat2DBContext.getConnection()
-                MetaData metaData = Chat2DBContext.getMetaData();
+                    : param.getConnection();
+                MetaData metaData = Chat2DBContext.getMetaData(param.getDbType());
                 return metaData.databases(connection);
             });
         return ListResult.of(databases);
@@ -53,7 +54,8 @@ public class DatabaseServiceImpl implements DatabaseService {
         List<Schema> schemas = CacheManage.getList(getSchemasKey(param.getDataSourceId(), param.getDataBaseName()),
             Schema.class,
             (key) -> param.isRefresh(), (key) -> {
-                Connection connection = Chat2DBContext.getConnection();
+                Connection connection = param.getConnection() == null ? Chat2DBContext.getConnection()
+                    : param.getConnection();
                 MetaData metaData = Chat2DBContext.getMetaData();
                 return metaData.schemas(connection, param.getDataBaseName());
             });

chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/DatabaseController.java

@@ -23,6 +23,9 @@ import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
+/**
+ * database controller
+ */
 @ConnectionInfoAspect
 @RequestMapping("/api/rdb/database")
 @RestController

chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/SchemaController.java

@@ -23,6 +23,9 @@ import org.springframework.web.bind.annotation.RequestBody;
 import org.springframework.web.bind.annotation.RequestMapping;
 import org.springframework.web.bind.annotation.RestController;
 
+/**
+ * shema controller
+ */
 @ConnectionInfoAspect
 @RequestMapping("/api/rdb/schema")
 @RestController

chat2db-server/chat2db-spi/src/main/java/ai/chat2db/spi/sql/Chat2DBContext.java

@@ -16,6 +16,7 @@ import ai.chat2db.spi.config.DriverConfig;
 import com.jcraft.jsch.JSchException;
 import com.jcraft.jsch.Session;
 import lombok.extern.slf4j.Slf4j;
+import org.apache.commons.lang3.StringUtils;
 
 /**
  * @author jipengfei
@@ -54,6 +55,13 @@ public class Chat2DBContext {
         return PLUGIN_MAP.get(getConnectInfo().getDbType()).getMetaData();
     }
 
+    public static MetaData getMetaData(String dbType) {
+        if(StringUtils.isBlank(dbType)){
+            return getMetaData();
+        }
+        return PLUGIN_MAP.get(dbType).getMetaData();
+    }
+
     public static DBConfig getDBConfig() {
         return PLUGIN_MAP.get(getConnectInfo().getDbType()).getDBConfig();
     }