From 611d44fcceab96c42140162270d963507a668b04 Mon Sep 17 00:00:00 2001
From: Nick Alexander <nalexander@mozilla.com>
Date: Fri, 9 Sep 2016 15:26:13 -0700
Subject: [PATCH] Process lookup-refs in batches.  Fixes #25.

This uses a common table expression and multiple SQL calls rather than a
temporary table, since transactions with huge numbers of distinct
lookup-refs are likely to be very rare.

We mark lookup-refs with `lookup-ref`, which is a little awkward because
binding `(let [[a v] lookup-ref] ...)` doesn't directly work, but avoids
some ambiguity present in Datomic and DataScript around interpreting
lookup-refs as multiple value lists.  (Which bit the tests in an earlier
version of this patch!)
---
 src/common/datomish/api.cljc              |  2 +
 src/common/datomish/db.cljc               | 79 +++++++++++++++++++++++
 src/common/datomish/db/debug.cljc         |  9 +++
 src/common/datomish/transact.cljc         | 71 +++++++++-----------
 src/common/datomish/transact/explode.cljc |  1 +
 test/datomish/db_test.cljc                | 59 ++++++++++++++++-
 6 files changed, 178 insertions(+), 43 deletions(-)

diff --git a/src/common/datomish/api.cljc b/src/common/datomish/api.cljc
index 6ceb93e4..567335e0 100644
--- a/src/common/datomish/api.cljc
+++ b/src/common/datomish/api.cljc
@@ -32,6 +32,8 @@
 
 (def id-literal db/id-literal)
 
+(def lookup-ref db/lookup-ref)
+
 (def db transact/db)
 
 (def entid db/entid)
diff --git a/src/common/datomish/db.cljc b/src/common/datomish/db.cljc
index 005ece4e..fd926685 100644
--- a/src/common/datomish/db.cljc
+++ b/src/common/datomish/db.cljc
@@ -63,6 +63,23 @@
 (defn id-literal? [x]
   (instance? TempId x))
 
+(defrecord LookupRef [a v])
+
+(defn lookup-ref
+  [a v]
+  (if (and
+        (or (keyword? a)
+            (integer? a))
+        v)
+    (->LookupRef a v)
+    (raise (str "Lookup-ref with bad attribute " a " or value " v
+                {:error :transact/bad-lookup-ref, :a a, :v v}))))
+
+(defn lookup-ref? [x]
+  "Return `x` if `x` is like [:attr value], nil otherwise."
+  (when (instance? LookupRef x)
+    x))
+
 (defprotocol IClock
   (now
     [clock]
@@ -105,6 +122,13 @@
     [db a v]
     "Search for a single matching datom using the AVET index.")
 
+  (<avs
+    [db avs]
+    "Search for many matching datoms using the AVET index.
+
+    Take [[a0 v0] [a1 v1] ...] and return a map {[a0 v0] e0}.  If no datom [e1 a1 v1] exists, the
+    key [a1 v1] is not present in the returned map.")
+
   (<apply-entities
     [db tx entities]
     "Apply entities to the store, returning sequence of datoms transacted.")
@@ -634,6 +658,61 @@
           <?
           yield-datom))))
 
+  (<avs
+    [db avs]
+    {:pre [(sequential? avs)]}
+
+    (go-pair
+      (let [schema
+            (.-schema db)
+
+            values-part
+            "(?, ?, ?, ?)"
+
+            repeater
+            (memoize (fn [n] (interpose ", " (repeat n values-part))))
+
+            exec
+            (partial s/execute! (:sqlite-connection db))
+
+            ;; Map [a v] -> searchid.
+            av->searchid
+            (into {} (map vector avs (range)))
+
+            ;; Each query takes 4 variables per item. So we partition into max-sql-vars / 4.
+            qs
+            (map
+              (fn [chunk]
+                (cons
+                  ;; Query string.
+                  (apply str "WITH t(searchid, a, v, value_type_tag) AS (VALUES "
+                         (apply str (repeater (count chunk))) ;; TODO: join?
+                         ") SELECT t.searchid, d.e
+                           FROM t, datoms AS d
+                           WHERE d.index_avet IS NOT 0 AND d.a = t.a AND d.value_type_tag = t.value_type_tag AND d.v = t.v")
+
+                  ;; Bindings.
+                  (mapcat (fn [[[a v] searchid]]
+                            (let [a (entid db a)
+                                  [v tag] (ds/->SQLite schema a v)]
+                              [searchid a v tag]))
+                          chunk)))
+
+              (partition-all (quot max-sql-vars 4) av->searchid))
+
+            ;; Map searchid -> e.  There's a generic reduce that takes [pair-chan] lurking in here.
+            searchid->e
+            (loop [coll {}
+                   qs qs]
+              (let [[q & qs] qs]
+                (if q
+                  (let [rs (<? (s/all-rows (:sqlite-connection db) q))
+                        coll* (into coll (map (juxt :searchid :e)) rs)]
+                    (recur coll* qs))
+                  coll)))
+            ]
+        (util/mapvals (partial get searchid->e) av->searchid))))
+
   (<apply-entities [db tx entities]
     {:pre [(db? db) (sequential? entities)]}
     (-<apply-entities db tx entities))
diff --git a/src/common/datomish/db/debug.cljc b/src/common/datomish/db/debug.cljc
index 09145048..9503f64f 100644
--- a/src/common/datomish/db/debug.cljc
+++ b/src/common/datomish/db/debug.cljc
@@ -29,6 +29,15 @@
       (filter #(not (= :db/txInstant (second %))))
       (set))))
 
