Adding KeySelectors to get-range

.gitignore

@@ -9,3 +9,4 @@ pom.xml.asc
 /.nrepl-port
 .hgignore
 .hg/
+.dir-locals.el

deps.edn

@@ -0,0 +1,18 @@
+{:paths ["src"]
+
+ :deps
+ {org.foundationdb/fdb-java           {:mvn/version "6.3.23"}}
+
+ :aliases
+ {:dev
+  {:extra-paths ["dev"]
+   :extra-deps  {org.clojure/clojure            {:mvn/version "1.11.1"}
+                 com.lambdaisland/classpath     {:mvn/version "0.0.27"}}}
+
+  :test
+  {:extra-paths ["test"]
+   :extra-deps  {org.clj-commons/byte-streams {:mvn/version "0.3.1"}}}}
+
+ :mvn/repos
+ {"central"        {:url "https://repo1.maven.org/maven2/"}
+  "clojars"        {:url "https://clojars.org/repo"}}}

project.clj

@@ -3,10 +3,10 @@
   :url "https://vedang.github.io/clj_fdb/"
   :license {:name "Eclipse Public License"
             :url "http://www.eclipse.org/legal/epl-v10.html"}
-  :dependencies [[org.foundationdb/fdb-java "6.3.13"]]
-  :profiles {:dev {:dependencies [[org.clojure/clojure "1.10.3"]
-                                  [org.clojure/tools.logging "1.1.0"]
-                                  [byte-streams "0.2.4"]]
+  :dependencies [[org.foundationdb/fdb-java "6.3.23"]]
+  :profiles {:test {:dependencies [[org.clj-commons/byte-streams "0.3.1"]]}
+             :dev {:dependencies [[org.clojure/clojure "1.11.1"]
+                                  [org.clojure/tools.logging "1.1.0"]]
                    :jvm-opts ["-Dclojure.tools.logging.factory=clojure.tools.logging.impl/slf4j-factory"]
                    :plugins [[lein-codox "0.10.7"]]}}
   :target-path "target/%s"

src/me/vedang/clj_fdb/core.clj

@@ -1,11 +1,12 @@
 (ns me.vedang.clj-fdb.core
   (:refer-clojure :exclude [get set range])
   (:require [me.vedang.clj-fdb.impl :as fimpl]
-            [me.vedang.clj-fdb.subspace.subspace :as fsub]
+            [me.vedang.clj-fdb.key-selector :as sut]
             [me.vedang.clj-fdb.mutation-type :as fmut]
+            [me.vedang.clj-fdb.subspace.subspace :as fsub]
             [me.vedang.clj-fdb.transaction :as ftr]
             [me.vedang.clj-fdb.tuple.tuple :as ftup])
-  (:import [com.apple.foundationdb KeyValue Range Transaction TransactionContext]
+  (:import [com.apple.foundationdb KeySelector KeyValue Range Transaction TransactionContext]
            com.apple.foundationdb.subspace.Subspace
            com.apple.foundationdb.tuple.Tuple
            java.lang.IllegalArgumentException))
@@ -112,17 +113,32 @@
        (fsub/range s t))
      (range arg2))))
 
