Add Estimated Number of Keys Function

We need this functionality for all key-value instead of only for
RocksDB because we will use it in the CQL expression cache.

Needed in: #1051

modules/admin-api/src/blaze/admin_api.clj

@@ -4,6 +4,7 @@
    [blaze.admin-api.validation]
    [blaze.anomaly :as ba :refer [if-ok]]
    [blaze.async.comp :as ac :refer [do-sync]]
+   [blaze.db.kv :as kv]
    [blaze.db.kv.rocksdb :as rocksdb]
    [blaze.elm.expression :as-alias expr]
    [blaze.fhir.response.create :as create-response]
@@ -86,7 +87,7 @@
 (defn- column-family-data [db column-family]
   (let [long-property (partial rocksdb/long-property db column-family)]
     {:name (name column-family)
-     :estimate-num-keys (long-property "rocksdb.estimate-num-keys")
+     :estimate-num-keys (kv/estimate-num-keys db column-family)
      :estimate-live-data-size (long-property "rocksdb.estimate-live-data-size")
      :live-sst-files-size (long-property "rocksdb.live-sst-files-size")
      :size-all-mem-tables (long-property "rocksdb.size-all-mem-tables")}))

modules/kv/src/blaze/db/kv.clj

@@ -183,3 +183,10 @@
   Writes are atomic. Blocks."
   [store entries]
   (p/-write store entries))
+
+(defn estimate-num-keys
+  "Returns the estimated number of keys in `column-family` of `store`.
+
+  Returns an anomaly if the column-family was not found."
+  [store column-family]
+  (p/-estimate-num-keys store column-family))

modules/kv/src/blaze/db/kv/mem.clj

@@ -110,14 +110,17 @@
     (set! closed? true)))
 
 (defn- column-family-not-found-msg [column-family]
-  (format "column family `%s` not found" (name column-family)))
+  (format "Column family `%s` not found." (name column-family)))
+
+(defn- column-family-not-found-anom [column-family]
+  (ba/not-found (column-family-not-found-msg column-family)))
 
 (deftype MemKvSnapshot [db]
   p/KvSnapshot
   (-new-iterator [_ column-family]
     (if-let [db (get db column-family)]
       (->MemKvIterator db (atom {:rest (seq db)}) false)
-      (throw-anom (ba/not-found (column-family-not-found-msg column-family)))))
+      (throw-anom (column-family-not-found-anom column-family))))
 
   (-snapshot-get [_ column-family k]
     (some-> (get-in db [column-family k]) (copy)))
@@ -127,7 +130,7 @@
 
 (defn- assoc-copy [m column-family k v]
   (when (nil? m)
-    (throw-anom (ba/not-found (column-family-not-found-msg column-family))))
+    (throw-anom (column-family-not-found-anom column-family)))
   (assoc m (copy k) (copy v)))
 
 (defn- put-entries [db entries]
@@ -172,7 +175,12 @@
 
   (-write [_ entries]
     (swap! db write-entries entries)
-    nil))
+    nil)
+
+  (-estimate-num-keys [_ column-family]
+    (if-let [m (get @db column-family)]
+      (count m)
+      (column-family-not-found-anom column-family))))
 
 (def ^:private bytes-cmp
   (reify Comparator

modules/kv/src/blaze/db/kv/protocols.clj

@@ -43,4 +43,6 @@
 
   (-delete [store entries])
 
-  (-write [store entries]))
+  (-write [store entries])
+
+  (-estimate-num-keys [store column-family]))

modules/kv/src/blaze/db/kv_spec.clj

@@ -4,7 +4,8 @@
    [blaze.coll.spec :as cs]
    [blaze.db.kv :as kv]
    [blaze.db.kv.spec]
-   [clojure.spec.alpha :as s]))
+   [clojure.spec.alpha :as s]
+   [cognitect.anomalies :as anom]))
 
 (s/fdef kv/valid?
   :args (s/cat :iter ::kv/iterator)
@@ -79,3 +80,7 @@
 (s/fdef kv/write!
   :args (s/cat :kv-store :blaze.db/kv-store
                :entries (cs/coll-of ::kv/write-entry)))
+
+(s/fdef kv/estimate-num-keys
+  :args (s/cat :kv-store :blaze.db/kv-store :column-family simple-keyword?)
+  :ret (s/or :estimate-num-keys nat-int? :anomaly ::anom/anomaly))

modules/kv/test/blaze/db/kv/mem_test.clj

@@ -15,13 +15,14 @@
    [clojure.test :as test :refer [deftest is testing]]
    [cognitect.anomalies :as anom]
    [integrant.core :as ig]
+   [juxt.iota :refer [given]]
    [taoensso.timbre :as log])
   (:import
    [java.lang AutoCloseable]))
 
 (set! *warn-on-reflection* true)
 (st/instrument)
-(log/set-level! :trace)
+(log/set-min-level! :trace)
 
 (test/use-fixtures :each tu/fixture)
 
@@ -576,5 +577,18 @@
       (kv/write! kv-store [[:delete :default (ba 0x00)]])
       (is (nil? (kv/get kv-store :default (ba 0x00)))))))
 
+(deftest estimate-num-keys-test
+  (with-system [{kv-store ::kv/mem} config]
+    (is (zero? (kv/estimate-num-keys kv-store :default)))
+
+    (given (kv/estimate-num-keys kv-store :foo)
+      ::anom/category := ::anom/not-found
+      ::anom/message := "Column family `foo` not found."))
+
+  (with-system-data [{kv-store ::kv/mem} config]
+    [[:default (ba 0x00) (ba 0x10)]]
+
+    (is (= 1 (kv/estimate-num-keys kv-store :default)))))
+
 (deftest init-component-test
   (is (kv/store? (ig/init-key ::kv/mem {}))))

modules/rocksdb/src/blaze/db/kv/rocksdb.clj

@@ -206,6 +206,9 @@
       (impl/write-wb! cfhs wb entries)
       (.write db write-opts wb)))
 
+  (-estimate-num-keys [store column-family]
+    (p/-long-property store column-family "rocksdb.estimate-num-keys"))
+
   p/Rocks
   (-path [_]
     path)

modules/rocksdb/test/blaze/db/kv/rocksdb_test.clj

@@ -27,7 +27,7 @@
 
 (set! *warn-on-reflection* true)
 (st/instrument)
-(log/set-level! :trace)
+(log/set-min-level! :trace)
 
 (test/use-fixtures :each tu/fixture)
 
@@ -624,6 +624,19 @@
       (kv/write! db [[:delete :default (ba 0x00)]])
       (is (nil? (kv/get db :default (ba 0x00)))))))
 
+(deftest estimate-num-keys-test
+  (with-system [{db ::kv/rocksdb} (config (new-temp-dir!))]
+    (is (zero? (kv/estimate-num-keys db :default)))
+
+    (given (kv/estimate-num-keys db :foo)
+      ::anom/category := ::anom/not-found
+      ::anom/message := "Column family `foo` not found."))
+
+  (with-system-data [{db ::kv/rocksdb} (config (new-temp-dir!))]
+    [[:default (ba 0x00) (ba 0x10)]]
+
+    (is (= 1 (kv/estimate-num-keys db :default)))))
+
 (deftest path-test
   (with-system [{db ::kv/rocksdb} (config (new-temp-dir!))]
     (is (string? (rocksdb/path db)))))