+(defn <datoms>= [db tx]
+  (go-pair
+    (->>
+      (s/all-rows (:sqlite-connection db) ["SELECT e, a, v, tx FROM datoms WHERE tx >= ?" tx])
+      (<?)
+      (mapv #(vector (:e %) (db/ident db (:a %)) (:v %)))
+      (filter #(not (= :db/txInstant (second %))))
+      (set))))
+
 (defn <datoms [db]
   (<datoms-after db 0))
 
diff --git a/src/common/datomish/transact.cljc b/src/common/datomish/transact.cljc
index 683ee431..af963a45 100644
--- a/src/common/datomish/transact.cljc
+++ b/src/common/datomish/transact.cljc
@@ -127,7 +127,7 @@
             (db/entid db v))]
     (when (and a (not (integer? a)))
       (raise "Unknown attribute " a
-             {:form orig :attribute a}))
+             {:form orig :attribute a :entity orig}))
     [op e a v]))
 
 (defrecord Transaction [db tempids entities])
@@ -250,49 +250,40 @@
       ;; Extract the current txInstant for the report.
       (->> (update-txInstant db*)))))
 
-(defn- lookup-ref? [x]
-  "Return `x` if `x` is like [:attr value], false otherwise."
-  (and (sequential? x)
-       (= (count x) 2)
-       (or (keyword? (first x))
-           (integer? (first x)))
-       x))
-
 (defn <resolve-lookup-refs [db report]
   {:pre [(db/db? db) (report? report)]}
 
-  (let [entities (:entities report)]
-    ;; TODO: meta.
+  (let [unique-identity? (memoize (partial ds/unique-identity? (db/schema db)))
+        groups           (group-by (partial keep db/lookup-ref?) (:entities report))
+        entities         (get groups (lazy-seq)) ;; No lookup-refs? Pass through.
+        to-resolve       (dissoc groups (lazy-seq)) ;; The ones with lookup-refs.
+        avs              (set (map (juxt :a :v) (apply concat (keys to-resolve))))
+        ->av             (fn [r] ;; Conditional (juxt :a :v) that passes through nil.
+                           (when r [(:a r) (:v r)]))]
     (go-pair
-      (if (empty? entities)
-        report
-        (assoc-in
-          report [:entities]
-          ;; We can't use `for` because go-pair doesn't traverse function boundaries.
-          ;; Apologies for the tortured nested loop.
-          (loop [[op & entity] (first entities)
-                 next (rest entities)
-                 acc []]
-            (if (nil? op)
-              acc
-              (recur (first next)
-                     (rest next)
-                     (conj acc
-                           (loop [field (first entity)
-                                  rem (rest entity)
-                                  acc [op]]
-                             (if (nil? field)
-                               acc
-                               (recur (first rem)
-                                      (rest rem)
-                                      (conj acc
-                                            (if-let [[a v] (lookup-ref? field)]
-                                              (or
-                                                ;; The lookup might fail! If so, throw.
-                                                (:e (<? (db/<av db a v)))
-                                                (raise "No entity found with attr " a " and val " v "."
-                                                       {:a a :v v}))
-                                              field))))))))))))))
+      (let [av->e    (<? (db/<avs db avs))
+            resolve1 (fn [field]
+                       (if-let [[a v] (->av (db/lookup-ref? field))]
+                         (if-not (unique-identity? (db/entid db a))
+                           (raise "Lookup-ref found with non-unique-identity attribute " a " and value " v
+                                  {:error :transact/lookup-ref-with-non-unique-identity-attribute
+                                   :a     a
+                                   :v     v})
+                           (or
+                             (get av->e [a v])
+                             (raise "No entity found for lookup-ref with attribute " a " and value " v
+                                    {:error :transact/lookup-ref-not-found
+                                     :a     a
+                                     :v     v})))
+                         field))
+            resolve  (fn [entity]
+                       (mapv resolve1 entity))]
+        (assoc
+          report
+          :entities
+          (concat
+            entities
+            (map resolve (apply concat (vals to-resolve)))))))))
 
 (declare <resolve-id-literals)
 
