Read test query results before reporting JDBC container as ready

modules/jdbc/src/main/java/org/testcontainers/containers/JdbcDatabaseContainer.java

@@ -15,7 +15,10 @@
 
 import java.sql.Connection;
 import java.sql.Driver;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.sql.SQLTimeoutException;
 import java.sql.Statement;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -190,8 +193,28 @@ protected void waitUntilContainerStarted() {
                 Thread.sleep(100L);
             } else {
                 try (Connection connection = createConnection(""); Statement statement = connection.createStatement()) {
+                    long remainingNanos = TimeUnit.SECONDS.toNanos(startupTimeoutSeconds) - (System.nanoTime() - start);
+                    if (remainingNanos <= 0) {
+                        throw new SQLTimeoutException("Startup timeout exceeded before executing the test query");
+                    }
+                    try {
+                        // JDBC uses whole seconds; zero would disable the timeout.
+                        statement.setQueryTimeout((int) Math.max(1, TimeUnit.NANOSECONDS.toSeconds(remainingNanos)));
+                    } catch (SQLFeatureNotSupportedException e) {
+                        logger().debug("JDBC driver does not support a query timeout", e);
+                    }
                     boolean testQuerySucceeded = statement.execute(this.getTestQueryString());
                     if (testQuerySucceeded) {
+                        // Some databases (e.g. Trino) accept the query but only fail once its results are fetched,
+                        // so the database is only considered ready after the results have been read successfully
+                        try (ResultSet resultSet = statement.getResultSet()) {
+                            if (resultSet != null) {
+                                while (resultSet.next()) {
+                                    checkStartupTimeout(start);
+                                }
+                            }
+                        }
+                        checkStartupTimeout(start);
                         return;
                     }
                 } catch (NoDriverFoundException e) {
@@ -215,6 +238,12 @@ protected void waitUntilContainerStarted() {
         );
     }
 
+    private void checkStartupTimeout(long start) throws SQLTimeoutException {
+        if (System.nanoTime() - start >= TimeUnit.SECONDS.toNanos(startupTimeoutSeconds)) {
+            throw new SQLTimeoutException("Startup timeout exceeded while reading the test query results");
+        }
+    }
+
     @Override
     protected void containerIsStarted(InspectContainerResponse containerInfo) {
         logger().info("Container is started (JDBC URL: {})", this.getJdbcUrl());

modules/jdbc/src/test/java/org/testcontainers/containers/JdbcDatabaseContainerTest.java

@@ -5,10 +5,20 @@
 import org.slf4j.Logger;
 
 import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.sql.SQLTimeoutException;
+import java.sql.Statement;
+import java.util.concurrent.atomic.AtomicInteger;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatExceptionOfType;
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 class JdbcDatabaseContainerTest {
 
@@ -20,6 +30,129 @@ void anExceptionIsThrownIfJdbcIsNotAvailable() {
         assertThatExceptionOfType(IllegalStateException.class).isThrownBy(jdbcContainer::waitUntilContainerStarted);
     }
 
+    @Test
+    void testQueryIsRetriedIfReadingItsResultsFails() {
+        // Some databases (e.g. Trino) accept the test query but only fail once its results are fetched
+        FailingResultSetJdbcDatabaseContainerStub jdbcContainer = new FailingResultSetJdbcDatabaseContainerStub();
+        jdbcContainer.withStartupTimeoutSeconds(5);
+
+        jdbcContainer.waitUntilContainerStarted();
+
+        assertThat(jdbcContainer.connectionAttempts).hasValue(2);
+    }
+
+    @Test
+    void testQueryTimeoutUsesRemainingStartupTime() throws SQLException {
+        Statement statement = mock(Statement.class);
+        when(statement.execute("SELECT 1")).thenReturn(true);
+        JdbcDatabaseContainer<?> jdbcContainer = containerWithStatement(statement).withStartupTimeoutSeconds(2);
+
+        jdbcContainer.waitUntilContainerStarted();
+
+        verify(statement).setQueryTimeout(1);
+    }
+
+    @Test
+    void unsupportedQueryTimeoutDoesNotPreventStartup() throws SQLException {
+        Statement statement = mock(Statement.class);
+        doThrow(new SQLFeatureNotSupportedException()).when(statement).setQueryTimeout(anyInt());
+        when(statement.execute("SELECT 1")).thenReturn(true);
+
+        containerWithStatement(statement).waitUntilContainerStarted();
+
+        verify(statement).execute("SELECT 1");
+    }
+
+    @Test
+    void resultsCompletedAfterStartupTimeoutDoNotReportReady() throws SQLException {
+        Statement statement = mock(Statement.class);
+        when(statement.execute("SELECT 1")).thenReturn(true);
+        ResultSet resultSet = mock(ResultSet.class);
+        when(statement.getResultSet()).thenReturn(resultSet);
+        when(resultSet.next())
+            .thenAnswer(invocation -> {
+                Thread.sleep(1100);
+                return false;
+            });
+        JdbcDatabaseContainer<?> jdbcContainer = containerWithStatement(statement).withStartupTimeoutSeconds(1);
+
+        assertThatExceptionOfType(IllegalStateException.class)
+            .isThrownBy(jdbcContainer::waitUntilContainerStarted)
+            .withCauseInstanceOf(SQLTimeoutException.class);
+
+        verify(resultSet).close();
+        verify(statement).close();
+    }
+
+    @Test
+    void resultDrainingStopsAfterStartupTimeout() throws SQLException {
+        Statement statement = mock(Statement.class);
+        when(statement.execute("SELECT 1")).thenReturn(true);
+        ResultSet resultSet = mock(ResultSet.class);
+        when(statement.getResultSet()).thenReturn(resultSet);
+        when(resultSet.next())
+            .thenAnswer(invocation -> {
+                Thread.sleep(1100);
+                return true;
+            })
+            .thenReturn(false);
+        JdbcDatabaseContainer<?> jdbcContainer = containerWithStatement(statement).withStartupTimeoutSeconds(1);
+
+        assertThatExceptionOfType(IllegalStateException.class)
+            .isThrownBy(jdbcContainer::waitUntilContainerStarted)
+            .withCauseInstanceOf(SQLTimeoutException.class);
+
+        verify(resultSet).next();
+        verify(resultSet).close();
+        verify(statement).close();
+    }
+
+    private JdbcDatabaseContainer<?> containerWithStatement(Statement statement) throws SQLException {
+        Connection connection = mock(Connection.class);
+        when(connection.createStatement()).thenReturn(statement);
+        return new JdbcDatabaseContainerStub("mysql:latest") {
+            @Override
+            protected String getTestQueryString() {
+                return "SELECT 1";
+            }
+
+            @Override
+            public Connection createConnection(String queryString) {
+                return connection;
+            }
+        };
+    }
+
+    static class FailingResultSetJdbcDatabaseContainerStub extends JdbcDatabaseContainerStub {
+
+        private final AtomicInteger connectionAttempts = new AtomicInteger();
+
+        FailingResultSetJdbcDatabaseContainerStub() {
+            super("mysql:latest");
+        }
+
+        @Override
+        protected String getTestQueryString() {
+            return "SELECT 1";
+        }
+
+        @Override
+        public Connection createConnection(String queryString) throws SQLException, NoDriverFoundException {
+            ResultSet resultSet = mock(ResultSet.class);
+            if (connectionAttempts.incrementAndGet() == 1) {
+                when(resultSet.next()).thenThrow(new SQLException("No nodes available to run query"));
+            } else {
+                when(resultSet.next()).thenReturn(true, false);
+            }
+            Statement statement = mock(Statement.class);
+            when(statement.execute("SELECT 1")).thenReturn(true);
+            when(statement.getResultSet()).thenReturn(resultSet);
+            Connection connection = mock(Connection.class);
+            when(connection.createStatement()).thenReturn(statement);
+            return connection;
+        }
+    }
+
     static class JdbcDatabaseContainerStub extends JdbcDatabaseContainer {
 
         public JdbcDatabaseContainerStub(@NonNull String dockerImageName) {