From 85313e9c623ee1018997a7d103455ef9aedd0cb2 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Tue, 11 Aug 2026 13:29:29 +0800 Subject: [PATCH] fix(spark): preserve comments in generated table DDL Closes #5853 --- .../java/com/dbx/agent/spark/SparkAgent.java | 40 ++++ .../com/dbx/agent/spark/SparkAgentTest.java | 196 ++++++++++++++++++ 2 files changed, 236 insertions(+) create mode 100644 agents/drivers/spark/src/test/java/com/dbx/agent/spark/SparkAgentTest.java diff --git a/agents/drivers/spark/src/main/java/com/dbx/agent/spark/SparkAgent.java b/agents/drivers/spark/src/main/java/com/dbx/agent/spark/SparkAgent.java index e134b680aa..5cf2aa6a41 100644 --- a/agents/drivers/spark/src/main/java/com/dbx/agent/spark/SparkAgent.java +++ b/agents/drivers/spark/src/main/java/com/dbx/agent/spark/SparkAgent.java @@ -4,6 +4,7 @@ import com.dbx.agent.AbstractJdbcAgent; import com.dbx.agent.ColumnInfo; import com.dbx.agent.ConnectParams; import com.dbx.agent.DatabaseInfo; +import com.dbx.agent.DdlBuilder; import com.dbx.agent.ForeignKeyInfo; import com.dbx.agent.IndexInfo; import com.dbx.agent.JdbcIdentifiers; @@ -188,6 +189,31 @@ public final class SparkAgent extends AbstractJdbcAgent { return Collections.emptyList(); } + @Override + public String getTableDdl(String schema, String table) { + try { + String tableReference = qualifiedTableReference(schema, table); + try (java.sql.Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW CREATE TABLE " + tableReference)) { + if (rs.next()) { + String ddl = trimToNull(rs.getString(1)); + if (ddl != null) { + return ddl; + } + } + } + } catch (Exception ignored) { + // Some Spark-compatible Thrift endpoints may not expose SHOW CREATE TABLE. + } + return DdlBuilder.buildTableDdl( + schema, + table, + getColumns(schema, table), + Collections.emptyList(), + Collections.emptyList() + ); + } + @Override public String setSchemaSQL(String schema) { // When a catalog is configured, qualify the USE statement with the @@ -272,6 +298,20 @@ public final class SparkAgent extends AbstractJdbcAgent { return result; } + private String qualifiedTableReference(String schema, String table) { + if (configuredCatalog != null && !configuredCatalog.isEmpty()) { + String reference = JdbcIdentifiers.INSTANCE.backtick(configuredCatalog); + if (schema != null && !schema.isEmpty() && !schema.equals(configuredCatalog)) { + reference += "." + JdbcIdentifiers.INSTANCE.backtick(schema); + } + return reference + "." + JdbcIdentifiers.INSTANCE.backtick(table); + } + if (schema != null && !schema.isEmpty()) { + return JdbcIdentifiers.INSTANCE.backtick(schema) + "." + JdbcIdentifiers.INSTANCE.backtick(table); + } + return JdbcIdentifiers.INSTANCE.backtick(table); + } + private static List getColumnsFromMetadata(Connection conn, String schema, String table) throws Exception { List result = new ArrayList<>(); DatabaseMetaData meta = conn.getMetaData(); diff --git a/agents/drivers/spark/src/test/java/com/dbx/agent/spark/SparkAgentTest.java b/agents/drivers/spark/src/test/java/com/dbx/agent/spark/SparkAgentTest.java new file mode 100644 index 0000000000..ea9b529953 --- /dev/null +++ b/agents/drivers/spark/src/test/java/com/dbx/agent/spark/SparkAgentTest.java @@ -0,0 +1,196 @@ +package com.dbx.agent.spark; + +import com.dbx.agent.test.TestSupport; +import java.lang.reflect.Field; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Method; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.ResultSet; +import java.sql.ResultSetMetaData; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class SparkAgentTest { + @Test + void getTableDdlUsesNativeSparkDdlWithCommentsAndProviderClauses() { + List sql = new ArrayList<>(); + String nativeDdl = "CREATE TABLE spark_catalog.default.events (\n" + + " id INT COMMENT '编号注释',\n" + + " name STRING,\n" + + " amount DECIMAL(12,2) COMMENT '金额')\n" + + "USING parquet\n" + + "COMMENT 'Spark 表注释'"; + SparkAgent agent = new SparkAgent(); + TestSupport.setPrivateConnection(agent, connection(sql, nativeDdl, false)); + + String ddl = agent.getTableDdl("default", "events"); + + Assertions.assertEquals(nativeDdl, ddl); + Assertions.assertEquals( + List.of("SHOW CREATE TABLE `default`.`events`"), + sql + ); + } + + @Test + void getTableDdlQualifiesConfiguredCatalogWithoutChangingContext() throws Exception { + List sql = new ArrayList<>(); + SparkAgent agent = new SparkAgent(); + setConfiguredCatalog(agent, "lake`catalog"); + TestSupport.setPrivateConnection(agent, connection(sql, "CREATE TABLE native_ddl", false)); + + String ddl = agent.getTableDdl("analytics", "orders`archive"); + + Assertions.assertEquals("CREATE TABLE native_ddl", ddl); + Assertions.assertEquals( + List.of("SHOW CREATE TABLE `lake``catalog`.`analytics`.`orders``archive`"), + sql + ); + } + + @Test + void getTableDdlFallsBackToMetadataWhenShowCreateFails() { + List sql = new ArrayList<>(); + SparkAgent agent = new SparkAgent(); + TestSupport.setPrivateConnection(agent, connection(sql, null, true)); + + String ddl = agent.getTableDdl("default", "events"); + + Assertions.assertEquals( + "CREATE TABLE \"default\".\"events\" (\n" + + " \"id\" int,\n" + + " \"name\" string\n" + + ");\n", + ddl + ); + Assertions.assertEquals( + List.of( + "SHOW CREATE TABLE `default`.`events`", + "USE `default`", + "DESCRIBE `events`" + ), + sql + ); + } + + private static Connection connection(List sql, String nativeDdl, boolean failShowCreate) { + Statement statement = proxy(Statement.class, (method, args) -> { + if ("execute".equals(method.getName())) { + sql.add(String.valueOf(args[0])); + return false; + } + if ("executeQuery".equals(method.getName())) { + String query = String.valueOf(args[0]); + sql.add(query); + if (query.startsWith("SHOW CREATE TABLE")) { + if (failShowCreate) { + throw new SQLException("SHOW CREATE TABLE is unavailable"); + } + return resultSet( + new String[]{"createtab_stmt"}, + new Object[][]{{nativeDdl}} + ); + } + if (query.startsWith("DESCRIBE")) { + return resultSet( + new String[]{"col_name", "data_type", "comment"}, + new Object[][]{ + {"id", "int", "编号注释"}, + {"name", "string", null} + } + ); + } + } + if ("close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + return proxy(Connection.class, (method, args) -> { + if ("createStatement".equals(method.getName())) { + return statement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + } + + private static ResultSet resultSet(String[] columns, Object[][] rows) { + int[] index = {-1}; + ResultSetMetaData metadata = proxy(ResultSetMetaData.class, (method, args) -> { + if ("getColumnCount".equals(method.getName())) { + return columns.length; + } + return defaultValue(method.getReturnType()); + }); + return proxy(ResultSet.class, (method, args) -> { + if ("next".equals(method.getName())) { + index[0] += 1; + return index[0] < rows.length; + } + if ("getString".equals(method.getName())) { + Object key = args[0]; + int column = key instanceof Number + ? ((Number) key).intValue() - 1 + : findColumn(columns, String.valueOf(key)); + Object value = rows[index[0]][column]; + return value == null ? null : String.valueOf(value); + } + if ("getMetaData".equals(method.getName())) { + return metadata; + } + if ("close".equals(method.getName())) { + return null; + } + return defaultValue(method.getReturnType()); + }); + } + + private static int findColumn(String[] columns, String name) { + for (int i = 0; i < columns.length; i++) { + if (columns[i].equalsIgnoreCase(name)) { + return i; + } + } + throw new IllegalArgumentException("Unknown column: " + name); + } + + private static void setConfiguredCatalog(SparkAgent agent, String catalog) throws Exception { + Field field = SparkAgent.class.getDeclaredField("configuredCatalog"); + field.setAccessible(true); + field.set(agent, catalog); + } + + private static T proxy(Class type, MethodHandler handler) { + InvocationHandler invocationHandler = new InvocationHandler() { + @Override + public Object invoke(Object proxy, Method method, Object[] args) throws Throwable { + return handler.handle(method, args == null ? new Object[0] : args); + } + }; + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, invocationHandler)); + } + + private static Object defaultValue(Class type) { + if (type == Boolean.TYPE) return false; + if (type == Byte.TYPE) return (byte) 0; + if (type == Short.TYPE) return (short) 0; + if (type == Integer.TYPE) return 0; + if (type == Long.TYPE) return 0L; + if (type == Float.TYPE) return 0f; + if (type == Double.TYPE) return 0d; + if (type == Character.TYPE) return (char) 0; + return null; + } + + private interface MethodHandler { + Object handle(Method method, Object[] args) throws Throwable; + } +}