diff --git a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java index 0da63753db..0b11f87d85 100644 --- a/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java +++ b/agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java @@ -113,24 +113,28 @@ public final class JdbcExecutor { return unchecked(() -> { ResultSetMetaData meta = rs.getMetaData(); int colCount = meta.getColumnCount(); - List columns = new ArrayList<>(); - List columnTypes = new ArrayList<>(); + List columns = new ArrayList<>(colCount); + List columnTypes = new ArrayList<>(colCount); + // Cache JDBC column metadata once; some drivers resolve it lazily, + // and row reading is the hot path. + int[] sqlTypeByIndex = new int[colCount]; String[] typeNameByIndex = new String[colCount]; for (int i = 1; i <= colCount; i++) { columns.add(meta.getColumnLabel(i)); + sqlTypeByIndex[i - 1] = safeColumnSqlType(meta, i); String typeName = safeColumnTypeName(meta, i); columnTypes.add(typeName); typeNameByIndex[i - 1] = typeName; } - List> rows = new ArrayList<>(); + List> rows = new ArrayList<>(initialRowCapacity(maxRows)); boolean truncated = false; while (rs.next()) { if (rows.size() >= maxRows) { truncated = true; break; } - rows.add(rowValues(rs, valueReader, typeNameByIndex)); + rows.add(rowValues(rs, valueReader, sqlTypeByIndex, typeNameByIndex)); } return new QueryResult(columns, columnTypes, rows, 0L, executionTimeMs, truncated); @@ -221,9 +225,12 @@ public final class JdbcExecutor { int colCount = meta.getColumnCount(); List columns = new ArrayList<>(colCount); List columnTypes = new ArrayList<>(colCount); + // Cache JDBC column metadata once; fetching a page should only read row values. + int[] sqlTypeByIndex = new int[colCount]; String[] typeNameByIndex = new String[colCount]; for (int i = 1; i <= colCount; i++) { columns.add(meta.getColumnLabel(i)); + sqlTypeByIndex[i - 1] = safeColumnSqlType(meta, i); String typeName = safeColumnTypeName(meta, i); columnTypes.add(typeName); typeNameByIndex[i - 1] = typeName; @@ -234,6 +241,7 @@ public final class JdbcExecutor { rs, columns, columnTypes, + sqlTypeByIndex, typeNameByIndex, Math.max(options.getMaxRows(), 1), valueReader @@ -449,7 +457,7 @@ public final class JdbcExecutor { return unchecked(() -> { session.lastAccessedAtMillis = System.currentTimeMillis(); int effectivePageSize = Math.max(pageSize, 1); - List> rows = new ArrayList<>(); + List> rows = new ArrayList<>(initialRowCapacity(effectivePageSize)); if (session.pendingRow != null) { rows.add(session.pendingRow); @@ -461,7 +469,7 @@ public final class JdbcExecutor { closeSession(targetSessions, session.id); return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, null, false); } - rows.add(rowValues(session.resultSet, session.valueReader, session.typeNameByIndex)); + rows.add(rowValues(session.resultSet, session.valueReader, session.sqlTypeByIndex, session.typeNameByIndex)); session.rowsRead += 1; } @@ -477,7 +485,7 @@ public final class JdbcExecutor { return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, null, false); } - session.pendingRow = rowValues(session.resultSet, session.valueReader, session.typeNameByIndex); + session.pendingRow = rowValues(session.resultSet, session.valueReader, session.sqlTypeByIndex, session.typeNameByIndex); session.rowsRead += 1; return new QueryPageResult(session.columns, session.columnTypes, rows, 0L, executionTimeMs, false, session.id, true); }); @@ -517,17 +525,29 @@ public final class JdbcExecutor { } } - private List rowValues(ResultSet rs, ResultValueReader valueReader, String[] typeNameByIndex) throws SQLException { - ResultSetMetaData meta = rs.getMetaData(); - int colCount = meta.getColumnCount(); + static int safeColumnSqlType(ResultSetMetaData meta, int columnIndex) { + try { + return meta.getColumnType(columnIndex); + } catch (SQLException ignored) { + return Types.OTHER; + } + } + + private List rowValues( + ResultSet rs, + ResultValueReader valueReader, + int[] sqlTypeByIndex, + String[] typeNameByIndex + ) throws SQLException { + int colCount = sqlTypeByIndex.length; List row = new ArrayList<>(colCount); for (int i = 1; i <= colCount; i++) { - int sqlType = meta.getColumnType(i); + int sqlType = sqlTypeByIndex[i - 1]; Object value; if (valueReader instanceof ColumnAwareResultValueReader) { String typeName = typeNameByIndex != null && i - 1 < typeNameByIndex.length ? typeNameByIndex[i - 1] - : safeColumnTypeName(meta, i); + : ""; value = ((ColumnAwareResultValueReader) valueReader).read(rs, i, sqlType, typeName); } else { value = valueReader.read(rs, i, sqlType); @@ -537,6 +557,13 @@ public final class JdbcExecutor { return row; } + private static int initialRowCapacity(int requestedRows) { + if (requestedRows <= 0) { + return 0; + } + return Math.min(requestedRows, 1024); + } + private void applySchema(Connection conn, String schema, Function setSchemaSql) throws SQLException { if (schema == null || schema.trim().isEmpty()) { return; @@ -664,6 +691,7 @@ public final class JdbcExecutor { private final ResultSet resultSet; private final List columns; private final List columnTypes; + private final int[] sqlTypeByIndex; private final String[] typeNameByIndex; private final int maxRows; private final ResultValueReader valueReader; @@ -677,6 +705,7 @@ public final class JdbcExecutor { ResultSet resultSet, List columns, List columnTypes, + int[] sqlTypeByIndex, String[] typeNameByIndex, int maxRows, ResultValueReader valueReader @@ -686,6 +715,7 @@ public final class JdbcExecutor { this.resultSet = resultSet; this.columns = columns; this.columnTypes = columnTypes; + this.sqlTypeByIndex = sqlTypeByIndex; this.typeNameByIndex = typeNameByIndex; this.maxRows = maxRows; this.valueReader = valueReader; diff --git a/agents/common/src/main/java/com/dbx/agent/QueryResult.java b/agents/common/src/main/java/com/dbx/agent/QueryResult.java index 172966e5a6..c38881aa77 100644 --- a/agents/common/src/main/java/com/dbx/agent/QueryResult.java +++ b/agents/common/src/main/java/com/dbx/agent/QueryResult.java @@ -99,7 +99,7 @@ public final class QueryResult { if (input == null) { return Collections.emptyList(); } - List> normalized = new ArrayList<>(); + List> normalized = new ArrayList<>(input.size()); for (List row : input) { normalized.add(row == null ? Collections.emptyList() : new ArrayList(row)); } diff --git a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java index 9acee273bc..f6c125523f 100644 --- a/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java +++ b/agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java @@ -6,7 +6,10 @@ import java.lang.reflect.InvocationHandler; import java.lang.reflect.Method; import java.lang.reflect.Proxy; import java.sql.ResultSet; +import java.sql.ResultSetMetaData; import java.sql.Types; +import java.util.Arrays; +import java.util.concurrent.atomic.AtomicInteger; import javax.sql.rowset.serial.SerialBlob; import static org.junit.jupiter.api.Assertions.assertEquals; @@ -38,10 +41,90 @@ class JdbcExecutorTest { assertEquals("0x0a0b", JdbcExecutor.INSTANCE.defaultResultValue(rs, 1, Types.OTHER)); } + @Test + void readResultSetCachesColumnTypeMetadataAcrossRows() { + CountingResultSetFixture fixture = countingResultSet(new Object[][]{ + {1, "Ada"}, + {2, "Grace"} + }); + + QueryResult result = JdbcExecutor.INSTANCE.readResultSet( + fixture.resultSet(), + 12L, + 10, + JdbcExecutor.INSTANCE::defaultResultValue + ); + + assertEquals(Arrays.asList("id", "name"), result.getColumns()); + assertEquals(Arrays.asList("INTEGER", "VARCHAR"), result.getColumn_types()); + assertEquals(Arrays.asList(Arrays.asList(1, "Ada"), Arrays.asList(2, "Grace")), result.getRows()); + assertEquals(1, fixture.getMetaDataCalls()); + assertEquals(2, fixture.getColumnTypeCalls()); + } + private static ResultSet resultSet(byte[] bytes, StringSupplier stringSupplier) { return resultSet(bytes, stringSupplier, false); } + private static CountingResultSetFixture countingResultSet(Object[][] rows) { + String[] labels = {"id", "name"}; + int[] sqlTypes = {Types.INTEGER, Types.VARCHAR}; + String[] typeNames = {"INTEGER", "VARCHAR"}; + AtomicInteger cursor = new AtomicInteger(-1); + AtomicInteger getMetaDataCalls = new AtomicInteger(); + AtomicInteger getColumnTypeCalls = new AtomicInteger(); + + InvocationHandler metaHandler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "getColumnCount": + return labels.length; + case "getColumnLabel": + return labels[(Integer) args[0] - 1]; + case "getColumnType": + getColumnTypeCalls.incrementAndGet(); + return sqlTypes[(Integer) args[0] - 1]; + case "getColumnTypeName": + return typeNames[(Integer) args[0] - 1]; + default: + return defaultValue(method.getReturnType()); + } + }; + ResultSetMetaData metadata = (ResultSetMetaData) Proxy.newProxyInstance( + ResultSetMetaData.class.getClassLoader(), + new Class[]{ResultSetMetaData.class}, + metaHandler + ); + + InvocationHandler resultSetHandler = (Object unused, Method method, Object[] args) -> { + switch (method.getName()) { + case "getMetaData": + getMetaDataCalls.incrementAndGet(); + return metadata; + case "next": + return cursor.incrementAndGet() < rows.length; + case "getInt": + return ((Number) currentCell(rows, cursor.get(), (Integer) args[0])).intValue(); + case "getString": + Object value = currentCell(rows, cursor.get(), (Integer) args[0]); + return value == null ? null : value.toString(); + case "wasNull": + return false; + default: + return defaultValue(method.getReturnType()); + } + }; + ResultSet resultSet = (ResultSet) Proxy.newProxyInstance( + ResultSet.class.getClassLoader(), + new Class[]{ResultSet.class}, + resultSetHandler + ); + return new CountingResultSetFixture(resultSet, getMetaDataCalls, getColumnTypeCalls); + } + + private static Object currentCell(Object[][] rows, int rowIndex, int columnIndex) { + return rows[rowIndex][columnIndex - 1]; + } + private static ResultSet resultSet(Object objectValue, StringSupplier stringSupplier, boolean wasNull) { InvocationHandler handler = (Object unused, Method method, Object[] args) -> { switch (method.getName()) { @@ -95,4 +178,32 @@ class JdbcExecutorTest { private interface StringSupplier { String get() throws Exception; } + + private static final class CountingResultSetFixture { + private final ResultSet resultSet; + private final AtomicInteger getMetaDataCalls; + private final AtomicInteger getColumnTypeCalls; + + private CountingResultSetFixture( + ResultSet resultSet, + AtomicInteger getMetaDataCalls, + AtomicInteger getColumnTypeCalls + ) { + this.resultSet = resultSet; + this.getMetaDataCalls = getMetaDataCalls; + this.getColumnTypeCalls = getColumnTypeCalls; + } + + private ResultSet resultSet() { + return resultSet; + } + + private int getMetaDataCalls() { + return getMetaDataCalls.get(); + } + + private int getColumnTypeCalls() { + return getColumnTypeCalls.get(); + } + } }