perf(execution): prewarm tab session pools and report perceived wait

agents/common/src/main/java/com/dbx/agent/JdbcExecutor.java

@@ -175,6 +175,13 @@ public final class JdbcExecutor {
                     try (ResultSet rs = stmt.getResultSet()) {
                         result = readResultSet(rs, elapsed, effectiveMaxRows, valueReader);
                     }
+                    // `elapsed` stops at `stmt.execute()`. Retrieving and
+                    // converting the rows above is frequently the dominant
+                    // cost of a statement, so report the whole duration the
+                    // same way the native drivers do. Without this a
+                    // fetch-heavy result shows up in the UI as a few
+                    // milliseconds while the caller waited for seconds.
+                    result.setExecution_time_ms(System.currentTimeMillis() - start);
                 } else {
                     result = new QueryResult(
                         Collections.emptyList(),
@@ -464,7 +471,11 @@ public final class JdbcExecutor {
                 );
                 createdSession = session;
                 targetSessions.put(sessionId, session);
-                return readSessionPage(targetSessions, session, options.getPageSize(), elapsed);
+                QueryPageResult page = readSessionPage(targetSessions, session, options.getPageSize(), elapsed);
+                // Same reason as `execute`: pulling the first page's rows is
+                // part of the duration the caller waited for.
+                page.setExecution_time_ms(System.currentTimeMillis() - start);
+                return page;
             } catch (Exception e) {
                 if (createdSession != null) {
                     closeSession(targetSessions, createdSession.id);
@@ -711,8 +722,14 @@ public final class JdbcExecutor {
             throw new IllegalArgumentException(missingMessage);
         }
         synchronized (session) {
+            long fetchStart = System.currentTimeMillis();
             try {
-                return readSessionPage(targetSessions, session, pageSize, 0L);
+                QueryPageResult page = readSessionPage(targetSessions, session, pageSize, 0L);
+                // A later page reads its rows after `stmt.execute()` has long
+                // returned, so it has to measure the fetch itself instead of
+                // reporting 0ms while the caller waits.
+                page.setExecution_time_ms(System.currentTimeMillis() - fetchStart);
+                return page;
             } catch (RuntimeException | Error error) {
                 closeSession(targetSessions, sessionId);
                 throw error;

agents/common/src/test/java/com/dbx/agent/JdbcExecutorTest.java

@@ -119,6 +119,26 @@ class JdbcExecutorTest {
         assertEquals(2, fixture.getColumnTypeCalls());
     }
 
+    @Test
+    void executeCountsRowRetrievalInExecutionTime() {
+        // `stmt.execute()` returns before the rows are consumed, so a driver
+        // that only timed that call reported a handful of milliseconds for a
+        // statement whose fetch took seconds. Row retrieval has to be part of
+        // the duration the caller waited for.
+        QueryResult result = JdbcExecutor.INSTANCE.execute(
+            executionConnection(true, -1, null, new AtomicInteger(), slowResultSet(30L), null),
+            "SELECT * FROM slow_table",
+            "",
+            schema -> ""
+        );
+
+        assertEquals(Arrays.asList(Arrays.asList(1, "Ada"), Arrays.asList(2, "Grace")), result.getRows());
+        assertTrue(
+            result.getExecution_time_ms() >= 25L,
+            "row retrieval must count toward execution_time_ms but was " + result.getExecution_time_ms() + "ms"
+        );
+    }
+
     @Test
     void executeReturnsMultipleStatementWarningsForNoResultStatements() {
         SQLWarning first = new SQLWarning("identity value is 443", "S0003", 7998);
@@ -432,6 +452,69 @@ class JdbcExecutorTest {
         return rows[rowIndex][columnIndex - 1];
     }
 
+    /** Result set whose cursor advances slowly, emulating a remote fetch. */
+    private static ResultSet slowResultSet(long perRowDelayMs) {
+        Object[][] rows = {
+            {1, "Ada"},
+            {2, "Grace"}
+        };
+        String[] labels = {"id", "name"};
+        int[] sqlTypes = {Types.INTEGER, Types.VARCHAR};
+        String[] typeNames = {"INTEGER", "VARCHAR"};
+        AtomicInteger cursor = new AtomicInteger(-1);
+
+        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":
+                    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":
+                    return metadata;
+                case "next":
+                    int nextIndex = cursor.incrementAndGet();
+                    if (nextIndex >= rows.length) {
+                        return false;
+                    }
+                    Thread.sleep(perRowDelayMs);
+                    return true;
+                case "getObject":
+                case "getInt":
+                case "getString":
+                    Object cell = currentCell(rows, cursor.get(), (Integer) args[0]);
+                    if (method.getName().equals("getInt")) {
+                        return ((Number) cell).intValue();
+                    }
+                    return method.getName().equals("getObject") ? cell : String.valueOf(cell);
+                case "wasNull":
+                    return false;
+                default:
+                    return defaultValue(method.getReturnType());
+            }
+        };
+        return (ResultSet) Proxy.newProxyInstance(
+            ResultSet.class.getClassLoader(),
+            new Class<?>[]{ResultSet.class},
+            resultSetHandler
+        );
+    }
+
     private static ResultSet resultSet(Object objectValue, StringSupplier stringSupplier, boolean wasNull) {
         byte[] bytesValue = objectValue instanceof byte[] ? (byte[]) objectValue : null;
         return resultSet(objectValue, bytesValue, stringSupplier, wasNull, null);