support for mongodb

2025-08-02 21:50:43 +08:00 · 2023-11-26 16:14:15 +08:00
parent 5549908aaa
commit f1176af805
3 changed files with 124 additions and 85 deletions
--- a/chat2db-server/chat2db-plugins/chat2db-mongodb/src/main/java/ai/chat2db/plugin/mongodb/MongodbManage.java
+++ b/chat2db-server/chat2db-plugins/chat2db-mongodb/src/main/java/ai/chat2db/plugin/mongodb/MongodbManage.java
@ -5,17 +5,25 @@ import java.sql.SQLException;

 import ai.chat2db.spi.DBManage;
 import ai.chat2db.spi.jdbc.DefaultDBManage;
+import ai.chat2db.spi.sql.Chat2DBContext;
+import ai.chat2db.spi.sql.ConnectInfo;
 import ai.chat2db.spi.sql.SQLExecutor;
+import org.apache.commons.lang3.ObjectUtils;
 import org.springframework.util.StringUtils;

 public class MongodbManage extends DefaultDBManage implements DBManage {
    @Override
    public void connectDatabase(Connection connection, String database) {
-        if (StringUtils.isEmpty(database)) {
+        ConnectInfo connectInfo = Chat2DBContext.getConnectInfo();
+        if (ObjectUtils.anyNull(connectInfo) || StringUtils.isEmpty(connectInfo.getSchemaName())) {
+            return;
+        }
+        String schemaName = connectInfo.getSchemaName();
+        if (StringUtils.isEmpty(schemaName)) {
            return;
        }
        try {
-            SQLExecutor.getInstance().execute(connection, "use " + database + ";");
+            SQLExecutor.getInstance().execute(connection, "use " + schemaName + ";");
        } catch (SQLException e) {
            throw new RuntimeException(e);
        }
--- a/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java
+++ b/chat2db-server/chat2db-server-domain/chat2db-server-domain-core/src/main/java/ai/chat2db/server/domain/core/impl/DlTemplateServiceImpl.java
@ -2,18 +2,14 @@ package ai.chat2db.server.domain.core.impl;

 import java.sql.Connection;
 import java.sql.SQLException;
-import java.util.*;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;

-import ai.chat2db.server.domain.api.param.*;
-import ai.chat2db.server.domain.api.param.operation.OperationLogCreateParam;
-import ai.chat2db.server.domain.api.service.OperationLogService;
-import ai.chat2db.server.domain.api.service.TableService;
-import ai.chat2db.server.domain.core.util.MetaNameUtils;
-import ai.chat2db.spi.MetaData;
-import ai.chat2db.spi.ValueHandler;
-import ai.chat2db.spi.model.*;
-import ai.chat2db.spi.sql.ConnectInfo;
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.PagerUtils;
 import com.alibaba.druid.sql.SQLUtils;
@ -21,16 +17,34 @@ import com.alibaba.druid.sql.ast.SQLStatement;
 import com.alibaba.druid.sql.ast.statement.SQLSelectStatement;
 import com.alibaba.druid.sql.parser.ParserException;

+import ai.chat2db.server.domain.api.param.DlCountParam;
+import ai.chat2db.server.domain.api.param.DlExecuteParam;
+import ai.chat2db.server.domain.api.param.SelectResultOperation;
+import ai.chat2db.server.domain.api.param.TableQueryParam;
+import ai.chat2db.server.domain.api.param.UpdateSelectResultParam;
+import ai.chat2db.server.domain.api.param.operation.OperationLogCreateParam;
 import ai.chat2db.server.domain.api.service.DlTemplateService;
+import ai.chat2db.server.domain.api.service.OperationLogService;
+import ai.chat2db.server.domain.api.service.TableService;
+import ai.chat2db.server.domain.core.util.MetaNameUtils;
 import ai.chat2db.server.tools.base.constant.EasyToolsConstant;
+import ai.chat2db.server.tools.base.enums.DataSourceTypeEnum;
 import ai.chat2db.server.tools.base.excption.BusinessException;
 import ai.chat2db.server.tools.base.wrapper.result.DataResult;
 import ai.chat2db.server.tools.base.wrapper.result.ListResult;
 import ai.chat2db.server.tools.common.util.EasyCollectionUtils;
 import ai.chat2db.server.tools.common.util.I18nUtils;
+import ai.chat2db.spi.MetaData;
+import ai.chat2db.spi.ValueHandler;
 import ai.chat2db.spi.enums.DataTypeEnum;
 import ai.chat2db.spi.enums.SqlTypeEnum;
+import ai.chat2db.spi.model.ExecuteResult;
+import ai.chat2db.spi.model.Header;
+import ai.chat2db.spi.model.TableColumn;
+import ai.chat2db.spi.model.TableIndex;
+import ai.chat2db.spi.model.TableIndexColumn;
 import ai.chat2db.spi.sql.Chat2DBContext;
+import ai.chat2db.spi.sql.ConnectInfo;
 import ai.chat2db.spi.sql.SQLExecutor;
 import ai.chat2db.spi.util.JdbcUtils;
 import ai.chat2db.spi.util.SqlUtils;
@ -71,7 +85,6 @@ public class DlTemplateServiceImpl implements DlTemplateService {

        List<String> sqlList = SqlUtils.parse(param.getSql(), dbType);

-
        if (CollectionUtils.isEmpty(sqlList)) {
            throw new BusinessException("dataSource.sqlAnalysisError");
        }
@ -127,29 +140,34 @@ public class DlTemplateServiceImpl implements DlTemplateService {
        param.setSql(sql);
    }

-
    private ExecuteResult executeSQL(String originalSql, DbType dbType, DlExecuteParam param) {
        int pageNo = 1;
        int pageSize = 0;
        Integer offset = null;
        Integer count = null;
        String sqlType = SqlTypeEnum.UNKNOWN.getCode();
-
+        // 解析sql
+        String type = Chat2DBContext.getConnectInfo().getDbType();
+        boolean supportDruid = !DataSourceTypeEnum.MONGODB.getCode().equals(type);
        // 解析sql分页
-        SQLStatement sqlStatement;
+        SQLStatement sqlStatement = null;
+        if (supportDruid) {
            try {
                sqlStatement = SQLUtils.parseSingleStatement(originalSql, dbType);
-            // 是否需要代码帮忙分页
-            if (sqlStatement instanceof SQLSelectStatement) {
+            } catch (ParserException e) {
+                log.warn("解析sql失败:{}", originalSql, e);
+            }
+        }
+
+        // Mongodb is currently unable to recognize it, so every time a page is transmitted
+        if (!supportDruid || (sqlStatement instanceof SQLSelectStatement)) {
            pageNo = Optional.ofNullable(param.getPageNo()).orElse(1);
            pageSize = Optional.ofNullable(param.getPageSize()).orElse(EasyToolsConstant.MAX_PAGE_SIZE);
            offset = (pageNo - 1) * pageSize;
            count = pageSize;
            sqlType = SqlTypeEnum.SELECT.getCode();
        }
-        } catch (ParserException e) {
-            log.warn("解析sql失败:{}", originalSql, e);
-        }
+
        ExecuteResult executeResult = null;
        if (SqlTypeEnum.SELECT.getCode().equals(sqlType) && !SqlUtils.hasPageLimit(originalSql, dbType)) {
            String pageLimit = Chat2DBContext.getSqlBuilder().pageLimit(originalSql, offset, pageNo, pageSize);
@ -161,14 +179,18 @@ public class DlTemplateServiceImpl implements DlTemplateService {
            executeResult = execute(originalSql, offset, count);
        }

-
        executeResult.setSqlType(sqlType);
        executeResult.setOriginalSql(originalSql);
+
+        boolean supportJsqlParser = !DataSourceTypeEnum.MONGODB.getCode().equals(type);
+        if (supportJsqlParser) {
            try {
                SqlUtils.buildCanEditResult(originalSql, dbType, executeResult);
            } catch (Exception e) {
                log.warn("buildCanEditResult error", e);
            }
+        }
+
        if (SqlTypeEnum.SELECT.getCode().equals(sqlType)) {
            executeResult.setPageNo(pageNo);
            executeResult.setPageSize(pageSize);
@ -181,8 +203,9 @@ public class DlTemplateServiceImpl implements DlTemplateService {
        }

        List<Header> headers = executeResult.getHeaderList();
-        if (executeResult.getSuccess() && executeResult.isCanEdit() && CollectionUtils.isNotEmpty(headers)){
-            headers = setColumnInfo(headers, executeResult.getTableName(), param.getSchemaName(), param.getDatabaseName());
+        if (executeResult.getSuccess() && executeResult.isCanEdit() && CollectionUtils.isNotEmpty(headers)) {
+            headers = setColumnInfo(headers, executeResult.getTableName(), param.getSchemaName(),
+                param.getDatabaseName());
        }
        Header rowNumberHeader = Header.builder()
            .name(I18nUtils.getMessage("sqlResult.rowNumber"))
@ -285,7 +308,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {
        return keyColumns;
    }

-    private String getDeleteSql(UpdateSelectResultParam param, List<String> row, MetaData metaSchema, List<String> keyColumns) {
+    private String getDeleteSql(UpdateSelectResultParam param, List<String> row, MetaData metaSchema,
+        List<String> keyColumns) {
        StringBuilder script = new StringBuilder();
        script.append("DELETE FROM ").append(param.getTableName()).append("");

@ -365,8 +389,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {

    }

-
-    private String getUpdateSql(UpdateSelectResultParam param, List<String> row, List<String> odlRow, MetaData metaSchema,
+    private String getUpdateSql(UpdateSelectResultParam param, List<String> row, List<String> odlRow,
+        MetaData metaSchema,
        List<String> keyColumns, boolean copy) {
        StringBuilder script = new StringBuilder();
        if (CollectionUtils.isEmpty(row) || CollectionUtils.isEmpty(odlRow)) {
@ -402,7 +426,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {
            if (CollectionUtils.isEmpty(columns)) {
                return headers;
            }
-            Map<String, TableColumn> columnMap = columns.stream().collect(Collectors.toMap(TableColumn::getName, tableColumn -> tableColumn));
+            Map<String, TableColumn> columnMap = columns.stream().collect(
+                Collectors.toMap(TableColumn::getName, tableColumn -> tableColumn));
            List<TableIndex> tableIndices = tableService.queryIndexes(tableQueryParam);
            if (!CollectionUtils.isEmpty(tableIndices)) {
                for (TableIndex tableIndex : tableIndices) {
@ -431,8 +456,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {
                }
            }

-        }catch (Exception e){
-            log.error("setColumnInfo error:",e);
+        } catch (Exception e) {
+            log.error("setColumnInfo error:", e);
        }
        return headers;
    }
@ -441,7 +466,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {
        ExecuteResult executeResult;
        try {
            ValueHandler valueHandler = Chat2DBContext.getMetaData().getValueHandler();
-            executeResult = SQLExecutor.getInstance().execute(sql, Chat2DBContext.getConnection(), true, offset, count,valueHandler);
+            executeResult = SQLExecutor.getInstance().execute(sql, Chat2DBContext.getConnection(), true, offset, count,
+                valueHandler);
        } catch (SQLException e) {
            log.warn("执行sql:{}异常", sql, e);
            executeResult = ExecuteResult.builder()
@ -467,7 +493,8 @@ public class DlTemplateServiceImpl implements DlTemplateService {
            createParam.setSchemaName(connectInfo.getSchemaName());
            createParam.setUseTime(executeResult.getDuration());
            createParam.setType(connectInfo.getDbType());
-            createParam.setOperationRows(executeResult.getUpdateCount() != null ? Long.valueOf(executeResult.getUpdateCount()) : null);
+            createParam.setOperationRows(
+                executeResult.getUpdateCount() != null ? Long.valueOf(executeResult.getUpdateCount()) : null);
            operationLogService.create(createParam);
        } catch (Exception e) {
            log.error("addOperationLog error:", e);
--- a/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java
+++ b/chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/rdb/RdbDmlExportController.java
@ -7,7 +7,6 @@ import java.nio.charset.StandardCharsets;
 import java.time.LocalDateTime;
 import java.util.List;

-import ai.chat2db.spi.jdbc.DefaultValueHandler;
 import com.alibaba.druid.DbType;
 import com.alibaba.druid.sql.SQLUtils;
 import com.alibaba.druid.sql.SQLUtils.FormatOption;
@ -28,11 +27,11 @@ import ai.chat2db.server.domain.api.enums.ExportSizeEnum;
 import ai.chat2db.server.domain.api.enums.ExportTypeEnum;
 import ai.chat2db.server.tools.base.excption.BusinessException;
 import ai.chat2db.server.tools.common.exception.ParamBusinessException;
-import ai.chat2db.server.tools.common.util.ConfigUtils;
 import ai.chat2db.server.tools.common.util.EasyCollectionUtils;
 import ai.chat2db.server.tools.common.util.EasyEnumUtils;
 import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect;
 import ai.chat2db.server.web.api.controller.rdb.request.DataExportRequest;
+import ai.chat2db.spi.jdbc.DefaultValueHandler;
 import ai.chat2db.spi.sql.Chat2DBContext;
 import ai.chat2db.spi.sql.SQLExecutor;
 import ai.chat2db.spi.util.JdbcUtils;
@ -92,11 +91,17 @@ public class RdbDmlExportController {
            throw new ParamBusinessException("exportSize");
        }
        DbType dbType = JdbcUtils.parse2DruidDbType(Chat2DBContext.getConnectInfo().getDbType());
+        String tableName;
+        if (dbType != null) {
            SQLStatement sqlStatement = SQLUtils.parseSingleStatement(sql, dbType);
            if (!(sqlStatement instanceof SQLSelectStatement)) {
                throw new BusinessException("dataSource.sqlAnalysisError");
            }
-        String tableName = SqlUtils.getTableName(sql, dbType);
+            tableName = SqlUtils.getTableName(sql, dbType);
+        } else {
+            tableName = StringUtils.join(Lists.newArrayList(request.getDatabaseName(), request.getSchemaName()), "_");
+        }
+
        response.setCharacterEncoding("utf-8");
        String fileName = URLEncoder.encode(
                tableName + "_" + LocalDateTime.now().format(DatePattern.PURE_DATETIME_FORMATTER),
@ -108,7 +113,6 @@ public class RdbDmlExportController {
        } else {
            doExportInsert(sql, response, fileName, dbType, tableName);
        }
-        String SS = ConfigUtils.APP_PATH;
    }

    private void doExportCsv(String sql, HttpServletResponse response, String fileName)
@ -131,7 +135,7 @@ public class RdbDmlExportController {
                List<List<String>> writeDataList = Lists.newArrayList();
                writeDataList.add(dataList);
                excelWrapper.getExcelWriter().write(writeDataList, excelWrapper.getWriteSheet());
-            }, false,new DefaultValueHandler());
+            }, false, new DefaultValueHandler());
        } finally {
            if (excelWrapper.getExcelWriter() != null) {
                excelWrapper.getExcelWriter().finish();
@ -162,7 +166,7 @@ public class RdbDmlExportController {
                    sqlInsertStatement.setValues(valuesClause);

                    printWriter.println(SQLUtils.toSQLString(sqlInsertStatement, dbType, INSERT_FORMAT_OPTION) + ";");
-                }, false,new DefaultValueHandler());
+                }, false, new DefaultValueHandler());
        }
    }