Only quote roles for connection impersonation when they contain special characters (#34858)

modules/drivers/snowflake/src/metabase/driver/snowflake.clj

@@ -581,7 +581,11 @@
 
 (defmethod driver.sql/set-role-statement :snowflake
   [_ role]
-  (format "USE ROLE \"%s\";" role))
+  (let [special-chars-pattern #"[^a-zA-Z0-9_]"
+        needs-quote           (re-find special-chars-pattern role)]
+    (if needs-quote
+      (format "USE ROLE \"%s\";" role)
+      (format "USE ROLE %s;" role))))
 
 (defmethod driver.sql/default-database-role :snowflake
   [_ database]

modules/drivers/snowflake/test/metabase/driver/snowflake_test.clj

@@ -5,6 +5,7 @@
    [clojure.string :as str]
    [clojure.test :refer :all]
    [metabase.driver :as driver]
+   [metabase.driver.sql :as driver.sql]
    [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
    [metabase.models :refer [Table]]
    [metabase.models.database :refer [Database]]
@@ -286,3 +287,14 @@
                                                       :regionid "us-west-1"}}]
         (is (= {:account "my-instance.us-west-1"}
                (:details db)))))))
+
+(deftest set-role-statement-test
+  (testing "set-role-statement should return a USE ROLE command, with the role quoted if it contains special characters"
+    ;; No special characters
+    (is (= "USE ROLE MY_ROLE;"        (driver.sql/set-role-statement :snowflake "MY_ROLE")))
+    (is (= "USE ROLE ROLE123;"        (driver.sql/set-role-statement :snowflake "ROLE123")))
+    (is (= "USE ROLE lowercase_role;" (driver.sql/set-role-statement :snowflake "lowercase_role")))
+
+    ;; Special characters
+    (is (= "USE ROLE \"Role.123\";"   (driver.sql/set-role-statement :snowflake "Role.123")))
+    (is (= "USE ROLE \"$role\";"      (driver.sql/set-role-statement :snowflake "$role")))))

src/metabase/driver/postgres.clj

@@ -851,9 +851,11 @@
 
 (defmethod driver.sql/set-role-statement :postgres
   [_ role]
-  (if (= (u/upper-case-en role) "NONE")
-   (format "SET ROLE %s;" role)
-   (format "SET ROLE \"%s\";" role)))
+  (let [special-chars-pattern #"[^a-zA-Z0-9_]"
+        needs-quote           (re-find special-chars-pattern role)]
+    (if needs-quote
+      (format "SET ROLE \"%s\";" role)
+      (format "SET ROLE %s;" role))))
 
 (defmethod driver.sql/default-database-role :postgres
   [_ _]

test/metabase/driver/postgres_test.clj

@@ -13,6 +13,7 @@
    [metabase.driver :as driver]
    [metabase.driver.postgres :as postgres]
    [metabase.driver.postgres.actions :as postgres.actions]
+   [metabase.driver.sql :as driver.sql]
    [metabase.driver.sql-jdbc.actions :as sql-jdbc.actions]
    [metabase.driver.sql-jdbc.actions-test :as sql-jdbc.actions-test]
    [metabase.driver.sql-jdbc.connection :as sql-jdbc.conn]
@@ -1266,3 +1267,18 @@
                             "REVOKE ALL PRIVILEGES ON SCHEMA \"dotted.schema\" FROM privilege_rows_test_example_role;"
                             "DROP ROLE privilege_rows_test_example_role;"]]
                 (jdbc/execute! conn-spec stmt)))))))))
+
+(deftest set-role-statement-test
+  (testing "set-role-statement should return a SET ROLE command, with the role quoted if it contains special characters"
+    ;; No special characters
+    (is (= "SET ROLE MY_ROLE;"        (driver.sql/set-role-statement :postgres "MY_ROLE")))
+    (is (= "SET ROLE ROLE123;"        (driver.sql/set-role-statement :postgres "ROLE123")))
+    (is (= "SET ROLE lowercase_role;" (driver.sql/set-role-statement :postgres "lowercase_role")))
+
+    ;; None (special role in Postgres to revert back to login role; should not be quoted)
+    (is (= "SET ROLE none;"      (driver.sql/set-role-statement :postgres "none")))
+    (is (= "SET ROLE NONE;"      (driver.sql/set-role-statement :postgres "NONE")))
+
+    ;; Special characters
+    (is (= "SET ROLE \"Role.123\";"   (driver.sql/set-role-statement :postgres "Role.123")))
+    (is (= "SET ROLE \"$role\";"      (driver.sql/set-role-statement :postgres "$role")))))