diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnum.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnum.java index 914b07746..794b0cda3 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnum.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnum.java @@ -159,6 +159,10 @@ public static boolean isOracleDialect(String typeName) { return isDialect(typeName, DatabaseDialectEnum.ORACLE.getCode()); } + public static boolean isHiveDialect(String typeName) { + return isDialect(typeName, DatabaseDialectEnum.HIVE.getCode()); + } + public static boolean isAdbPg(String typeName) { BizDataSourceTypeEnum te = fromTypeName(typeName); if (te == null) { diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/schema/TableMetadataService.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/schema/TableMetadataService.java index a75c591d9..01da71cfa 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/schema/TableMetadataService.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/service/schema/TableMetadataService.java @@ -233,7 +233,7 @@ private Map> fetchTableSampleData(DbConfigBO dbConfig, Acce try { // 构建批量查询SQL,一次查询多个列的样本数据 - String columnNames = columns.stream().map(ColumnInfoBO::getName).collect(Collectors.joining(", ")); + String columnNames = SqlUtil.buildColumnsSql(dbConfig.getDialectType(), columns); String sql = SqlUtil.buildSelectSql(dbConfig.getDialectType(), tableName, columnNames, 5); DbQueryParameter batchParam = new DbQueryParameter(); diff --git a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/SqlUtil.java b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/SqlUtil.java index bd297c59b..5b309d248 100644 --- a/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/SqlUtil.java +++ b/data-agent-management/src/main/java/com/alibaba/cloud/ai/dataagent/util/SqlUtil.java @@ -15,8 +15,13 @@ */ package com.alibaba.cloud.ai.dataagent.util; +import com.alibaba.cloud.ai.dataagent.bo.schema.ColumnInfoBO; import com.alibaba.cloud.ai.dataagent.enums.BizDataSourceTypeEnum; import lombok.experimental.UtilityClass; +import org.apache.commons.collections.CollectionUtils; + +import java.util.List; +import java.util.stream.Collectors; /** * SQL 工具类 @@ -27,6 +32,22 @@ @UtilityClass public class SqlUtil { + /** + * 构建Columns SQL块 + * @param typeName 数据源类型 + * @param columnInfoBOs 列名 + * @return Columns SQL块 + */ + public static String buildColumnsSql(String typeName, List columnInfoBOs) { + if (CollectionUtils.isEmpty(columnInfoBOs)) { + return "*"; + } + if (BizDataSourceTypeEnum.isMysqlDialect(typeName) || BizDataSourceTypeEnum.isHiveDialect(typeName)) { + // mysql/hive使用``标识对象 + return columnInfoBOs.stream().map(ColumnInfoBO::getName).map(name -> String.format("`%s`", name)).collect(Collectors.joining(", ")); + } + return columnInfoBOs.stream().map(ColumnInfoBO::getName).map(name -> String.format("\"%s\"", name)).collect(Collectors.joining(", ")); + } /** * 构建SELECT SQL语句 * @param typeName 数据源类型 diff --git a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnumTest.java b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnumTest.java index c4d710153..67ba96a78 100644 --- a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnumTest.java +++ b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/enums/BizDataSourceTypeEnumTest.java @@ -146,6 +146,16 @@ void isOracleDialect_falseForOthers() { assertFalse(BizDataSourceTypeEnum.isOracleDialect("mysql")); } + @Test + void isHiveDialect_trueForHive() { + assertTrue(BizDataSourceTypeEnum.isHiveDialect("hive")); + } + + @Test + void isHiveDialect_falseForOthers() { + assertFalse(BizDataSourceTypeEnum.isHiveDialect("mysql")); + } + @Test void isAdbPg_trueForAdgPg() { assertTrue(BizDataSourceTypeEnum.isAdbPg("adg_pg")); diff --git a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/util/SqlUtilTest.java b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/util/SqlUtilTest.java index ac54eb82b..e05a2f6e6 100644 --- a/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/util/SqlUtilTest.java +++ b/data-agent-management/src/test/java/com/alibaba/cloud/ai/dataagent/util/SqlUtilTest.java @@ -15,8 +15,12 @@ */ package com.alibaba.cloud.ai.dataagent.util; +import com.alibaba.cloud.ai.dataagent.bo.schema.ColumnInfoBO; import org.junit.jupiter.api.Test; +import java.util.Arrays; +import java.util.List; + import static org.junit.jupiter.api.Assertions.*; class SqlUtilTest { @@ -73,4 +77,81 @@ void buildSelectSql_emptyTableName_throwsException() { assertThrows(IllegalArgumentException.class, () -> SqlUtil.buildSelectSql("mysql", "", "*", 10)); } + @Test + void buildColumnsSql_mysql_appendsLimit() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("mysql", columnInfoBOS); + assertEquals("`id`, `name`, `desc`", sql); + } + @Test + void buildColumnsSql_hive_appendsLimit() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("hive", columnInfoBOS); + assertEquals("`id`, `name`, `desc`", sql); + } + + @Test + void buildColumnsSql_postgresql_appendsLimit() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("postgresql", columnInfoBOS); + assertEquals("\"id\", \"name\", \"desc\"", sql); + } + + @Test + void buildColumnsSql_h2_appendsLimit() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("h2", columnInfoBOS); + assertEquals("\"id\", \"name\", \"desc\"", sql); + } + + @Test + void buildColumnsSql_sqlserver_prependsTop() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("sqlserver", columnInfoBOS); + assertEquals("\"id\", \"name\", \"desc\"", sql); + } + + @Test + void buildColumnsSql_oracle_appendsFetchFirst() { + List columnInfoBOS = Arrays.asList( + new ColumnInfoBO("id", "users", "id", "int", true, true, null), + new ColumnInfoBO("name", "users", "name", "varchar", false, false, null), + new ColumnInfoBO("desc", "users", "desc", "varchar", false, false, null) + ); + String sql = SqlUtil.buildColumnsSql("oracle", columnInfoBOS); + assertEquals("\"id\", \"name\", \"desc\"", sql); + } + + @Test + void buildColumnsSql_nullColumnNames_defaultsToStar() { + String sql = SqlUtil.buildColumnsSql("mysql", null); + assertEquals("*", sql); + } + + @Test + void buildColumnsSql_emptyColumnNames_defaultsToStar() { + String sql = SqlUtil.buildColumnsSql("mysql", Arrays.asList()); + assertEquals("*", sql); + } + }