Merge pull request #3279 from metabase/more-efficient-hydration

More efficient batched hydration :car: 

src/metabase/api/database.clj

@@ -149,7 +149,6 @@
   [id]
   (->404 (Database id)
          read-check
-         ;; TODO - this is a bit slow due to the nested hydration.  needs some optimizing.
          (hydrate [:tables [:fields :target :values] :segments :metrics])))
 
 (defendpoint GET "/:id/autocomplete_suggestions"

src/metabase/models/field.clj

@@ -90,18 +90,44 @@
   (db/cascade-delete! 'FieldValues :field_id id)
   (db/cascade-delete! 'MetricImportantField :field_id id))
 
-(defn ^:hydrate target
+
+(defn target
   "Return the FK target `Field` that this `Field` points to."
   [{:keys [special_type fk_target_field_id]}]
   (when (and (= :fk special_type)
              fk_target_field_id)
     (Field fk_target_field_id)))
 
-(defn ^:hydrate values
+(defn values
   "Return the `FieldValues` associated with this FIELD."
   [{:keys [id]}]
   (db/select [FieldValues :field_id :values], :field_id id))
 
+(defn with-values
+  "Efficiently hydrate the `FieldValues` for a collection of FIELDS."
+  {:batched-hydrate :values}
+  [fields]
+  (let [field-ids        (set (map :id fields))
+        id->field-values (u/key-by :field_id (when (seq field-ids)
+                                               (db/select FieldValues :field_id [:in field-ids])))]
+    (for [field fields]
+      (assoc field :values (get id->field-values (:id field) [])))))
+
+(defn with-targets
+  "Efficiently hydrate the FK target fields for a collection of FIELDS."
+  {:batched-hydrate :target}
+  [fields]
+  (let [target-field-ids (set (for [field fields
+                                    :when (and (= :fk (:special_type field))
+                                               (:fk_target_field_id field))]
+                                (:fk_target_field_id field)))
+        id->target-field (u/key-by :id (when (seq target-field-ids)
+                                         (db/select Field :id [:in target-field-ids])))]
+    (for [field fields
+          :let  [target-id (:fk_target_field_id field)]]
+      (assoc field :target (id->target-field target-id)))))
+
+
 (defn qualified-name-components
   "Return the pieces that represent a path to FIELD, of the form `[table-name parent-fields-name* field-name]`."
   [{field-name :name, table-id :table_id, parent-id :parent_id}]

src/metabase/models/metric.clj

@@ -107,8 +107,7 @@
   ([table-id]
    (retrieve-metrics table-id :active))
   ([table-id state]
-   {:pre [(integer? table-id)
-          (keyword? state)]}
+   {:pre [(integer? table-id) (keyword? state)]}
    (-> (if (= :all state)
          (db/select Metric, :table_id table-id, {:order-by [[:name :asc]]})
          (db/select Metric, :table_id table-id, :is_active (= :active state), {:order-by [[:name :asc]]}))

src/metabase/models/table.clj

@@ -17,6 +17,9 @@
   "Valid values for `Table.visibility_type` (field may also be `nil`)."
   #{:hidden :technical :cruft})
 
+
+;;; ------------------------------------------------------------ Entity ------------------------------------------------------------
+
 (i/defentity Table :metabase_table)
 
 (defn- pre-insert [table]
@@ -29,23 +32,36 @@
   (db/cascade-delete! Field :table_id id)
   (db/cascade-delete! 'Card :table_id id))
 
-(defn ^:hydrate fields
-  "Return the `FIELDS` belonging to TABLE."
+(u/strict-extend (class Table)
+  i/IEntity (merge i/IEntityDefaults
+                   {:hydration-keys     (constantly [:table])
+                    :types              (constantly {:entity_type :keyword, :visibility_type :keyword, :description :clob})
+                    :timestamped?       (constantly true)
+                    :can-read?          (constantly true)
+                    :can-write?         i/superuser?
+                    :pre-insert         pre-insert
+                    :pre-cascade-delete pre-cascade-delete}))
+
+
+;;; ------------------------------------------------------------ Hydration ------------------------------------------------------------
+
+(defn fields
+  "Return the `FIELDS` belonging to a single TABLE."
   [{:keys [id]}]
-  (db/select Field, :table_id id, :visibility_type [:not= "retired"], {:order-by [[:position :asc] [:name :asc]]}))
+  (db/select Field, :table_id id :visibility_type [:not= "retired"], {:order-by [[:position :asc] [:name :asc]]}))
 
-(defn ^:hydrate metrics
-  "Retrieve the metrics for TABLE."
+(defn metrics
+  "Retrieve the `Metrics` for a single TABLE."
   [{:keys [id]}]
   (retrieve-metrics id :all))
 
-(defn ^:hydrate segments
-  "Retrieve the segments for TABLE."
+(defn segments
+  "Retrieve the `Segments` for a single TABLE."
   [{:keys [id]}]
   (retrieve-segments id :all))
 
 (defn field-values
-  "Return the `FieldValues` for all `Fields` belonging to TABLE."
+  "Return the `FieldValues` for all `Fields` belonging to a single TABLE."
   {:hydrate :field_values, :arglists '([table])}
   [{:keys [id]}]
   (let [field-ids (db/select-ids Field
@@ -61,6 +77,44 @@
   [{:keys [id]}]
   (db/select-one-id Field, :table_id id, :special_type "id", :visibility_type [:not-in ["sensitive" "retired"]]))
 
+
+(defn- with-objects [hydration-key fetch-objects-fn tables]
+  (let [table-ids         (set (map :id tables))
+        table-id->objects (group-by :table_id (when (seq table-ids)
+                                                (fetch-objects-fn table-ids)))]
+    (for [table tables]
+      (assoc table hydration-key (get table-id->objects (:id table) [])))))
+
+(defn with-segments
+  "Efficiently hydrate the `Segments` for a collection of TABLES."
+  {:batched-hydrate :segments}
+  [tables]
+  (with-objects :segments
+    (fn [table-ids]
+      (db/select Segment :table_id [:in table-ids], {:order-by [[:name :asc]]}))
+    tables))
+
+(defn with-metrics
+  "Efficiently hydrate the `Metrics` for a collection of TABLES."
+  {:batched-hydrate :metrics}
+  [tables]
+  (with-objects :metrics
+    (fn [table-ids]
+      (db/select Metric :table_id [:in table-ids], {:order-by [[:name :asc]]}))
+    tables))
+
+(defn with-fields
+  "Efficiently hydrate the `Fields` for a collection of TABLES."
+  {:batched-hydrate :fields}
+  [tables]
+  (with-objects :fields
+    (fn [table-ids]
+      (db/select Field :table_id [:in table-ids], :visibility_type [:not= "retired"], {:order-by [[:position :asc] [:name :asc]]}))
+    tables))
+
+
+;;; ------------------------------------------------------------ Convenience Fns ------------------------------------------------------------
+
 (defn qualified-identifier
   "Return a keyword identifier for TABLE in the form `:schema.table-name` (if the Table has a non-empty `:schema` field) or `:table-name` (if the Table has no `:schema`)."
   ^clojure.lang.Keyword [{schema :schema, table-name :name}]
@@ -73,19 +127,6 @@
   [table]
   (Database (:db_id table)))
 
-(u/strict-extend (class Table)
-  i/IEntity (merge i/IEntityDefaults
-                   {:hydration-keys     (constantly [:table])
-                    :types              (constantly {:entity_type :keyword, :visibility_type :keyword, :description :clob})
-                    :timestamped?       (constantly true)
-                    :can-read?          (constantly true)
-                    :can-write?         i/superuser?
-                    :pre-insert         pre-insert
-                    :pre-cascade-delete pre-cascade-delete}))
-
-
-;; ## Persistence Functions
-
 (defn table-id->database-id
   "Retrieve the `Database` ID for the given table-id."
   [table-id]
@@ -93,6 +134,10 @@
   (db/select-one-field :db_id Table, :id table-id))
 
 
+;;; ------------------------------------------------------------ Persistence Functions ------------------------------------------------------------
+
+
+
 (defn retire-tables!
   "Retire all `Tables` in the list of TABLE-IDs along with all of each tables `Fields`."
   [table-ids]

test/metabase/models/hydrate_test.clj

@@ -8,27 +8,39 @@
             [metabase.test.data.users :refer :all]
             [metabase.util :as u]))
 
-;; we'll swap out hydration-key-> with one that will track the keys used, and we'll warn about unused ones at the end
+;; we'll swap out hydration-key->f with one that will track the keys used, and we'll warn about unused ones at the end
 
-(def ^:private used-fn-keys (atom #{}))
-
-(let [hydration-key->f @(resolve 'metabase.models.hydrate/hydration-key->f)]
+(let [used-fn-keys     (atom #{})
+      hydration-key->f @(resolve 'metabase.models.hydrate/hydration-key->f)]
   (intern 'metabase.models.hydrate 'hydration-key->f
           (fn [k]
             (swap! used-fn-keys conj k)
-            (hydration-key->f k))))
-
-(defn- check-all-hydration-keys-used
-  {:expectations-options :after-run}
-  []
-  (let [hydrate-k->f   @@(resolve 'metabase.models.hydrate/k->f)
-        available-keys (set (keys hydrate-k->f))
-        used-keys      @used-fn-keys
-        unused-keys    (set/difference available-keys used-keys)]
-    (when (seq unused-keys)
+            (hydration-key->f k)))
+  (defn- check-all-hydration-keys-used
+    {:expectations-options :after-run}
+    []
+    (let [hydrate-k->f   @@(resolve 'metabase.models.hydrate/k->f)
+          available-keys (set (keys hydrate-k->f))
+          used-keys      @used-fn-keys
+          unused-keys    (set/difference available-keys used-keys)]
+      (doseq [k unused-keys]
+        (println (u/format-color 'red "WARNING: %s is marked `^:hydrate` but is never used for hydration." (hydrate-k->f k)))))))
+
+(let [used-fn-keys             (atom #{})
+      hydration-key->batched-f @(resolve 'metabase.models.hydrate/hydration-key->batched-f)]
+  (intern 'metabase.models.hydrate 'hydration-key->batched-f
+          (fn [k]
+            (swap! used-fn-keys conj k)
+            (hydration-key->batched-f k)))
+  (defn- check-all-batched-hydration-keys-used
+    {:expectations-options :after-run}
+    []
+    (let [hydrate-k->f   @@(resolve 'metabase.models.hydrate/k->batched-f)
+          available-keys (set (keys hydrate-k->f))
+          used-keys      @used-fn-keys
+          unused-keys    (set/difference available-keys used-keys)]
       (doseq [k unused-keys]
-        (println (u/format-color 'red "%s is marked `^:hydrate` but is never used for hydration." (hydrate-k->f k))))
-      (System/exit -1))))
+        (println (u/format-color 'red "WARNING: %s is marked `^:batched-hydrate` but is never used for hydration." (hydrate-k->f k)))))))
 
 
 (defn- ^:hydrate x [{:keys [id]}]
@@ -52,20 +64,20 @@
 (expect :toucan_id
   (k->k_id :toucan))
 
-;; ### can-batched-hydrate?
-(def can-batched-hydrate? (ns-resolve 'metabase.models.hydrate 'can-batched-hydrate?))
+;; ### can-automagically-batched-hydrate?
+(def can-automagically-batched-hydrate? (ns-resolve 'metabase.models.hydrate 'can-automagically-batched-hydrate?))
 
 ;; should fail for unknown keys
 (expect false
-  (can-batched-hydrate? [{:a_id 1} {:a_id 2}] :a))
+  (can-automagically-batched-hydrate? [{:a_id 1} {:a_id 2}] :a))
 
 ;; should work for known keys if k_id present in every map
 (expect true
-  (can-batched-hydrate? [{:user_id 1} {:user_id 2}] :user))
+  (can-automagically-batched-hydrate? [{:user_id 1} {:user_id 2}] :user))
 
 ;; should fail for known keys if k_id isn't present in every map
 (expect false
-  (can-batched-hydrate? [{:user_id 1} {:user_id 2} {:x 3}] :user))
+  (can-automagically-batched-hydrate? [{:user_id 1} {:user_id 2} {:x 3}] :user))
 
 ;; ### valid-hydration-form?
 (def valid-hydration-form? (ns-resolve 'metabase.models.hydrate 'valid-hydration-form?))