feat(metadata): optimize agent metadata loading
This commit is contained in:
t8y2 committed
1 parent
8883a6b3b7
commit
e74782f651
65 files changed
+6192
-571
No files matched your search
@@ -3,11 +3,13 @@ package com.dbx.agent.databend;
|
||||
import com.dbx.agent.ConfiguredJdbcAgent;
|
||||
import com.dbx.agent.JdbcAgentProfile;
|
||||
import com.dbx.agent.JsonRpcServer;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.ObjectSource;
|
||||
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.SQLException;
|
||||
import java.sql.PreparedStatement;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
@@ -54,6 +56,23 @@ public final class DatabendAgent extends ConfiguredJdbcAgent {
|
||||
});
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<ObjectInfo> listObjects(String schema, MetadataListConstraints constraints) {
|
||||
MetadataListConstraints normalized = MetadataListConstraints.orNone(constraints);
|
||||
boolean includeTables = normalized.includesTableLikeTypes();
|
||||
boolean includeProcedures = includesProcedures(normalized);
|
||||
if (includeProcedures && !includeTables) {
|
||||
return listProcedures(schema, normalized);
|
||||
}
|
||||
if (includeTables && !includeProcedures) {
|
||||
return super.listObjects(schema, normalized);
|
||||
}
|
||||
if (!includeTables && !includeProcedures) {
|
||||
return List.of();
|
||||
}
|
||||
return normalized.filterObjects(listObjects(schema));
|
||||
}
|
||||
|
||||
@Override
|
||||
public ObjectSource getObjectSource(String schema, String name, String objectType) {
|
||||
if (!"PROCEDURE".equalsIgnoreCase(objectType)) {
|
||||
@@ -80,6 +99,53 @@ public final class DatabendAgent extends ConfiguredJdbcAgent {
|
||||
}
|
||||
}
|
||||
|
||||
private List<ObjectInfo> listProcedures(String schema, MetadataListConstraints constraints) {
|
||||
return unchecked(() -> {
|
||||
useSchema(schema);
|
||||
List<ObjectInfo> result = new ArrayList<>();
|
||||
StringBuilder sql = new StringBuilder("SELECT name, comment FROM system.procedures");
|
||||
List<Object> args = new ArrayList<>();
|
||||
if (constraints.hasFilter()) {
|
||||
sql.append(" WHERE UPPER(name) LIKE ? ESCAPE '\\\\'");
|
||||
args.add(constraints.fuzzyLikePattern().toUpperCase(Locale.ROOT));
|
||||
}
|
||||
sql.append(" ORDER BY name");
|
||||
if (constraints.hasLimit()) {
|
||||
sql.append(" LIMIT ?");
|
||||
args.add(constraints.getLimit());
|
||||
if (constraints.hasOffset()) {
|
||||
sql.append(" OFFSET ?");
|
||||
args.add(constraints.getOffset());
|
||||
}
|
||||
}
|
||||
try (PreparedStatement stmt = requireConnected().prepareStatement(sql.toString())) {
|
||||
bind(stmt, args);
|
||||
try (ResultSet rs = stmt.executeQuery()) {
|
||||
while (rs.next()) {
|
||||
result.add(new ObjectInfo(rs.getString("name"), "PROCEDURE", schema, rs.getString("comment")));
|
||||
}
|
||||
}
|
||||
}
|
||||
MetadataListConstraints guard = constraints.hasLimit() ? constraints.withoutPaging() : constraints;
|
||||
return guard.filterObjects(result);
|
||||
});
|
||||
}
|
||||
|
||||
private static boolean includesProcedures(MetadataListConstraints constraints) {
|
||||
return !constraints.hasObjectTypes() || constraints.objectTypeAllowed("PROCEDURE");
|
||||
}
|
||||
|
||||
private static void bind(PreparedStatement stmt, List<Object> args) throws SQLException {
|
||||
for (int index = 0; index < args.size(); index += 1) {
|
||||
Object arg = args.get(index);
|
||||
if (arg instanceof Integer) {
|
||||
stmt.setInt(index + 1, (Integer) arg);
|
||||
} else {
|
||||
stmt.setString(index + 1, String.valueOf(arg));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private ProcedureMetadata findProcedure(String name) throws SQLException {
|
||||
try (java.sql.PreparedStatement stmt = requireConnected().prepareStatement(
|
||||
"SELECT arguments, comment FROM system.procedures WHERE name = ? ORDER BY procedure_id LIMIT 1"
|
||||
|
||||
@@ -2,10 +2,24 @@ package com.dbx.agent.databend;
|
||||
|
||||
import com.dbx.agent.ConnectParams;
|
||||
import com.dbx.agent.DatabaseAgent;
|
||||
import com.dbx.agent.MetadataListConstraints;
|
||||
import com.dbx.agent.ObjectInfo;
|
||||
import com.dbx.agent.test.JdbcFakeExecutionBehaviorTest;
|
||||
import com.dbx.agent.test.TestSupport;
|
||||
import org.junit.jupiter.api.Assertions;
|
||||
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.DatabaseMetaData;
|
||||
import java.sql.PreparedStatement;
|
||||
import java.sql.ResultSet;
|
||||
import java.sql.Statement;
|
||||
import java.util.ArrayList;
|
||||
import java.util.List;
|
||||
|
||||
class DatabendAgentTest extends JdbcFakeExecutionBehaviorTest {
|
||||
@Override
|
||||
protected DatabaseAgent createAgent() {
|
||||
@@ -51,4 +65,132 @@ class DatabendAgentTest extends JdbcFakeExecutionBehaviorTest {
|
||||
DatabendAgent.DATABEND_PROFILE.schemaSwitchSql("analytics\"prod")
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
void constrainedListObjectsKeepsDatabendProcedures() {
|
||||
List<String> sql = new ArrayList<>();
|
||||
DatabendAgent agent = new DatabendAgent();
|
||||
TestSupport.setPrivateConnection(agent, databendMetadataConnection(sql, resultSet(
|
||||
new String[]{"name", "comment"},
|
||||
new Object[][]{
|
||||
{"refresh_orders", "proc comment"},
|
||||
{"format_order", null}
|
||||
}
|
||||
)));
|
||||
|
||||
List<ObjectInfo> objects = agent.listObjects(
|
||||
"default",
|
||||
new MetadataListConstraints("format", 1, null, List.of("PROCEDURE"))
|
||||
);
|
||||
|
||||
Assertions.assertEquals(1, objects.size());
|
||||
Assertions.assertEquals("format_order", objects.get(0).getName());
|
||||
Assertions.assertEquals("PROCEDURE", objects.get(0).getObject_type());
|
||||
Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("FROM system.procedures")), sql.toString());
|
||||
Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("UPPER(name) LIKE ?")), sql.toString());
|
||||
Assertions.assertTrue(sql.stream().anyMatch(value -> value.contains("LIMIT ?")), sql.toString());
|
||||
}
|
||||
|
||||
private static Connection databendMetadataConnection(List<String> sql, ResultSet procedures) {
|
||||
DatabaseMetaData metadata = proxy(DatabaseMetaData.class, (method, args) -> {
|
||||
if ("getTableTypes".equals(method.getName()) || "getTables".equals(method.getName())) {
|
||||
return emptyResultSet();
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
Statement statement = proxy(Statement.class, (method, args) -> {
|
||||
if ("execute".equals(method.getName()) || "close".equals(method.getName())) {
|
||||
return "execute".equals(method.getName()) ? false : null;
|
||||
}
|
||||
if ("executeQuery".equals(method.getName())) {
|
||||
return procedures;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
PreparedStatement preparedStatement = proxy(PreparedStatement.class, (method, args) -> {
|
||||
if ("executeQuery".equals(method.getName())) {
|
||||
return procedures;
|
||||
}
|
||||
if ("setString".equals(method.getName()) || "setInt".equals(method.getName()) || "close".equals(method.getName())) {
|
||||
return null;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
return proxy(Connection.class, (method, args) -> {
|
||||
if ("getMetaData".equals(method.getName())) {
|
||||
return metadata;
|
||||
}
|
||||
if ("createStatement".equals(method.getName())) {
|
||||
return statement;
|
||||
}
|
||||
if ("prepareStatement".equals(method.getName())) {
|
||||
sql.add(String.valueOf(args[0]));
|
||||
return preparedStatement;
|
||||
}
|
||||
if ("isClosed".equals(method.getName())) {
|
||||
return false;
|
||||
}
|
||||
return defaultValue(method.getReturnType());
|
||||
});
|
||||
}
|
||||
|
||||
private static ResultSet emptyResultSet() {
|
||||
return resultSet(new String[]{"VALUE"}, new Object[][]{});
|
||||
}
|
||||
|
||||
private static ResultSet resultSet(String[] columns, Object[][] rows) {
|
||||
int[] index = {-1};
|
||||
return proxy(ResultSet.class, (method, args) -> {
|
||||
switch (method.getName()) {
|
||||
case "next":
|
||||
index[0] += 1;
|
||||
return index[0] < rows.length;
|
||||
case "getString":
|
||||
Object value = columnValue(columns, rows[index[0]], args[0]);
|
||||
return value == null ? null : String.valueOf(value);
|
||||
case "close":
|
||||
return null;
|
||||
default:
|
||||
return defaultValue(method.getReturnType());
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private static Object columnValue(String[] columns, Object[] row, Object key) {
|
||||
if (key instanceof Number) {
|
||||
return row[((Number) key).intValue() - 1];
|
||||
}
|
||||
for (int i = 0; i < columns.length; i++) {
|
||||
if (columns[i].equalsIgnoreCase(String.valueOf(key))) {
|
||||
return row[i];
|
||||
}
|
||||
}
|
||||
return null;
|
||||
}
|
||||
|
||||
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;
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user