Merge pull request #1649 from ClickHouse/fix_1647

Reset roles by "SET ROLE NONE"

CHANGELOG.md

@@ -2,6 +2,13 @@
 
 ## 0.6.0-patch4
 
+### Bug Fixes
+- Roles (memorized by client) can be reset with 'SET ROLE NONE' query (https://github.com/ClickHouse/clickhouse-java/issues/1647)
+- MaxBufferSize can be greater than internal MAX value now
+- Updated example project to use the latest version of the client
+
+## 0.6.0-patch4
+
 ### New Features
 - Added possibility to set client ID in `Referer` HTTP Header (https://github.com/ClickHouse/clickhouse-java/issues/1572)
 - [HTTP] Persistence of a role after it is set by `SET ROLE <role>`

clickhouse-http-client/src/main/java/com/clickhouse/client/http/ClickHouseHttpClient.java

@@ -30,6 +30,7 @@
 import java.util.Set;
 import java.util.concurrent.CompletionException;
 import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.stream.Collectors;
 
 public class ClickHouseHttpClient extends AbstractClient<ClickHouseHttpConnection> {
     private static final Logger log = LoggerFactory.getLogger(ClickHouseHttpClient.class);
@@ -87,7 +88,7 @@ public ClickHouseHttpClient() {
         }
     }
 
-    private ConcurrentSkipListSet<String> roles = new ConcurrentSkipListSet<>();
+    protected ConcurrentSkipListSet<String> roles = new ConcurrentSkipListSet<>();
 
     @Override
     protected boolean checkConnection(ClickHouseHttpConnection connection, ClickHouseNode requestServer,
@@ -150,9 +151,7 @@ protected String buildQueryParams(Map<String, String> params) {
     private Map<String, Serializable> buildAdditionalReqParams(ClickHouseRequest<?> sealedRequest) {
         ClickHouseConfig config = sealedRequest.getConfig();
         if (config.getBoolOption(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES)) {
-            if (sealedRequest.hasSetting("_set_roles_stmt")) {
-                return Collections.singletonMap("_roles", sealedRequest.getSettings().get("_set_roles_stmt"));
-            } else if (!roles.isEmpty()) {
+            if (!(sealedRequest.hasSetting("_set_roles_stmt") || roles.isEmpty())) {
                 return Collections.singletonMap("_roles", roles);
             }
         }
@@ -220,8 +219,18 @@ public final Class<? extends ClickHouseOption> getOptionClass() {
         return ClickHouseHttpOption.class;
     }
 
-    private void rememberRoles(Set<String> requestedRoles) {
+    public void rememberRoles(Set<String> requestedRoles) {
+        roles.clear();
+        if (requestedRoles != null) {
+            roles.addAll(requestedRoles.stream().filter(r -> !"NONE".equalsIgnoreCase(r)).collect(Collectors.toList()));
+        }
+    }
+
+    public Set<String> getRoles() {
+        return roles;
+    }
+
+    public void clearRoles() {
         roles.clear();
-        roles.addAll(requestedRoles);
     }
 }

clickhouse-jdbc/src/main/java/com/clickhouse/jdbc/internal/ClickHouseStatementImpl.java

@@ -138,6 +138,7 @@ private ClickHouseResponse getLastResponse(Map<ClickHouseOption, Serializable> o
             } catch (Exception e) {
                 throw SqlExceptionUtils.handle(e);
             } finally {
+                request.removeSetting("_set_roles_stmt");
                 if (response == null) {
                     // something went wrong
                 } else if (i + 1 < len) {

clickhouse-jdbc/src/test/java/com/clickhouse/jdbc/AccessManagementTest.java

@@ -1,6 +1,5 @@
 package com.clickhouse.jdbc;
 
-import com.clickhouse.client.ClickHouseException;
 import com.clickhouse.client.ClickHouseProtocol;
 import com.clickhouse.client.http.config.ClickHouseHttpOption;
 import com.clickhouse.client.http.config.HttpConnectionProvider;
@@ -28,6 +27,11 @@ public void testSetRoleDifferentConnections(String[] roles, String setRoleExpr,
         properties.setProperty(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES.getKey(), "true");
         properties.setProperty(ClickHouseHttpOption.CONNECTION_PROVIDER.getKey(), connectionProvider);
         ClickHouseDataSource dataSource = new ClickHouseDataSource(url, properties);
+        String serverVersion = getServerVersion(dataSource.getConnection());
+        if (ClickHouseVersion.of(serverVersion).check("(,24.3]")) {
+            System.out.println("Test is skipped: feature is supported since 24.4");
+            return;
+        }
 
         try (Connection connection = dataSource.getConnection("access_dba", "123")) {
             Statement st = connection.createStatement();
@@ -37,22 +41,18 @@ public void testSetRoleDifferentConnections(String[] roles, String setRoleExpr,
             st.execute("CREATE ROLE " + String.join(", ", roles));
             st.execute("CREATE USER some_user IDENTIFIED WITH no_password");
             st.execute("GRANT " + String.join(", ", roles) + " TO some_user");
+            st.execute("SET DEFAULT ROLE NONE TO some_user");
         } catch (Exception e) {
-            Assert.fail("Failed", e);
+            Assert.fail("Failed to prepare for the test", e);
         }
 
         try (Connection connection = dataSource.getConnection("some_user", "")) {
             Statement st = connection.createStatement();
             st.execute(setRoleExpr);
             assertRolesEquals(connection, activeRoles);
-        } catch (SQLException e) {
-            if (e.getErrorCode() == ClickHouseException.ERROR_UNKNOWN_SETTING) {
-                String serverVersion = getServerVersion(dataSource.getConnection());
-                if (ClickHouseVersion.of(serverVersion).check("(,24.3]")) {
-                    return;
-                }
-            }
-            Assert.fail("Failed", e);
+            // Check roles are reset
+            st.execute("SET ROLE NONE");
+            assertRolesEquals(connection);
         } catch (Exception e) {
             Assert.fail("Failed", e);
         }
@@ -74,7 +74,7 @@ private static Object[][] setRolesArgsForTestSetRole() {
         };
     }
 
-    private void assertRolesEquals(Connection connection, String... expected) {
+    private void assertRolesEquals(Connection connection, String... expected) throws SQLException {
         try {
             Statement st = connection.createStatement();
             ResultSet resultSet = st.executeQuery("select currentRoles()");
@@ -83,13 +83,9 @@ private void assertRolesEquals(Connection connection, String... expected) {
             String[] roles = (String[]) resultSet.getArray(1).getArray();
             Arrays.sort(roles);
             Arrays.sort(expected);
-            System.out.print("Roles: ");
-            for (String role : roles) {
-                System.out.print("'" + role + "', ");
-            }
-            System.out.println();
-            Assert.assertEquals(roles, expected);
-
+            Assert.assertEquals(roles, expected,
+                    "Memorized roles: " + Arrays.toString(roles) + " != Expected: " + Arrays.toString(expected));
+            System.out.println("Roles: " + Arrays.toString(roles));
         } catch (Exception e) {
             Assert.fail("Failed", e);
         }
@@ -106,4 +102,78 @@ private String getServerVersion(Connection connection) {
         }
         return null;
     }
+
+    @Test
+    public void testSetRolesAccessingTableRows() throws SQLException {
+        String httpEndpoint = "http://" + getServerAddress(ClickHouseProtocol.HTTP) + "/";
+        String url = String.format("jdbc:ch:%s", httpEndpoint);
+        Properties properties = new Properties();
+        properties.setProperty(ClickHouseHttpOption.REMEMBER_LAST_SET_ROLES.getKey(), "true");
+        ClickHouseDataSource dataSource = new ClickHouseDataSource(url, properties);
+        String serverVersion = getServerVersion(dataSource.getConnection());
+        if (ClickHouseVersion.of(serverVersion).check("(,24.3]")) {
+            System.out.println("Test is skipped: feature is supported since 24.4");
+            return;
+        }
+
+        try (Connection connection = dataSource.getConnection("access_dba", "123")) {
+            Statement st = connection.createStatement();
+            st.execute("DROP ROLE IF EXISTS row_a");
+            st.execute("DROP USER IF EXISTS some_user");
+
+            st.execute("CREATE ROLE row_a, row_b");
+            st.execute("CREATE USER some_user IDENTIFIED WITH no_password");
+            st.execute("GRANT row_a, row_b TO some_user");
+
+            st.execute("CREATE OR REPLACE TABLE test_table (`s` String ) ENGINE = MergeTree ORDER BY tuple();");
+            st.execute("INSERT INTO test_table VALUES ('a'), ('b')");
+
+            st.execute("GRANT SELECT ON test_table TO some_user");
+            st.execute("CREATE ROW POLICY OR REPLACE policy_row_b ON test_table FOR SELECT USING s = 'b' TO row_b;");
+            st.execute("CREATE ROW POLICY OR REPLACE policy_row_a ON test_table FOR SELECT USING s = 'a' TO row_a;");
+        } catch (Exception e) {
+            Assert.fail("Failed on setup", e);
+        }
+
+        try (Connection connection = dataSource.getConnection("some_user", "")) {
+            Statement st = connection.createStatement();
+            ResultSet rs = st.executeQuery("SELECT * FROM test_table");
+            Assert.assertTrue(rs.next());
+            Assert.assertTrue(rs.next());
+
+            st.execute("SET ROLE row_a");
+            rs = st.executeQuery("SELECT * FROM test_table");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString(1), "a");
+            Assert.assertFalse(rs.next());
+
+            st.execute("SET ROLE row_b");
+            rs = st.executeQuery("SELECT * FROM test_table");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString(1), "b");
+            Assert.assertFalse(rs.next());
+
+
+            st.execute("SET ROLE row_a, row_b");
+            rs = st.executeQuery("SELECT * FROM test_table ORDER BY s");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString(1), "a");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString(1), "b");
+            Assert.assertFalse(rs.next());
+
+            st.execute("SET ROLE row_b");
+            rs = st.executeQuery("SELECT * FROM test_table");
+            Assert.assertTrue(rs.next());
+            Assert.assertEquals(rs.getString(1), "b");
+            Assert.assertFalse(rs.next());
+
+            st.execute("SET ROLE NONE");
+            rs = st.executeQuery("SELECT * FROM test_table");
+            Assert.assertTrue(rs.next());
+            Assert.assertTrue(rs.next());
+        } catch (Exception e) {
+            Assert.fail("Failed to check roles", e);
+        }
+    }
 }