Repository: crate/crate
Pull request #17600: Re-implement ensureGreen/Yellow using sys tables
Author: login: mfussenegger
Labels: name: cla-signed
Created: 2025-03-12T15:12:10Z | Merged: 2025-03-20T15:01:18Z



Diff:
diff --git a/server/src/test/java/io/crate/integrationtests/GroupByAggregateTest.java b/server/src/test/java/io/crate/integrationtests/GroupByAggregateTest.java
index f74365ce1d99..d727552a6bee 100644
--- a/server/src/test/java/io/crate/integrationtests/GroupByAggregateTest.java
+++ b/server/src/test/java/io/crate/integrationtests/GroupByAggregateTest.java
@@ -50,7 +50,7 @@ public class GroupByAggregateTest extends IntegTestCase {
     private final Setup setup = new Setup(sqlExecutor);
 
     @Before
-    public void initTestData() {
+    public void initTestData() throws Exception {
         setup.setUpEmployees();
     }
 
diff --git a/server/src/test/java/io/crate/integrationtests/SubSelectIntegrationTest.java b/server/src/test/java/io/crate/integrationtests/SubSelectIntegrationTest.java
index b2970564ea20..c5e9068d1931 100644
--- a/server/src/test/java/io/crate/integrationtests/SubSelectIntegrationTest.java
+++ b/server/src/test/java/io/crate/integrationtests/SubSelectIntegrationTest.java
@@ -782,7 +782,7 @@ public void test_quoted_subscript_with_sub_select() {
     }
 
     @Test
-    public void test_non_recursive_with_query() {
+    public void test_non_recursive_with_query() throws Exception {
         setup.setUpCharacters();
 
         execute("WITH ch AS (SELECT * FROM characters WHERE female = true) " +
diff --git a/server/src/test/java/io/crate/integrationtests/VersionHandlingIntegrationTest.java b/server/src/test/java/io/crate/integrationtests/VersionHandlingIntegrationTest.java
index a9cd5d59c845..06013c64b1fa 100644
--- a/server/src/test/java/io/crate/integrationtests/VersionHandlingIntegrationTest.java
+++ b/server/src/test/java/io/crate/integrationtests/VersionHandlingIntegrationTest.java
@@ -24,8 +24,7 @@
 import static io.crate.protocols.postgres.PGErrorStatus.INTERNAL_ERROR;
 import static io.crate.testing.Asserts.assertThat;
 import static io.netty.handler.codec.http.HttpResponseStatus.BAD_REQUEST;
-
-import java.io.IOException;
+import static org.assertj.core.api.Assertions.assertThat;
 
 import org.elasticsearch.test.IntegTestCase;
 import org.junit.Test;
@@ -39,7 +38,7 @@ public class VersionHandlingIntegrationTest extends IntegTestCase {
     private final Setup setup = new Setup(sqlExecutor);
 
     @Test
-    public void selectMultiGetRequestWithColumnAlias() throws IOException {
+    public void selectMultiGetRequestWithColumnAlias() throws Exception {
         this.setup.createTestTableWithPrimaryKey();
         execute("insert into test (pk_col, message) values ('1', 'foo')");
         execute("insert into test (pk_col, message) values ('2', 'bar')");
diff --git a/server/src/testFixtures/java/io/crate/integrationtests/Setup.java b/server/src/testFixtures/java/io/crate/integrationtests/Setup.java
index 54edfa98f2a7..0f15f73fafd1 100644
--- a/server/src/testFixtures/java/io/crate/integrationtests/Setup.java
+++ b/server/src/testFixtures/java/io/crate/integrationtests/Setup.java
@@ -152,7 +152,7 @@ public void groupBySetup() throws Exception {
         groupBySetup("integer");
     }
 
-    public void groupBySetup(String numericType) {
+    public void groupBySetup(String numericType) throws Exception {
         transportExecutor.exec(String.format(Locale.ENGLISH, "create table characters (" +
                                                              " race string," +
                                                              " gender string," +
@@ -184,7 +184,7 @@ public void groupBySetup(String numericType) {
         transportExecutor.exec("refresh table characters");
     }
 
-    public void setUpEmployees() {
+    public void setUpEmployees() throws Exception {
         transportExecutor.exec("create table employees (" +
                                " name string, " +
                                " department string," +
@@ -238,7 +238,7 @@ public void setUpObjectTable() {
         transportExecutor.exec("refresh table ot");
     }
 
-    public void setUpArrayTables() {
+    public void setUpArrayTables() throws Exception {
         transportExecutor.exec("create table any_table (" +
                                "  id int primary key," +
                                "  temps array(double)," +
@@ -260,11 +260,11 @@ public void setUpArrayTables() {
         transportExecutor.exec("refresh table any_table");
     }
 
-    public void partitionTableSetup() {
+    public void partitionTableSetup() throws Exception {
         partitionTableSetup(null);
     }
 
-    public void partitionTableSetup(@Nullable Integer numberOfShards) {
+    public void partitionTableSetup(@Nullable Integer numberOfShards) throws Exception {
         String stmt = "create table parted (" +
             "id int primary key," +
             "date timestamp with time zone primary key," +
@@ -282,7 +282,7 @@ public void partitionTableSetup(@Nullable Integer numberOfShards) {
         transportExecutor.exec("refresh table parted");
     }
 
-    public void createTestTableWithPrimaryKey() {
+    public void createTestTableWithPrimaryKey() throws Exception {
         transportExecutor.exec("create table test (" +
                                "  pk_col string primary key, " +
                                "  message string" +
@@ -290,7 +290,7 @@ public void createTestTableWithPrimaryKey() {
         transportExecutor.ensureGreen();
     }
 
-    public void setUpCharacters() {
+    public void setUpCharacters() throws Exception {
         transportExecutor.exec("create table characters (id int primary key, name string, female boolean, details object)");
         transportExecutor.ensureYellowOrGreen();
         transportExecutor.execBulk("insert into characters (id, name, female) values (?, ?, ?)",
@@ -304,7 +304,7 @@ public void setUpCharacters() {
         transportExecutor.exec("refresh table characters");
     }
 
-    public void setUpJobs() {
+    public void setUpJobs() throws Exception {
         transportExecutor.exec("create table jobs (id int primary key, department string, min_salary double, max_salary double)");
         transportExecutor.ensureYellowOrGreen();
         transportExecutor.execBulk("insert into jobs (id, department, min_salary, max_salary) values (?, ?, ?, ?)",
diff --git a/server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java b/server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
index fbca1036b1e4..80e9f1fce67a 100644
--- a/server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
+++ b/server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
@@ -51,12 +51,11 @@
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.ElasticsearchTimeoutException;
 import org.elasticsearch.action.ActionListener;
-import org.elasticsearch.action.admin.cluster.health.ClusterHealthRequest;
-import org.elasticsearch.action.admin.cluster.health.ClusterHealthResponse;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.cluster.health.ClusterHealthStatus;
-import org.elasticsearch.common.Priority;
 import org.elasticsearch.common.util.concurrent.FutureUtils;
+import org.elasticsearch.discovery.MasterNotDiscoveredException;
+import org.elasticsearch.test.ESTestCase;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.VisibleForTesting;
@@ -499,48 +498,86 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        assertThat(state)
+            .as("ensureState can only be used to check for GREEN or YELLOW state")
+            .isNotEqualTo(ClusterHealthStatus.RED);
+        ESTestCase.assertBusy(() -> {
+            SQLResponse response;
+            try {
+                response = exec(
+                    """
+                        select
+                            (select count(*) from sys.nodes) as num_nodes,
+                            (select
+                                {
+                                    health = health,
+                                    pending_tasks = pending_tasks
+                                } as cluster_health
+                                FROM sys.cluster_health
+                            ),
+                            (select
+                                {
+                                    all_states = array_unique(array_agg(current_state)),
+                                    primary_states = array_unique(array_agg(current_state) FILTER (WHERE primary = true)),
+                                    relocating = count(*) FILTER (WHERE current_state = 'RELOCATING')
+                                } as shards_state
+                                FROM sys.allocations
+                            )
+                        """
+                );
+            } catch (Exception e) {
+                if (e instanceof MasterNotDiscoveredException) {
+                    // retry
+                    throw new AssertionError("Master not discovered", e);
+                }
+                throw e;
+            }
+            if (expectedNumNodes != null) {
+                assertThat(response.rows()[0][0]).isEqualTo((long) expectedNumNodes);
+            }
+            Map<String, Object> clusterHealth = (Map<String, Object>) response.rows()[0][1];
+            Map<String, Object> shardsState = (Map<String, Object>) response.rows()[0][2];
+
+            long pendingTasks = ((Number) clusterHealth.get("pending_tasks")).longValue();
+            assertThat(pendingTasks).isEqualTo(0);
+            long numRelocating = ((Number) shardsState.get("relocating")).longValue();
+            assertThat(numRelocating).isEqualTo(0);
+
+            String color = (String) clusterHealth.get("health");
+            List<String> statesToCheck;
+            if (state == ClusterHealthStatus.GREEN) {
+                assertThat(color).isEqualTo("GREEN");
+                statesToCheck = (List<String>) shardsState.get("all_states");
+            } else {
+                assertThat(color).satisfiesAnyOf(
+                    x -> assertThat(x).isEqualTo("GREEN"),
+                    x -> assertThat(x).isEqualTo("YELLOW")
+                );
+                statesToCheck = (List<String>) shardsState.get("primary_states");
+            }
+            assertThat(statesToCheck).satisfiesAnyOf(
+                x -> assertThat(x).isEmpty(),
+                x -> assertThat(x).isEqualTo(List.of("STARTED"))
+            );
+        });
     }
 
     public interface ClientProvider {


Reviews:
user: login: matriv
body: Thank you, left some comments/questions.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T12:49:14Z
user: login: seut
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T13:17:26Z
user: login: seut
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T13:19:21Z
user: login: seut
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T13:19:48Z
user: login: seut
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-18T13:20:27Z
user: login: matriv
body: nice, thank you!
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T13:59:54Z

Review comments:
diff_hunk: @@ -134,6 +136,7 @@ public void test_swap_table_partitioned_dotted_src_to_partitioned_dotted_target(
     }
 
     @Test
+    @Repeat (iterations = 20)
path: server/src/test/java/io/crate/integrationtests/FixCorruptedMetadataCommandITest.java
user: login: matriv
body: Please don't forget to remove, also on the test below.
created_at: 2025-03-18T12:25:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: matriv
body: minor: I see that currently we don't have a call with `RED`, so to make the code more robust, maybe throw an error before even the `assertBusy()` that `RED` is illegal?
created_at: 2025-03-18T12:37:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
+                assertion.satisfiesAnyOf(
+                    x -> assertThat(x).isEqualTo("GREEN"),
+                    x -> assertThat(x).isEqualTo("YELLOW")
+                );
+            }
+            // If the cluster is green/yellow, we can also check the state of the nodes & shards
+            var response = exec(
+                """
+                select
+                    (select count(*) from sys.nodes) as num_nodes,
+                    (select health from sys.cluster_health) as color,
+                    (select
+                        {
+                            all_states = array_unique(array_agg(current_state)),
+                            primary_states = array_unique(array_agg(current_state) FILTER (WHERE primary = true)),
+                            relocating = count(*) FILTER (WHERE current_state = 'RELOCATING')
+                        } as shards_state
+                        FROM sys.allocations
+                    )
+                """
+            );
+            if (expectedNumNodes != null) {
+                assertThat(response.rows()[0][0]).isEqualTo((long) expectedNumNodes);
+            }
+            color = (String) response.rows()[0][1];
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: matriv
body: why do we need to re-assert the color?
created_at: 2025-03-18T12:47:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
+                assertion.satisfiesAnyOf(
+                    x -> assertThat(x).isEqualTo("GREEN"),
+                    x -> assertThat(x).isEqualTo("YELLOW")
+                );
+            }
+            // If the cluster is green/yellow, we can also check the state of the nodes & shards
+            var response = exec(
+                """
+                select
+                    (select count(*) from sys.nodes) as num_nodes,
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: matriv
body: minor, but if we have a null `expectedNumNodes` we don't need to even run this query.
If you combine that we re-assert the color (unnecessarily?), I would split this into the sys.nodes query if required, and then do the `sys.allocations query`.
created_at: 2025-03-18T12:48:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
+                assertion.satisfiesAnyOf(
+                    x -> assertThat(x).isEqualTo("GREEN"),
+                    x -> assertThat(x).isEqualTo("YELLOW")
+                );
+            }
+            // If the cluster is green/yellow, we can also check the state of the nodes & shards
+            var response = exec(
+                """
+                select
+                    (select count(*) from sys.nodes) as num_nodes,
+                    (select health from sys.cluster_health) as color,
+                    (select
+                        {
+                            all_states = array_unique(array_agg(current_state)),
+                            primary_states = array_unique(array_agg(current_state) FILTER (WHERE primary = true)),
+                            relocating = count(*) FILTER (WHERE current_state = 'RELOCATING')
+                        } as shards_state
+                        FROM sys.allocations
+                    )
+                """
+            );
+            if (expectedNumNodes != null) {
+                assertThat(response.rows()[0][0]).isEqualTo((long) expectedNumNodes);
+            }
+            color = (String) response.rows()[0][1];
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: seut
body: To be sure nothing has changed in between the 2 SQL statements. Do you think this isn't needed?
created_at: 2025-03-18T13:17:26Z
author_association: MEMBER
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
+                assertion.satisfiesAnyOf(
+                    x -> assertThat(x).isEqualTo("GREEN"),
+                    x -> assertThat(x).isEqualTo("YELLOW")
+                );
+            }
+            // If the cluster is green/yellow, we can also check the state of the nodes & shards
+            var response = exec(
+                """
+                select
+                    (select count(*) from sys.nodes) as num_nodes,
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: seut
body: Using multiple queries will increase the possibility of state changes in-between. I don't think this is worth it, or?
created_at: 2025-03-18T13:19:21Z
author_association: MEMBER
diff_hunk: @@ -134,6 +136,7 @@ public void test_swap_table_partitioned_dotted_src_to_partitioned_dotted_target(
     }
 
     @Test
+    @Repeat (iterations = 20)
path: server/src/test/java/io/crate/integrationtests/FixCorruptedMetadataCommandITest.java
user: login: seut
body: right, forgot, thx!
created_at: 2025-03-18T13:19:47Z
author_association: MEMBER
diff_hunk: @@ -499,48 +497,80 @@ private BulkResponse executeBulk(String stmt, Object[][] bulkArgs, TimeValue tim
         }
     }
 
-    public ClusterHealthStatus ensureGreen() {
-        return ensureState(ClusterHealthStatus.GREEN, null);
+    public void ensureGreen() throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, null);
     }
 
-    public ClusterHealthStatus ensureGreen(Integer expectedNumNodes) {
-        return ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
+    public void ensureGreen(Integer expectedNumNodes) throws Exception {
+        ensureState(ClusterHealthStatus.GREEN, expectedNumNodes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen(Integer expectedNumNOdes) {
-        return ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
+    public void ensureYellowOrGreen(Integer expectedNumNOdes) throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, expectedNumNOdes);
     }
 
-    public ClusterHealthStatus ensureYellowOrGreen() {
-        return ensureState(ClusterHealthStatus.YELLOW, null);
+    public void ensureYellowOrGreen() throws Exception {
+        ensureState(ClusterHealthStatus.YELLOW, null);
     }
 
-    private ClusterHealthStatus ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) {
-        Client client = clientProvider.client();
-        ClusterHealthRequest request = new ClusterHealthRequest()
-            .waitForStatus(state)
-            .waitForEvents(Priority.LANGUID)
-            .waitForNoRelocatingShards(true)
-            .waitForNoInitializingShards(false);
-        if (expectedNumNodes != null) {
-            // We currently often use ensureGreen or ensureYellow to check whether the cluster is back in a good state after shutting down
-            // a node. If the node that is stopped is the master node, another node will become master and publish a cluster state where it
-            // is master but where the node that was stopped hasn't been removed yet from the cluster state. It will only subsequently
-            // publish a second state where the old master is removed. If the ensureGreen/ensureYellow is timed just right, it will get to
-            // execute before the second cluster state update removes the old master and the condition ensureGreen / ensureYellow will
-            // trivially hold if it held before the node was shut down. The following "waitForNodes" condition ensures that the node has
-            // been removed by the master so that the health check applies to the set of nodes we expect to be part of the cluster.
-            request.waitForNodes(Integer.toString(expectedNumNodes));
-        }
-        ClusterHealthResponse actionGet = FutureUtils.get(client.admin().cluster().health(request));
-        if (state == ClusterHealthStatus.YELLOW) {
-            assertThat(actionGet.getStatus()).satisfiesAnyOf(
-                s -> assertThat(s).isEqualTo(state),
-                s -> assertThat(s).isEqualTo(ClusterHealthStatus.GREEN));
-        } else {
-            assertThat(actionGet.getStatus()).isEqualTo(state);
-        }
-        return actionGet.getStatus();
+    @SuppressWarnings("unchecked")
+    private void ensureState(ClusterHealthStatus state, @Nullable Integer expectedNumNodes) throws Exception {
+        ESTestCase.assertBusy(() -> {
+            // First, check the global cluster health. If e.g. the cluster hasn't elected a master node, the 2nd query
+            // containing the `sys.allocations` table will fail this table is routed to the master node.
+            var clusterHealthResponse = exec("select health, description from sys.cluster_health");
+            String color = (String) clusterHealthResponse.rows()[0][0];
+            String description = (String) clusterHealthResponse.rows()[0][1];
+            var assertion = assertThat(color)
+                .withFailMessage("Expected cluster to be %s, but was %s, reason: %s", state.name(), color, description);
+            if (state == ClusterHealthStatus.GREEN) {
+                assertion.isEqualTo(state.name());
+            } else {
path: server/src/testFixtures/java/io/crate/testing/SQLTransportExecutor.java
user: login: seut
body: will add an assertion.
created_at: 2025-03-18T13:20:25Z
author_association: MEMBER

PR comments:
user: login: seut
created_at: 2025-03-18T11:32:17Z
body: @mfussenegger 
This is now ready for review, maybe you also wanna have a look.
I'd squash all commits afterwards.
author_association: MEMBER
user: login: mfussenegger
created_at: 2025-03-18T13:19:17Z
body: > @mfussenegger This is now ready for review, maybe you also wanna have a look. I'd squash all commits afterwards.

Looks good to merge to me once the repeats are removed again. Can't approve my own PR :)
author_association: MEMBER
user: login: mergify[bot]
created_at: 2025-03-20T14:00:34Z
body: This pull request has been removed from the queue for the following reason: <code>pull request dequeued</code>.

Pull request #17600 has been dequeued, merge conditions unmatch:
- [x] `label=ready-to-merge`
- [X] `#approved-reviews-by >= 1` [🛡 GitHub branch protection]
- [X] `#approved-reviews-by>=1`
- [X] `#changes-requested-reviews-by = 0` [🛡 GitHub branch protection]
- [X] any of [🛡 GitHub branch protection]:
  - [X] `check-success = ci/jenkins/pr_tests`
  - [ ] `check-neutral = ci/jenkins/pr_tests`
  - [ ] `check-skipped = ci/jenkins/pr_tests`

You should look at the reason for the failure and decide if the pull request needs to be fixed or if you want to requeue it.

If you want to requeue this pull request, you can post a `@mergifyio requeue` comment.

author_association: CONTRIBUTOR
user: login: mergify[bot]
created_at: 2025-03-20T14:40:06Z
body: This pull request has been removed from the queue for the following reason: <code>pull request dequeued</code>.

Pull request #17600 has been dequeued, merge conditions unmatch:
- [ ] `label=ready-to-merge`
- [ ] any of [🛡 GitHub branch protection]:
  - [ ] `check-neutral = ci/jenkins/pr_tests`
  - [ ] `check-skipped = ci/jenkins/pr_tests`
  - [ ] `check-success = ci/jenkins/pr_tests`
- [X] `#approved-reviews-by >= 1` [🛡 GitHub branch protection]
- [X] `#approved-reviews-by>=1`
- [X] `#changes-requested-reviews-by = 0` [🛡 GitHub branch protection]

You should look at the reason for the failure and decide if the pull request needs to be fixed or if you want to requeue it.

If you want to requeue this pull request, you can post a `@mergifyio requeue` comment.

author_association: CONTRIBUTOR