+(defn- create-range-args [arg1 arg2 {:keys [skip]}]
+  (if (instance? KeySelector arg1)
+    [(cond-> arg1 skip (sut/add skip)) arg2]
+    [(range arg1 arg2)]))
+
+(defn- get-key-decoder [arg1 arg2]
+  (cond (instance? KeySelector arg1)
+        fimpl/decode
+        (or arg1 (instance? Subspace arg2))
+        (partial fimpl/decode (or arg1 arg2))
+        :else fimpl/decode))
 
 (defn get-range
   "Takes the following:
   - TransactionContext `tc`
-  - Range of keys to fetch `rng` or a Subspace `subspace`
+  - Range of keys to fetch `rng`, Subspace `subspace` or a Keyselector
+  pair of `begin` and `end`.
   - In the case of a `subspace`, can also accept `t`, a Tuple within
   that Subspace
 
   The `opts` map takes the following option at the moment:
   - `keyfn` :
   - `valfn` : Functions to transform the key/value to the correct format.
+  - `limit` : limit the number of returned tuples
+  - `skip`  : skip a number of keys from the beginning of the range
+              (currently only works with KeySelectors)
 
   Note that the byte-arrays are always sent through the `fimpl/decode`
   function first. (So if you have stored a Tuple in FDB, the `valfn`
@@ -134,28 +150,29 @@
   namespaces.
 
   Returns a map of key/value pairs."
-  {:arglists '([tc rnge] [tc subspace] [tc k opts] [tc s k] [tc s k opts])}
+  {:arglists '([tc rnge] [tc subspace] [tc k opts] [tc s k]
+               [tc begin end] [tc s k opts] [tc begin end opts])}
   ([^TransactionContext tc arg1]
    (get-range tc arg1 default-opts))
   ([^TransactionContext tc arg1 arg2]
-   (let [[s k opts] (fimpl/handle-opts arg1 arg2)]
-     (get-range tc s k opts)))
-  ([^TransactionContext tc s k opts]
-   (let [rg (range s k)
-         key-decoder (if (or s (instance? Subspace k))
-                       (partial fimpl/decode (or s k))
-                       fimpl/decode)
-         keyfn (comp (:keyfn opts identity) key-decoder)
-         valfn (comp (:valfn opts identity) fimpl/decode)]
+   (let [[arg1 arg2 opts] (fimpl/handle-opts arg1 arg2)]
+     (get-range tc arg1 arg2 opts)))
+  ([^TransactionContext tc arg1 arg2 {:keys [limit keyfn valfn] :as opts}]
+   (let [range-args (cond-> (create-range-args arg1 arg2 opts)
+                      limit (conj limit))
+         key-decoder (get-key-decoder arg1 arg2)
+         keyfn (cond->> key-decoder
+                 keyfn (comp keyfn))
+         valfn (cond->> fimpl/decode
+                 valfn (comp valfn))]
      (ftr/read tc
                (fn [^Transaction tr]
                  (reduce (fn [acc ^KeyValue kv]
                            (assoc acc
                                   (keyfn (.getKey kv))
                                   (valfn (.getValue kv))))
                          {}
-                         (ftr/get-range tr rg)))))))
-
+                         (apply ftr/get-range tr range-args)))))))
 
 (defn clear-range
   "Takes the following:

src/me/vedang/clj_fdb/internal/util.clj

@@ -1,7 +1,8 @@
 (ns me.vedang.clj-fdb.internal.util
-  (:require [me.vedang.clj-fdb.core :as fc]
+  (:require [me.vedang.clj-fdb.FDB :as cfdb]
+            [me.vedang.clj-fdb.core :as fc]
             [me.vedang.clj-fdb.directory.directory :as fdir]
-            [me.vedang.clj-fdb.FDB :as cfdb]
+            [me.vedang.clj-fdb.range :as frange]
             [me.vedang.clj-fdb.tuple.tuple :as ftup])
   (:import com.apple.foundationdb.Database))
 
@@ -34,3 +35,11 @@
     (binding [*test-prefix* random-prefix]
       (test))
     (clear-all-with-prefix random-prefix)))
+
+(def ^:private smallest-ba (byte-array [(unchecked-byte 0x01)]))
+(def ^:private largest-ba (byte-array [(unchecked-byte 0xff)]))
+
+(defn clear-db
+  "WARNING! This clears the entire db."
+  [db]
+  (fc/clear-range db (frange/range smallest-ba largest-ba)))

src/me/vedang/clj_fdb/subspace/subspace.clj

@@ -5,10 +5,13 @@
            com.apple.foundationdb.subspace.Subspace
            com.apple.foundationdb.tuple.Tuple))
 
