From 33496c167c2cf97742bb66e84b29593287b39e42 Mon Sep 17 00:00:00 2001
From: sharankow <168414585+sharankow@users.noreply.github.com>
Date: Thu, 22 Aug 2024 15:16:50 +0200
Subject: [PATCH] Support quoted roles (#262) (#266)

* fix: add support to quote roles

* test: add a test case for a role containing hyphens
---
 src/metabase/driver/clickhouse.clj               | 13 +++++++++++--
 .../driver/clickhouse_impersonation_test.clj     | 16 ++++++++++++----
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/metabase/driver/clickhouse.clj b/src/metabase/driver/clickhouse.clj
index 72945f9..50a54f4 100644
--- a/src/metabase/driver/clickhouse.clj
+++ b/src/metabase/driver/clickhouse.clj
@@ -283,8 +283,17 @@
 
 (defmethod driver.sql/set-role-statement :clickhouse
   [_ role]
-  (format "SET ROLE %s;" role))
+  (let [default-role (driver.sql/default-database-role :clickhouse nil)
+        quote-if-needed (fn [r]
+                          (if (or (re-matches #"\".*\"" r) (= role default-role))
+                            r
+                            (format "\"%s\"" r)))
+        quoted-role (->> (clojure.string/split role #",")
+                         (map quote-if-needed)
+                         (clojure.string/join ","))]
+    (format "SET ROLE %s;" quoted-role))
+  )
 
 (defmethod driver.sql/default-database-role :clickhouse
   [_ _]
-  "NONE")
+  "NONE")
\ No newline at end of file
diff --git a/test/metabase/driver/clickhouse_impersonation_test.clj b/test/metabase/driver/clickhouse_impersonation_test.clj
index 239650c..6e8c8be 100644
--- a/test/metabase/driver/clickhouse_impersonation_test.clj
+++ b/test/metabase/driver/clickhouse_impersonation_test.clj
@@ -29,6 +29,12 @@
        (fn [^java.sql.Connection conn]
          (driver/set-role! :clickhouse conn "metabase_test_role")))
       (is true))
+    (testing "does not throw with a role containing hyphens"
+      (sql-jdbc.execute/do-with-connection-with-options
+        :clickhouse spec nil
+        (fn [^java.sql.Connection conn]
+          (driver/set-role! :clickhouse conn "metabase-test-role")))
+      (is true))
     (testing "does not throw with the default role"
       (sql-jdbc.execute/do-with-connection-with-options
        :clickhouse spec nil
@@ -79,9 +85,10 @@
                          "CREATE OR REPLACE TABLE `metabase_test_role_db`.`some_table` (i Int32) ENGINE = MergeTree ORDER BY (i);"
                          "INSERT INTO `metabase_test_role_db`.`some_table` VALUES (42), (144);"
                          "CREATE ROLE IF NOT EXISTS `metabase_test_role`;"
+                         "CREATE ROLE IF NOT EXISTS `metabase-test-role`;"
                          "CREATE USER IF NOT EXISTS `metabase_test_user` NOT IDENTIFIED;"
-                         "GRANT SELECT ON `metabase_test_role_db`.* TO `metabase_test_role`;"
-                         "GRANT `metabase_test_role` TO `metabase_test_user`;"]]
+                         "GRANT SELECT ON `metabase_test_role_db`.* TO `metabase_test_role`,`metabase-test-role`;"
+                         "GRANT `metabase_test_role`, `metabase-test-role` TO `metabase_test_user`;"]]
          (ctd/exec-statements statements single-node-port-details)
          (do-with-new-metadata-provider
           single-node-details
@@ -99,9 +106,10 @@
                           ORDER BY (i);"
                          "INSERT INTO `metabase_test_role_db`.`some_table` VALUES (42), (144);"
                          "CREATE ROLE IF NOT EXISTS `metabase_test_role` ON CLUSTER '{cluster}';"
+                         "CREATE ROLE IF NOT EXISTS `metabase-test-role` ON CLUSTER '{cluster}';"
                          "CREATE USER IF NOT EXISTS `metabase_test_user` ON CLUSTER '{cluster}' NOT IDENTIFIED;"
-                         "GRANT ON CLUSTER '{cluster}' SELECT ON `metabase_test_role_db`.* TO `metabase_test_role`;"
-                         "GRANT ON CLUSTER '{cluster}' `metabase_test_role` TO `metabase_test_user`;"]]
+                         "GRANT ON CLUSTER '{cluster}' SELECT ON `metabase_test_role_db`.* TO `metabase_test_role`, `metabase-test-role`;"
+                         "GRANT ON CLUSTER '{cluster}' `metabase_test_role`, `metabase-test-role` TO `metabase_test_user`;"]]
          (ctd/exec-statements statements cluster-port-details)
          (do-with-new-metadata-provider
           cluster-details