From c41511e51c47c0b0bd1c49a79540bb870fd6dea0 Mon Sep 17 00:00:00 2001
From: MikhailBurdukov <burdukvmikhail@gmail.com>
Date: Tue, 5 Dec 2023 15:11:32 +0000
Subject: [PATCH] ignore_on_cluster_for_grants

---
 .../Access/InterpreterGrantQuery.cpp          |  6 ++--
 .../removeOnClusterClauseIfNeeded.cpp         |  4 ++-
 .../integration/test_replicated_users/test.py | 35 +++++++++++++++++++
 3 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/src/Interpreters/Access/InterpreterGrantQuery.cpp b/src/Interpreters/Access/InterpreterGrantQuery.cpp
index 45e8ba9ea0d..259c6b39524 100644
--- a/src/Interpreters/Access/InterpreterGrantQuery.cpp
+++ b/src/Interpreters/Access/InterpreterGrantQuery.cpp
@@ -7,6 +7,7 @@
 #include <Access/RolesOrUsersSet.h>
 #include <Access/User.h>
 #include <Interpreters/Context.h>
+#include <Interpreters/removeOnClusterClauseIfNeeded.h>
 #include <Interpreters/QueryLog.h>
 #include <Interpreters/executeDDLQueryOnCluster.h>
 #include <boost/range/algorithm/copy.hpp>
@@ -396,7 +397,8 @@ namespace
 
 BlockIO InterpreterGrantQuery::execute()
 {
-    auto & query = query_ptr->as<ASTGrantQuery &>();
+    const auto updated_query = removeOnClusterClauseIfNeeded(query_ptr, getContext());
+    auto & query = updated_query->as<ASTGrantQuery &>();
 
     query.replaceCurrentUserTag(getContext()->getUserName());
     query.access_rights_elements.eraseNonGrantable();
@@ -430,7 +432,7 @@ BlockIO InterpreterGrantQuery::execute()
         current_user_access->checkGranteesAreAllowed(grantees);
         DDLQueryOnClusterParams params;
         params.access_to_check = std::move(required_access);
-        return executeDDLQueryOnCluster(query_ptr, getContext(), params);
+        return executeDDLQueryOnCluster(updated_query, getContext(), params);
     }
 
     /// Check if the current user has corresponding access rights granted with grant option.
diff --git a/src/Interpreters/removeOnClusterClauseIfNeeded.cpp b/src/Interpreters/removeOnClusterClauseIfNeeded.cpp
index 7dc452a0fcb..da3930d62a6 100644
--- a/src/Interpreters/removeOnClusterClauseIfNeeded.cpp
+++ b/src/Interpreters/removeOnClusterClauseIfNeeded.cpp
@@ -14,6 +14,7 @@
 #include <Parsers/Access/ASTCreateSettingsProfileQuery.h>
 #include <Parsers/Access/ASTCreateUserQuery.h>
 #include <Parsers/Access/ASTDropAccessEntityQuery.h>
+#include <Parsers/Access/ASTGrantQuery.h>
 
 
 namespace DB
@@ -33,7 +34,8 @@ static bool isAccessControlQuery(const ASTPtr & query)
         || query->as<ASTCreateRoleQuery>()
         || query->as<ASTCreateRowPolicyQuery>()
         || query->as<ASTCreateSettingsProfileQuery>()
-        || query->as<ASTDropAccessEntityQuery>();
+        || query->as<ASTDropAccessEntityQuery>()
+        || query->as<ASTGrantQuery>();
 }
 
 ASTPtr removeOnClusterClauseIfNeeded(const ASTPtr & query, ContextPtr context, const WithoutOnClusterASTRewriteParams & params)
diff --git a/tests/integration/test_replicated_users/test.py b/tests/integration/test_replicated_users/test.py
index 489724ed4fb..e34495a0071 100644
--- a/tests/integration/test_replicated_users/test.py
+++ b/tests/integration/test_replicated_users/test.py
@@ -114,6 +114,41 @@ def test_create_replicated_on_cluster_ignore(started_cluster, entity):
     node1.query(f"DROP {entity.keyword} {entity.name} {entity.options}")
 
 
+@pytest.mark.parametrize(
+    "use_on_cluster",
+    [
+        pytest.param(False, id="Without_on_cluster"),
+        pytest.param(True, id="With_ignored_on_cluster"),
+    ],
+)
+def test_grant_revoke_replicated(started_cluster, use_on_cluster: bool):
+    node1.replace_config(
+        "/etc/clickhouse-server/users.d/users.xml",
+        inspect.cleandoc(
+            f"""
+            <clickhouse>
+                <profiles>
+                    <default>
+                        <ignore_on_cluster_for_replicated_access_entities_queries>{int(use_on_cluster)}</ignore_on_cluster_for_replicated_access_entities_queries>
+                    </default>
+                </profiles>
+            </clickhouse>
+            """
+        ),
+    )
+    node1.query("SYSTEM RELOAD CONFIG")
+    on_cluster = "ON CLUSTER default" if use_on_cluster else ""
+
+    node1.query(f"CREATE USER theuser {on_cluster}")
+
+    assert node1.query(f"GRANT {on_cluster} SELECT ON *.* to theuser") == ""
+
+    assert node2.query(f"SHOW GRANTS FOR theuser") == "GRANT SELECT ON *.* TO theuser\n"
+
+    assert node1.query(f"REVOKE {on_cluster} SELECT ON *.* from theuser") == ""
+    node1.query(f"DROP USER theuser {on_cluster}")
+
+
 @pytest.mark.parametrize("entity", entities, ids=get_entity_id)
 def test_create_replicated_if_not_exists_on_cluster(started_cluster, entity):
     node1.query(