[yugabytedb] Change SQL wait strategy (#7784)

Co-authored-by: Eddú Meléndez <[email protected]>

Author: HarshDaryani896

modules/yugabytedb/src/main/java/org/testcontainers/containers/YugabyteDBYSQLContainer.java

@@ -154,4 +154,9 @@ public YugabyteDBYSQLContainer withPassword(final String password) {
         this.password = password;
         return this;
     }
+
+    @Override
+    protected void waitUntilContainerStarted() {
+        getWaitStrategy().waitUntilReady(this);
+    }
 }

modules/yugabytedb/src/main/java/org/testcontainers/containers/delegate/YugabyteDBYCQLDelegate.java

@@ -33,9 +33,20 @@ public void execute(
         boolean continueOnError,
         boolean ignoreFailedDrops
     ) {
+        final String containerInterfaceIP = container
+            .getContainerInfo()
+            .getNetworkSettings()
+            .getNetworks()
+            .entrySet()
+            .stream()
+            .findFirst()
+            .get()
+            .getValue()
+            .getIpAddress();
         try {
             ExecResult result = container.execInContainer(
                 BIN_PATH,
+                containerInterfaceIP,
                 "-u",
                 container.getUsername(),
                 "-p",

modules/yugabytedb/src/main/java/org/testcontainers/containers/strategy/YugabyteDBYCQLWaitStrategy.java

@@ -37,6 +37,16 @@ public final class YugabyteDBYCQLWaitStrategy extends AbstractWaitStrategy {
     public void waitUntilReady(WaitStrategyTarget target) {
         YugabyteDBYCQLContainer container = (YugabyteDBYCQLContainer) target;
         AtomicBoolean status = new AtomicBoolean(true);
+        final String containerInterfaceIP = container
+            .getContainerInfo()
+            .getNetworkSettings()
+            .getNetworks()
+            .entrySet()
+            .stream()
+            .findFirst()
+            .get()
+            .getValue()
+            .getIpAddress();
         retryUntilSuccess(
             (int) startupTimeout.getSeconds(),
             TimeUnit.SECONDS,
@@ -46,6 +56,7 @@ public void waitUntilReady(WaitStrategyTarget target) {
                         try {
                             ExecResult result = container.execInContainer(
                                 BIN_PATH,
+                                containerInterfaceIP,
                                 "-u",
                                 container.getUsername(),
                                 "-p",

modules/yugabytedb/src/main/java/org/testcontainers/containers/strategy/YugabyteDBYSQLWaitStrategy.java

@@ -8,6 +8,7 @@
 
 import java.sql.Connection;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.concurrent.TimeUnit;
 
 import static org.rnorth.ducttape.unreliables.Unreliables.retryUntilSuccess;
@@ -27,10 +28,13 @@
 @Slf4j
 public final class YugabyteDBYSQLWaitStrategy extends AbstractWaitStrategy {
 
-    private static final String YSQL_TEST_QUERY = "SELECT 1";
-
     private final WaitStrategyTarget target;
 
+    private static final String YSQL_EXTENDED_PROBE =
+        "CREATE TABLE IF NOT EXISTS YB_SAMPLE(k int, v int, primary key(k, v))";
+
+    private static final String YSQL_EXTENDED_PROBE_DROP_TABLE = "DROP TABLE IF EXISTS YB_SAMPLE";
+
     @Override
     public void waitUntilReady(WaitStrategyTarget target) {
         YugabyteDBYSQLContainer container = (YugabyteDBYSQLContainer) target;
@@ -40,10 +44,11 @@ public void waitUntilReady(WaitStrategyTarget target) {
             () -> {
                 getRateLimiter()
                     .doWhenReady(() -> {
-                        try (Connection con = container.createConnection(container.getJdbcUrl())) {
-                            con.createStatement().execute(YSQL_TEST_QUERY);
+                        try (Connection con = container.createConnection(""); Statement stmt = con.createStatement()) {
+                            stmt.execute(YSQL_EXTENDED_PROBE);
+                            stmt.execute(YSQL_EXTENDED_PROBE_DROP_TABLE);
                         } catch (SQLException ex) {
-                            log.error("Error connecting to the database", ex);
+                            throw new RuntimeException(ex);
                         }
                     });
                 return true;

modules/yugabytedb/src/test/java/org/testcontainers/junit/yugabytedb/YugabyteDBYCQLTest.java

@@ -15,6 +15,8 @@ public class YugabyteDBYCQLTest {
 
     private static final String IMAGE_NAME = "yugabytedb/yugabyte:2.14.4.0-b26";
 
+    private static final String IMAGE_NAME_2_18 = "yugabytedb/yugabyte:2.18.3.0-b75";
+
     private static final DockerImageName YBDB_TEST_IMAGE = DockerImageName.parse(IMAGE_NAME);
 
     @Test
@@ -106,6 +108,16 @@ public void testInitScript() {
         }
     }
 
+    @Test
+    public void shouldStartWhenContainerIpIsUsedInWaitStrategy() {
+        try (final YugabyteDBYCQLContainer ycqlContainer = new YugabyteDBYCQLContainer(IMAGE_NAME_2_18)) {
+            ycqlContainer.start();
+            boolean isQueryExecuted = performQuery(ycqlContainer, "SELECT release_version FROM system.local")
+                .wasApplied();
+            assertThat(isQueryExecuted).isTrue();
+        }
+    }
+
     private ResultSet performQuery(YugabyteDBYCQLContainer ycqlContainer, String cql) {
         try (
             CqlSession session = CqlSession

modules/yugabytedb/src/test/java/org/testcontainers/junit/yugabytedb/YugabyteDBYSQLTest.java

@@ -95,4 +95,20 @@ public void testWithCustomRole() throws SQLException {
                 .isEqualTo(1);
         }
     }
+
+    @Test
+    public void testWaitStrategy() throws SQLException {
+        try (final YugabyteDBYSQLContainer ysqlContainer = new YugabyteDBYSQLContainer(YBDB_TEST_IMAGE)) {
+            ysqlContainer.start();
+            assertThat(performQuery(ysqlContainer, "SELECT 1").getInt(1))
+                .as("A sample test query succeeds")
+                .isEqualTo(1);
+            boolean tableExists = performQuery(
+                ysqlContainer,
+                "SELECT EXISTS (SELECT FROM pg_tables WHERE tablename = 'YB_SAMPLE')"
+            )
+                .getBoolean(1);
+            assertThat(tableExists).as("yb_sample table does not exists").isFalse();
+        }
+    }
 }