make connection impersonation work when acquiring a connection via jdbc spec (#32238)

enterprise/backend/src/metabase_enterprise/advanced_permissions/driver/impersonation.clj

@@ -45,7 +45,7 @@
   Returns `nil` if connection impersonation should not be used for the current user. Throws an exception if multiple
   conflicting connection impersonation policies are found."
   [database-or-id]
-  (when-not api/*is-superuser?*
+  (when (and database-or-id (not api/*is-superuser?*))
     (let [group-ids           (t2/select-fn-set :group_id PermissionsGroupMembership :user_id api/*current-user-id*)
           conn-impersonations (enforced-impersonations
                                (when (seq group-ids)
@@ -81,8 +81,8 @@
   [driver ^Connection conn database]
   (when (driver/database-supports? driver :connection-impersonation database)
     (try
-      (let [default-role       (driver.sql/default-database-role driver database)
-            impersonation-role (connection-impersonation-role database)]
+      (let [default-role              (driver.sql/default-database-role driver database)
+            impersonation-role        (connection-impersonation-role database)]
         (driver/set-role! driver conn (or impersonation-role default-role)))
       (catch Throwable e
         (log/debug e (tru "Error setting role on connection"))

enterprise/backend/test/metabase_enterprise/advanced_permissions/driver/impersonation_test.clj

@@ -50,8 +50,13 @@
   (testing "Returns nil for superuser, even if they are in a group with an impersonation policy defined"
     (advanced-perms.api.tu/with-impersonations {:impersonations [{:db-id (mt/id) :attribute "impersonation_attr"}]
                                                 :attributes     {"impersonation_attr" "impersonation_role"}}
-      (is (= "impersonation_role"
-             (@#'impersonation/connection-impersonation-role (mt/db)))))))
+      (mw.session/as-admin
+       (is (nil? (@#'impersonation/connection-impersonation-role (mt/db)))))))
+
+  (testing "Does not throw an exception if passed a nil `database-or-id`"
+    (advanced-perms.api.tu/with-impersonations {:impersonations [{:db-id (mt/id) :attribute "impersonation_attr"}]
+                                                :attributes     {"impersonation_attr" "impersonation_role"}}
+      (is (nil? (@#'impersonation/connection-impersonation-role nil))))))
 
 (deftest conn-impersonation-test-postgres
   (mt/test-driver :postgres

src/metabase/driver/sql_jdbc/execute.clj

@@ -345,9 +345,8 @@
     (log/tracef "Setting default connection options with options %s" (pr-str options))
     (set-best-transaction-level! driver conn)
     (set-time-zone-if-supported! driver conn session-timezone)
-    (set-role-if-supported! driver conn (when (u/id db-or-id-or-spec) (if (integer? db-or-id-or-spec)
-                                                                        (t2/select-one Database db-or-id-or-spec)
-                                                                        db-or-id-or-spec)))
+    (set-role-if-supported! driver conn (cond (integer? db-or-id-or-spec) (t2/select-one Database db-or-id-or-spec)
+                                              (u/id db-or-id-or-spec)     db-or-id-or-spec))
     (let [read-only? (not write?)]
       (try
         ;; Setting the connection to read-only does not prevent writes on some databases, and is meant