+(def ^:private byte-array-class (class (byte-array 0)))
+
 (defn ^Subspace create
   "Constructor for a subspace formed with the specified prefix Tuple."
   ([prefix]
    (cond
+     (instance? byte-array-class prefix) (Subspace. prefix)
      (vector? prefix) (Subspace. (ftup/create prefix))
      (instance? Tuple prefix) (Subspace. ^Tuple prefix)
      :else (throw (IllegalArgumentException.
@@ -41,11 +44,12 @@
   suffix."
   ([^Subspace s]
    (.pack s))
-  ([^Subspace s t]
+  ([^Subspace s k]
    (cond
-     (instance? Tuple t) (.pack s ^Tuple t)
-     (vector? t) (.pack s ^Tuple (ftup/create t))
-     :else (.pack s ^Tuple (ftup/from t)))))
+     (instance? byte-array-class k) (.pack s k)
+     (instance? Tuple k) (.pack s ^Tuple k)
+     (vector? k) (.pack s ^Tuple (ftup/create k))
+     :else (.pack s ^Tuple (ftup/from k)))))
 
 
 (defn ^Tuple unpack

src/me/vedang/clj_fdb/transaction.clj

@@ -1,7 +1,7 @@
 (ns me.vedang.clj-fdb.transaction
   (:refer-clojure :exclude [get set read])
   (:import clojure.lang.IFn
-           [com.apple.foundationdb MutationType Range Transaction TransactionContext]
+           [com.apple.foundationdb MutationType KeySelector Range Transaction TransactionContext]
            com.apple.foundationdb.async.AsyncIterable
            java.util.concurrent.CompletableFuture
            [java.util.function Function Supplier]))
@@ -89,9 +89,16 @@
   "Gets an ordered range of keys and values from the database. The
   begin and end keys are specified by byte[] arrays, with the begin
   key inclusive and the end key exclusive. Ranges are returned from
-  calls to Tuple.range() and Range.startsWith(byte[])."
-  [^Transaction tr ^Range rg]
-  (.getRange tr rg))
+  calls to Tuple.range(), Range.startsWith(byte[]) or can be constructed
+  explicitly. The other option is to use KeySelectors to specify the
+  beginning and end of a Range to return."
+  {:arglists '([tc rnge] [tc rnge limit] [tc begin end] [tc begin end limit])}
+  ([^Transaction tr ^Range rg]
+   (.getRange tr rg))
+  ([^Transaction tr arg1 arg2]
+   (.getRange tr arg1 arg2))
+  ([^Transaction tr ^KeySelector begin ^KeySelector end limit]
+   (.getRange tr begin end limit)))
 
 
 (defn clear-key

test/me/vedang/clj_fdb/core_test.clj

@@ -1,10 +1,11 @@
 (ns me.vedang.clj-fdb.core-test
   (:require [byte-streams :as bs]
             [clojure.test :refer [deftest is testing use-fixtures]]
+            [me.vedang.clj-fdb.FDB :as cfdb]
             [me.vedang.clj-fdb.core :as fc]
             [me.vedang.clj-fdb.directory.directory :as fdir]
-            [me.vedang.clj-fdb.FDB :as cfdb]
             [me.vedang.clj-fdb.internal.util :as u]
+            [me.vedang.clj-fdb.key-selector :as sut]
             [me.vedang.clj-fdb.range :as frange]
             [me.vedang.clj-fdb.subspace.subspace :as fsub]
             [me.vedang.clj-fdb.transaction :as ftr]
@@ -61,6 +62,105 @@
         (is (= expected-map
                (fc/get-range db rg {:keyfn second :valfn bs/to-string})))))))
 
+(deftest get-range-with-range-test
+  ;; foo foo0 ... foo9 as keys
+  (let [input-keys (->> (for [i (range 10)]
+                          (str "foo" i))
+                        (into ["foo"]))
+        v "1"]
+    (with-open [^Database db (cfdb/open fdb)]
+      (ftr/run db
+        (fn [^Transaction tr]
+          (doseq [k input-keys]
+            (let [k (ftup/from u/*test-prefix* k)]
+              (fc/set tr k (bs/to-byte-array v))))))
+
+      (testing "get-range with frange/range"
+        (let [begin      (ftup/pack (ftup/from u/*test-prefix* "foo"))
+              end        (ftup/pack (ftup/from u/*test-prefix* "foo3"))
+              rg         (frange/range begin end)]
+          (is (= {"foo" v "foo0" v "foo1" v "foo2" v}
+                 (fc/get-range db rg {:keyfn second :valfn bs/to-string})))))
+
+      (testing "get-range with frange/range and limit"
+        (let [begin      (ftup/pack (ftup/from u/*test-prefix* "foo"))
+              end        (ftup/pack (ftup/from u/*test-prefix* "foo3"))
+              rg         (frange/range begin end)]
+          (is (= {"foo" v "foo0" v}
+                 (fc/get-range db rg {:keyfn second :valfn bs/to-string :limit 2}))))))))
+
+(deftest get-range-with-subspace-test
+  ;; foo foo0 ... foo9 as keys
+  (let [input-keys (->> (for [i (range 10)]
+                          (str "foo" i))
+                        (into ["foo"]))
+        subspace (fsub/create [u/*test-prefix* "the-store"])
+        v "1"]
+    (with-open [^Database db (cfdb/open fdb)]
+      (ftr/run db
+        (fn [^Transaction tr]
+          (doseq [k input-keys]
+            (let [k (ftup/from k)]
+              (fc/set tr subspace k (bs/to-byte-array v))))))
+
+      (testing "get-range with subspace"
+        (is (=  (zipmap input-keys (repeat v))
+                (fc/get-range db subspace {:keyfn first :valfn bs/to-string}))))
+
+      (testing "get-range with subspace limit"
+        (is (= {"foo" v "foo0" v}
+               (fc/get-range db subspace {:keyfn first :valfn bs/to-string :limit 2}))))
+
+      (testing "get-range with subspace and tuple"
+        (is (= (zipmap input-keys (repeat v))
+               (fc/get-range db subspace (ftup/from) {:keyfn first :valfn bs/to-string}))))
+
+      (testing "get-range with subspace and tuple and limit"
+        (is (= {"foo" v "foo0" v}
+               (fc/get-range db subspace (ftup/from) {:keyfn first :valfn bs/to-string :limit 2})))))))
+
+(defn- third [c] (nth c 2))
+
+(deftest get-range-with-keyselectors-test
+  ;; foo foo0 ... foo9 as keys
+  (let [input-keys (->> (for [i (range 10)]
+                          (str "foo" i))
+                        (into ["foo"]))
+        subspace (fsub/create [u/*test-prefix* "the-store"])
+        v "1"]
+    (with-open [^Database db (cfdb/open fdb)]
+      (ftr/run db
+        (fn [^Transaction tr]
+          (doseq [k input-keys]
+            (let [k (ftup/from k)]
+              (fc/set tr subspace k (bs/to-byte-array v))))))
+
+      (testing "get-range with keyselectors"
+        (is (=  (zipmap (take 6 input-keys) (repeat v))
+                (fc/get-range db
+                              (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo")))
+                              (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo5")))
+                              {:keyfn third :valfn bs/to-string}))))
+
+      (testing "get-range with keyselectors and limit"
+        (is (= {"foo" v "foo0" v}
+               (fc/get-range db
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo")))
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo5")))
+                             {:keyfn third :valfn bs/to-string :limit 2}))))
+
+      (testing "get-range with keyselectors and limit and skip"
+        (is (= {"foo1" v "foo2" v}
+               (fc/get-range db
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo")))
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo5")))
+                             {:keyfn third :valfn bs/to-string :limit 2 :skip 2})))
+        ;; case where skip goes beyond end keyselector
+        (is (= {"foo4" v }
+               (fc/get-range db
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo")))
+                             (sut/first-greater-or-equal (fsub/pack subspace (ftup/from "foo5")))
+                             {:keyfn third :valfn bs/to-string :limit 2 :skip 5})))))))
 
 (deftest clear-range-tests
   (testing "Test the best-case path for `fc/clear-range`. End is exclusive."