fix(spark): preserve comments in generated table DDL

Closes #5853
2026-11-15 22:53:23 +00:00 · 2026-08-11 13:29:29 +08:00
1 parent 0c407c2d03
commit 85313e9c62
@@ -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<ColumnInfo> getColumnsFromMetadata(Connection conn, String schema, String table) throws Exception {
        List<ColumnInfo> result = new ArrayList<>();
        DatabaseMetaData meta = conn.getMetaData();
@@ -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<String> 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<String> 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<String> 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<String> 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> T proxy(Class<T> 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;
+    }
+}