[tests] Wait until users, roles and grants reach all cluster replicas

clickhouse-client/src/test/java/com/clickhouse/client/BaseIntegrationTest.java

@@ -86,6 +86,25 @@ protected static String onCluster() {
         return ClickHouseServerForTest.isCluster() ? " ON CLUSTER 'test_cluster'" : "";
     }
 
+    /**
+     * Waits until each replica of the test cluster returns {@code expectedRows} rows of
+     * {@code system.<table>} that match {@code condition}, so that the next request sees a
+     * created or changed user, role, settings profile or grant on any replica.
+     * Does nothing when the tests do not run on a cluster.
+     */
+    protected static void waitForAllReplicas(String table, String condition, int expectedRows) {
+        ClickHouseServerForTest.waitForAllReplicas(table, condition, expectedRows);
+    }
+
+    /**
+     * Same check as {@link #waitForAllReplicas(String, String, int)}. {@code query} runs the
+     * replica SQL with the caller's client and returns one matching-row count per replica host.
+     */
+    protected static void waitForAllReplicas(String table, String condition, int expectedRows,
+            ClickHouseServerForTest.ReplicaQuery query) {
+        ClickHouseServerForTest.waitForAllReplicas(table, condition, expectedRows, query);
+    }
+
     protected boolean createDatabase(String dbName) {
         return ClickHouseServerForTest.runQuery("CREATE DATABASE IF NOT EXISTS `" + dbName + "`" + onCluster());
     }

clickhouse-client/src/test/java/com/clickhouse/client/ClickHouseServerForTest.java

@@ -4,19 +4,24 @@
 import org.slf4j.LoggerFactory;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.Network;
+import org.testng.Assert;
 import org.testng.annotations.AfterSuite;
 import org.testng.annotations.BeforeSuite;
 
 import com.clickhouse.client.ClickHouseTestEnvironment.ContainerMode;
+import com.clickhouse.data.ClickHouseRecord;
 
 import java.io.OutputStream;
 import java.net.HttpURLConnection;
 import java.net.InetSocketAddress;
 import java.net.URL;
 import java.nio.charset.StandardCharsets;
-import java.util.Locale;
+import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Base64;
 import java.util.Collections;
