diff --git a/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java b/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java index 2001ccc32..01c2766cd 100644 --- a/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java +++ b/agents/drivers/hive/src/main/java/com/dbx/agent/hive/HiveAgent.java @@ -17,6 +17,7 @@ import com.dbx.agent.QueryResult; import java.sql.Connection; import java.sql.DatabaseMetaData; import java.sql.ResultSet; +import java.sql.Statement; import java.util.ArrayList; import java.util.Collections; import java.util.Comparator; @@ -86,6 +87,26 @@ public final class HiveAgent extends AbstractJdbcAgent { }); } + @Override + public String getTableDdl(String schema, String table) { + return unchecked(() -> { + String qualifiedName = schema == null || schema.trim().isEmpty() + ? JdbcIdentifiers.INSTANCE.backtick(table) + : JdbcIdentifiers.INSTANCE.backtick(schema) + "." + JdbcIdentifiers.INSTANCE.backtick(table); + try (Statement stmt = requireConnected().createStatement(); + ResultSet rs = stmt.executeQuery("SHOW CREATE TABLE " + qualifiedName)) { + StringBuilder ddl = new StringBuilder(); + while (rs.next()) { + String line = rs.getString(1); + if (line != null) { + ddl.append(line).append('\n'); + } + } + return ddl.toString(); + } + }); + } + @Override public List listIndexes(String schema, String table) { return Collections.emptyList(); diff --git a/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java index f3496cdae..46691cb11 100644 --- a/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java +++ b/agents/drivers/hive/src/test/java/com/dbx/agent/hive/HiveAgentTest.java @@ -1,8 +1,18 @@ package com.dbx.agent.hive; import com.dbx.agent.ConnectParams; +import com.dbx.agent.test.TestSupport; import org.junit.jupiter.api.Test; +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.Statement; +import java.util.ArrayList; +import java.util.List; + import static org.junit.jupiter.api.Assertions.assertEquals; class HiveAgentTest { @@ -39,4 +49,79 @@ class HiveAgentTest { HiveAgent.buildUrl(params) ); } + + @Test + void getTableDdlUsesHiveShowCreateTable() { + HiveAgent agent = new HiveAgent(); + List queries = new ArrayList<>(); + String expectedDdl = "CREATE TABLE `hive_test`.`cleaned_data_table`(col string)\n" + + "ROW FORMAT SERDE 'org.apache.hadoop.hive.serde2.lazy.LazySimpleSerDe'\n" + + "LOCATION 'hdfs://warehouse/cleaned_data_table'"; + ResultSet resultSet = proxy(ResultSet.class, new InvocationHandler() { + private int row = -1; + + @Override + public Object invoke(Object proxy, Method method, Object[] args) { + if ("next".equals(method.getName())) { + return ++row == 0; + } + if ("getString".equals(method.getName())) { + return expectedDdl; + } + return defaultValue(method.getReturnType()); + } + }); + Statement statement = proxy(Statement.class, (proxy, method, args) -> { + if ("executeQuery".equals(method.getName())) { + queries.add((String) args[0]); + return resultSet; + } + return defaultValue(method.getReturnType()); + }); + Connection connection = proxy(Connection.class, (proxy, method, args) -> { + if ("createStatement".equals(method.getName())) { + return statement; + } + if ("isClosed".equals(method.getName())) { + return false; + } + return defaultValue(method.getReturnType()); + }); + TestSupport.setPrivateConnection(agent, connection); + + assertEquals(expectedDdl + "\n", agent.getTableDdl("hive_test", "cleaned_data_table")); + assertEquals(List.of("SHOW CREATE TABLE `hive_test`.`cleaned_data_table`"), queries); + } + + private static T proxy(Class type, InvocationHandler handler) { + return type.cast(Proxy.newProxyInstance(type.getClassLoader(), new Class[]{type}, handler)); + } + + private static Object defaultValue(Class type) { + if (type == boolean.class) { + return false; + } + if (type == int.class) { + return 0; + } + if (type == long.class) { + return 0L; + } + if (type == float.class) { + return 0F; + } + if (type == double.class) { + return 0D; + } + if (type == byte.class) { + return (byte) 0; + } + if (type == short.class) { + return (short) 0; + } + if (type == char.class) { + return '\0'; + } + return null; + } }