Fix N+1 in checking write permission for collection (#40345) (#40431)

Co-authored-by: Ngoc Khuat <qn.khuat@gmail.com>

enterprise/backend/src/metabase_enterprise/audit_app/api/user.clj

@@ -17,8 +17,8 @@
   Otherwise return an empty map."
   []
   (let [custom-reports     (audit-db/default-custom-reports-collection)
-        question-overview  (audit-db/entity-id->object :model/Dashboard audit-db/default-question-overview-entity-id)
-        dashboard-overview (audit-db/entity-id->object :model/Dashboard audit-db/default-dashboard-overview-entity-id)]
+        question-overview  (audit-db/memoized-select-audit-entity :model/Dashboard audit-db/default-question-overview-entity-id)
+        dashboard-overview (audit-db/memoized-select-audit-entity :model/Dashboard audit-db/default-dashboard-overview-entity-id)]
     (merge
      {}
      (when (mi/can-read? (audit-db/default-custom-reports-collection))

enterprise/backend/src/metabase_enterprise/audit_db.clj

@@ -21,6 +21,8 @@
    (java.util.jar JarEntry JarFile)
    (java.nio.file Path)))
 
+(def ^:private audit-installed? (atom false))
+
 (set! *warn-on-reflection* true)
 
 (defn- running-from-jar?
@@ -84,25 +86,30 @@
   "Default Dashboard Overview (this is a dashboard) entity id."
   "bJEYb0o5CXlfWFcIztDwJ")
 
-(defn entity-id->object
+(def ^{:arglists '([audit-installed? model entity-id])
+       :private  true} memoized-select-audit-entity*
+  (mdb.connection/memoize-for-application-db
+   (fn [audit-installed? model entity-id]
+     (when audit-installed?
+       (t2/select-one model :entity_id entity-id)))))
+
+(defn memoized-select-audit-entity
   "Returns the object from entity id and model. Memoizes from entity id.
   Should only be used for audit/pre-loaded objects."
   [model entity-id]
-  ((mdb.connection/memoize-for-application-db
-    (fn [entity-id]
-      (t2/select-one model :entity_id entity-id))) entity-id))
+  (memoized-select-audit-entity* @audit-installed? model entity-id))
 
 (defenterprise default-custom-reports-collection
   "Default custom reports collection."
   :feature :none
   []
-  (entity-id->object :model/Collection default-custom-reports-entity-id))
+  (memoized-select-audit-entity :model/Collection default-custom-reports-entity-id))
 
 (defenterprise default-audit-collection
   "Default audit collection (instance analytics) collection."
   :feature :none
   []
-  (entity-id->object :model/Collection default-audit-collection-entity-id))
+  (memoized-select-audit-entity :model/Collection default-audit-collection-entity-id))
 
 (defn- install-database!
   "Creates the audit db, a clone of the app db used for auditing purposes.
@@ -290,4 +297,6 @@
    (let [audit-db (t2/select-one :model/Database :is_audit true)]
        ;; prevent sync while loading
      ((sync-util/with-duplicate-ops-prevented :sync-database audit-db
-        (fn [] (maybe-load-analytics-content! audit-db)))))))
+        (fn []
+          (maybe-load-analytics-content! audit-db)
+          (reset! audit-installed? true)))))))

enterprise/backend/test/metabase_enterprise/audit_app/api/user_test.clj

@@ -8,9 +8,12 @@
    [metabase.models.permissions :as perms]
    [metabase.models.permissions-group :as perms-group]
    [metabase.test :as mt]
+   [metabase.test.fixtures :as fixtures]
    [toucan2.core :as t2]
    [toucan2.tools.with-temp :as t2.with-temp]))
 
+(use-fixtures :once (fixtures/initialize :db))
+
 (deftest delete-subscriptions-test
   (testing "DELETE /api/ee/audit-app/user/:id/subscriptions"
     (testing "Should require a token with `:audit-app`"

src/metabase/api/collection.clj

@@ -100,7 +100,7 @@
                ;; Order NULL collection types first so that audit collections are last
                :order-by [[[[:case [:= :type nil] 0 :else 1]] :asc]
                           [:%lower.name :asc]]})
-    exclude-other-user-collections (remove-other-users-personal-subcollections api/*current-user-id*)))
+   exclude-other-user-collections (remove-other-users-personal-subcollections api/*current-user-id*)))
 
 (api/defendpoint GET "/"
   "Fetch a list of all Collections that the current user has read permissions for (`:can_write` is returned as an
@@ -538,12 +538,7 @@
               (assoc row :name (collection/user->personal-collection-name (:personal_owner_id row) :user))
               (dissoc row :personal_owner_id)))]
     (for [row rows]
-      ;; Go through this rigamarole instead of hydration because we
-      ;; don't get models back from ulterior over-query
-      ;; Previous examination with logging to DB says that there's no N+1 query for this.
-      ;; However, this was only tested on H2 and Postgres
-      (-> row
-          (assoc :can_write (mi/can-write? Collection (:id row)))
+      (-> (t2/hydrate (t2/instance :model/Collection row) :can_write)
           (dissoc :collection_position :display :moderated_status :icon
                   :collection_preview :dataset_query :table_id :query_type :is_upload)
           update-personal-collection))))

src/metabase/models/collection.clj

@@ -67,13 +67,16 @@
   (derive ::mi/read-policy.full-perms-for-perms-set)
   (derive ::mi/write-policy.full-perms-for-perms-set))
 
+(defn- default-audit-collection?
+  [{:keys [id] :as _col}]
+  (= id (:id (perms/default-audit-collection))))
+
 (defmethod mi/can-write? Collection
   ([instance]
-   (mi/can-write? :model/Collection (:id instance)))
-  ([model pk]
-   (if (= pk (:id (perms/default-audit-collection)))
-     false
-     (mi/current-user-has-full-permissions? :write model pk))))
+   (and (not (default-audit-collection? instance))
+        (mi/current-user-has-full-permissions? :write instance)))
+  ([_model pk]
+   (mi/can-write? (t2/select-one :model/Collection pk))))
 
 (defmethod mi/can-read? Collection
   ([instance]