diff --git a/src/common/datomish/transact/explode.cljc b/src/common/datomish/transact/explode.cljc
index 4bcbe805..c07c99f0 100644
--- a/src/common/datomish/transact/explode.cljc
+++ b/src/common/datomish/transact/explode.cljc
@@ -47,6 +47,7 @@
       (explode-entity-a-v db entity v straight-a eid)
 
       (and (map? v)
+           (not (db/lookup-ref? v))
            (not (db/id-literal? v)))
       ;; Another entity is given as a nested map.
       (if (ds/ref? (db/schema db) straight-a*)
diff --git a/test/datomish/db_test.cljc b/test/datomish/db_test.cljc
index 1dea8c66..c1ca84cb 100644
--- a/test/datomish/db_test.cljc
+++ b/test/datomish/db_test.cljc
@@ -10,7 +10,7 @@
       [cljs.core.async.macros :as a :refer [go]]))
   (:require
    [datomish.api :as d]
-   [datomish.db.debug :refer [<datoms-after <transactions-after <shallow-entity <fulltext-values]]
+   [datomish.db.debug :refer [<datoms-after <datoms>= <transactions-after <shallow-entity <fulltext-values]]
    [datomish.util :as util #?(:cljs :refer-macros :clj :refer) [raise cond-let]]
    [datomish.schema :as ds]
    [datomish.simple-schema]
@@ -846,8 +846,61 @@
                    :db.install/_attribute :db.part/db}]]
 
     (testing "Simple schemas are expanded."
-             (is (= (map #(dissoc %1 :db/id) (datomish.simple-schema/simple-schema->schema in))
-                    expected)))))
+      (is (= (map #(dissoc %1 :db/id) (datomish.simple-schema/simple-schema->schema in))
+             expected)))))
 
+(deftest-db test-lookup-refs conn
+  (let [{tx0 :tx} (<? (d/<transact! conn test-schema))
+        {tx1 :tx} (<? (d/<transact! conn [[:db/add 1 :name "Ivan"]
+                                          [:db/add 2 :name "Phil"]
+                                          [:db/add 3 :name "Petr"]]))]
+    (testing "Looks up entity refs"
+      (let [{tx :tx} (<? (d/<transact! conn [[:db/add (d/lookup-ref :name "Ivan") :aka "Devil"]
+                                             [:db/add (d/lookup-ref :name "Phil") :email "@1"]]))]
+        (is (= #{[1 :name "Ivan"]
+                 [2 :name "Phil"]
+                 [3 :name "Petr"]
+                 [1 :aka "Devil"]
+                 [2 :email "@1"]}
+               (<? (<datoms>= (d/db conn) tx1))))))
+
+    (testing "Looks up value refs"
+      (let [{tx :tx} (<? (d/<transact! conn [[:db/add 1 :friends (d/lookup-ref :name "Petr")]
+                                             [:db/add 3 :friends (d/lookup-ref :name "Ivan")]]))]
+        (is (= #{[1 :friends 3]
+                 [3 :friends 1]}
+               (<? (<datoms>= (d/db conn) tx))))))
+
+    (testing "Looks up entity refs in maps"
+      (let [{tx :tx} (<? (d/<transact! conn [{:db/id (d/lookup-ref :name "Phil") :friends 1}]))]
+        (is (= #{[2 :friends 1]}
+               (<? (<datoms>= (d/db conn) tx))))))
+
+    (testing "Looks up value refs in maps"
+      (let [{tx :tx} (<? (d/<transact! conn [{:db/id 2 :friends (d/lookup-ref :name "Petr")}]))]
+        (is (=  #{[2 :friends 3]}
+                (<? (<datoms>= (d/db conn) tx))))))
+
+    (testing "Looks up value refs in sequences in maps"
+      (let [{tx :tx} (<? (d/<transact! conn [{:db/id 1 :friends [(d/lookup-ref :name "Ivan") (d/lookup-ref :name "Phil")]}]))]
+        (is (= #{[1 :friends 1]
+                 [1 :friends 2]}
+               (<? (<datoms>= (d/db conn) tx))))))
+
+    (testing "Fails for missing entities"
+      (is (thrown-with-msg?
+            ExceptionInfo #"No entity found for lookup-ref"
+            (<? (d/<transact! conn [[:db/add (d/lookup-ref :name "Mysterioso") :aka "The Magician"]]))))
+      (is (thrown-with-msg?
+            ExceptionInfo #"No entity found for lookup-ref"
+            (<? (d/<transact! conn [[:db/add 1 :friends (d/lookup-ref :name "Mysterioso")]])))))
+
+    (testing "Fails for non-identity attributes"
+      (is (thrown-with-msg?
+            ExceptionInfo #"Lookup-ref found with non-unique-identity attribute"
+            (<? (d/<transact! conn [[:db/add (d/lookup-ref :aka "The Magician") :email "@2"]]))))
+      (is (thrown-with-msg?
+            ExceptionInfo #"Lookup-ref found with non-unique-identity attribute"
+            (<? (d/<transact! conn [[:db/add 1 :friends (d/lookup-ref :aka "The Magician")]])))))))
 
 #_ (time (t/run-tests))