+import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 
 /**
@@ -33,6 +38,8 @@
 @SuppressWarnings("squid:S2187")
 public class ClickHouseServerForTest {
     private static final Logger LOGGER = LoggerFactory.getLogger(ClickHouseServerForTest.class);
+    private static final int REPLICA_CHECK_ATTEMPTS = 100;
+    private static final long REPLICA_CHECK_INTERVAL_MS = 100L;
 
     private static final Network network;
     private static final ClickHouseTestEnvironment environment;
@@ -279,6 +286,89 @@
         return false;
     }
 
+    /**
+     * Waits until each replica of the test cluster returns {@code expectedRows} rows of
+     * {@code system.<table>} that match {@code condition}. Does nothing outside cluster mode.
+     * Uses {@link ClickHouseClient}, which resolves to the HTTP client when that module is on the classpath.
+     */
+    public static void waitForAllReplicas(String table, String condition, int expectedRows) {
+        if (!isCluster()) {
+            return;
+        }
+        ClickHouseNode server = getClickHouseNode(ClickHouseProtocol.HTTP, isCloud(), ClickHouseNode.builder().build());
+        try (ClickHouseClient client = ClickHouseClient.newInstance(ClickHouseProtocol.HTTP)) {
+            waitForAllReplicas(table, condition, expectedRows, sql -> matchingRowCounts(client, server, sql));
+        }
+    }
+
+    /**
+     * Same check as {@link #waitForAllReplicas(String, String, int)}. {@code query} runs the
+     * replica SQL and returns one matching-row count per replica host.
+     */
+    public static void waitForAllReplicas(String table, String condition, int expectedRows, ReplicaQuery query) {
+        if (!isCluster()) {
+            return;
+        }
+        String sql = "SELECT hostName() AS host, toInt64(countIf(" + condition + ")) AS matching_rows"
+                + " FROM clusterAllReplicas('" + ClusterClickHouseContainerStarter.CLUSTER_NAME + "', system." + table
+                + ") GROUP BY host";
+        String state = "";
+        for (int attempt = 1; attempt <= REPLICA_CHECK_ATTEMPTS; attempt++) {
+            try {
+                long[] counts = query.matchingRowCounts(sql);
+                state = Arrays.toString(counts);
+                if (hasRowsOnAllReplicas(counts, expectedRows)) {
+                    return;
+                }
+            } catch (Exception e) {
+                state = e.getMessage();
+            }
+            try {
+                Thread.sleep(REPLICA_CHECK_INTERVAL_MS);
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                break;
+            }
+        }
+        Assert.fail("Expected " + expectedRows + " rows matching [" + condition + "] in system." + table + " on all "
+                + ClusterClickHouseContainerStarter.NODE_COUNT + " replicas, but got " + state);
+    }
+
+    private static long[] matchingRowCounts(ClickHouseClient client, ClickHouseNode server, String sql)
+            throws ClickHouseException {
+        try (ClickHouseResponse response = client.read(server).query(sql).executeAndWait()) {
+            List<Long> counts = new ArrayList<>();
+            for (ClickHouseRecord record : response.records()) {
+                counts.add(record.getValue(1).asLong());
+            }
+            long[] result = new long[counts.size()];
+            for (int i = 0; i < result.length; i++) {
+                result[i] = counts.get(i);
+            }
+            return result;
+        }
+    }
+
+    private static boolean hasRowsOnAllReplicas(long[] counts, int expectedRows) {
+        if (counts.length != ClusterClickHouseContainerStarter.NODE_COUNT) {
+            return false;
+        }
+        for (long count : counts) {
+            if (count != expectedRows) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Runs the replica-check SQL from {@link #waitForAllReplicas(String, String, int, ReplicaQuery)}.
+     */
+    @FunctionalInterface
+    public interface ReplicaQuery {
+        long[] matchingRowCounts(String sql) throws Exception;
+    }
+
     private static Endpoint resolve(ClickHouseProtocol protocol, int port, boolean applyOverride) {
         if (starter != null) {
             int containerPort = (port == ClickHouseProtocol.HTTP.getDefaultPort())

client-v2/src/test/java/com/clickhouse/client/ClientTests.java

@@ -434,28 +434,46 @@ private static Object[][] sessionRoles() {
     }
 
     @Test(groups = {"integration"}, dataProvider = "sessionRoles")
-    public void testOperationCustomRoles(String[] roles) throws Exception {
+    public void testOperationCustomRoles(String[] roleNames) throws Exception {
         if (isVersionMatch("(,24.3]", newClient().build())) {
             return;
         }
 
+        String suffix = "_" + RandomStringUtils.insecure().next(8, true, true);
+        String user = "some_user" + suffix;
+        String[] roles = Arrays.stream(roleNames).map(role -> role + suffix).toArray(String[]::new);
         String password = "^1A" + RandomStringUtils.random(12, true, true) + "3b$";
         final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
-        try (Client client = newClient().build()) {
-            client.execute("DROP ROLE IF EXISTS " + rolesList).get().close();
-            client.execute("CREATE ROLE " + rolesList).get().close();
-            client.execute("DROP USER IF EXISTS some_user").get().close();
-            client.execute("CREATE USER some_user IDENTIFIED BY '" + password + "'").get().close();
-            client.execute("GRANT " + rolesList + " TO some_user").get().close();
+        try (Client client = newClient().build();
+             Client userClient = newClient().setUsername(user).setPassword(password).build()) {
+            try {
+                createUserWithRoles(client, user, "IDENTIFIED BY '" + password + "'", roles);
+                QuerySettings settings = new QuerySettings().setDBRoles(Arrays.asList(roles));
+                List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()", settings);
+                Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
+                Set<String> currentRoles = new  HashSet<String> (resp.get(0).getList(1));
+                Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
+            } finally {
+                client.execute("DROP USER IF EXISTS " + user).get().close();
+                client.execute("DROP ROLE IF EXISTS " + rolesList).get().close();
+            }
         }
+    }
 
-        try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
-            QuerySettings settings = new QuerySettings().setDBRoles(Arrays.asList(roles));
-            List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()", settings);
-            Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
-            Set<String> currentRoles = new  HashSet<String> (resp.get(0).getList(1));
-            Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
-        }
+    private static void createUserWithRoles(Client client, String user, String identification, String[] roles)
+            throws Exception {
+        final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
+        final String roleNames = "'" + Strings.join("','", roles) + "'";
+        client.execute("CREATE ROLE " + rolesList).get().close();
+        client.execute("CREATE USER " + user + " " + identification).get().close();
+        waitForAllReplicas("roles", "name IN (" + roleNames + ")", roles.length,
+                sql -> client.queryAll(sql).stream().mapToLong(row -> row.getLong("matching_rows")).toArray());
+        waitForAllReplicas("users", "name = '" + user + "'", 1,
+                sql -> client.queryAll(sql).stream().mapToLong(row -> row.getLong("matching_rows")).toArray());
+        client.execute("GRANT " + rolesList + " TO " + user).get().close();
+        waitForAllReplicas("role_grants",
+                "user_name = '" + user + "' AND granted_role_name IN (" + roleNames + ")", roles.length,
+                sql -> client.queryAll(sql).stream().mapToLong(row -> row.getLong("matching_rows")).toArray());
     }
 
     @DataProvider(name = "clientSessionRoles")
@@ -466,27 +484,29 @@ private static Object[][] clientSessionRoles() {
         };
     }
     @Test(groups = {"integration"}, dataProvider = "clientSessionRoles")
-    public void testClientCustomRoles(String[] roles) throws Exception {
+    public void testClientCustomRoles(String[] roleNames) throws Exception {
         if (isVersionMatch("(,24.3]", newClient().build())) {
             return;
         }
 
+        String suffix = "_" + RandomStringUtils.insecure().next(8, true, true);
+        String user = "some_user" + suffix;
+        String[] roles = Arrays.stream(roleNames).map(role -> role + suffix).toArray(String[]::new);
         String password = "^1A" + RandomStringUtils.insecure().next(12, true, true) + "3B$";
         final String rolesList = "\"" + Strings.join("\",\"", roles) + "\"";
-        try (Client client = newClient().build()) {
-            client.execute("DROP ROLE IF EXISTS " + rolesList).get().close();
-            client.execute("CREATE ROLE " + rolesList).get().close();
-            client.execute("DROP USER IF EXISTS some_user").get().close();
-            client.execute("CREATE USER some_user IDENTIFIED WITH sha256_password BY '" + password + "'").get().close();
-            client.execute("GRANT " + rolesList + " TO some_user").get().close();
-        }
-
-        try (Client userClient = newClient().setUsername("some_user").setPassword(password).build()) {
-            userClient.setDBRoles(Arrays.asList(roles));
-            List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()");
-            Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
-            Set<String> currentRoles = new  HashSet<String> (resp.get(0).getList(1));
-            Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
+        try (Client client = newClient().build();
+             Client userClient = newClient().setUsername(user).setPassword(password).build()) {
+            try {
+                createUserWithRoles(client, user, "IDENTIFIED WITH sha256_password BY '" + password + "'", roles);
+                userClient.setDBRoles(Arrays.asList(roles));
+                List<GenericRecord> resp = userClient.queryAll("SELECT currentRoles()");
+                Set<String> roleSet = new HashSet<>(Arrays.asList(roles));
+                Set<String> currentRoles = new  HashSet<String> (resp.get(0).getList(1));
+                Assert.assertEquals(currentRoles, roleSet, "Roles " + roleSet + " not found in " + currentRoles);
+            } finally {
+                client.execute("DROP USER IF EXISTS " + user).get().close();
+                client.execute("DROP ROLE IF EXISTS " + rolesList).get().close();
+            }
         }
     }
 
@@ -507,36 +527,16 @@ public void testRuntimeCredentialChange() throws Exception {
                 adminClient.execute("DROP USER IF EXISTS " + user2).get().close();
                 adminClient.execute("CREATE USER " + user1 + " IDENTIFIED BY '" + password1 + "'").get().close();
                 adminClient.execute("CREATE USER " + user2 + " IDENTIFIED BY '" + password2 + "'").get().close();
+                waitForAllReplicas("users", "name IN ('" + user1 + "', '" + user2 + "')", 2,
+                        sql -> adminClient.queryAll(sql).stream().mapToLong(row -> row.getLong("matching_rows")).toArray());
 
                 try (Client userClient = newClient().setUsername(user1).setPassword(password1).build()) {
-                    List<GenericRecord> firstResponse = null;
-                    for (int i = 0; i < 20; i++) {
-                        try {
-                            firstResponse = userClient.queryAll("SELECT currentUser() AS user");
-                            break;
-                        } catch (ClientException e) {
-                            if (i == 19) {
-                                throw e;
-                            }
-                            Thread.sleep(250);
-                        }
-                    }
+                    List<GenericRecord> firstResponse = userClient.queryAll("SELECT currentUser() AS user");
                     Assert.assertEquals(firstResponse.get(0).getString("user"), user1);
 
                     userClient.updateUserAndPassword(user2, password2);
 
-                    List<GenericRecord> secondResponse = null;
-                    for (int i = 0; i < 20; i++) {
-                        try {
-                            secondResponse = userClient.queryAll("SELECT currentUser() AS user");
-                            break;
-                        } catch (ClientException e) {
-                            if (i == 19) {
-                                throw e;
-                            }
-                            Thread.sleep(250);
-                        }
-                    }
+                    List<GenericRecord> secondResponse = userClient.queryAll("SELECT currentUser() AS user");
                     Assert.assertEquals(secondResponse.get(0).getString("user"), user2);
                 }
             } finally {