Fix removeRange in InMemoryProvider

package.json

@@ -1,6 +1,6 @@
 {
   "name": "@microsoft/objectstoreprovider",
-  "version": "0.8.1",
+  "version": "0.8.2",
   "description": "A cross-browser object store library",
   "author": "Mukundan Kavanur Kidambi <[email protected]>",
   "scripts": {

src/InMemoryProvider.ts

@@ -64,12 +64,14 @@ export interface StoreData {
 
 export interface ILiveConsumerConfigs {
   usePushForGetRange: boolean;
+  usePrimaryKeyForGetKeysForRange: boolean;
 }
 
 export type GetLiveConsumerConfigsFn = () => ILiveConsumerConfigs;
 
 const defaultLiveConsumerConfigs: ILiveConsumerConfigs = {
   usePushForGetRange: false,
+  usePrimaryKeyForGetKeysForRange: false,
 };
 
 export class InMemoryProvider extends DbProvider {
@@ -915,23 +917,51 @@ class InMemoryIndex extends DbIndexFTSFromRangeQueries {
     lowRangeExclusive?: boolean,
     highRangeExclusive?: boolean
   ): string[] {
+    const usePrimaryKey = this.getLiveConfigs().usePrimaryKeyForGetKeysForRange;
     const keyLow = serializeKeyToString(keyLowRange, this._keyPath);
     const keyHigh = serializeKeyToString(keyHighRange, this._keyPath);
     const iterator = this._indexTree.entries();
-    const keys = [];
+    const keys: string[] = [];
+
     for (const entry of iterator) {
-      const key = entry.key;
+      let key = entry.key;
       if (key === undefined) {
         continue;
       }
 
-      if (
+      const isMatch =
         (key > keyLow || (key === keyLow && !lowRangeExclusive)) &&
-        (key < keyHigh || (key === keyHigh && !highRangeExclusive))
+        (key < keyHigh || (key === keyHigh && !highRangeExclusive));
+      if (!isMatch) {
+        continue;
+      }
+
+      // If the current index is not the primary one. We need to find primary key for each value and return that instead.
+      if (
+        usePrimaryKey &&
+        entry.value &&
+        this._keyPath !== this._primaryKeyPath
       ) {
+        for (const value of entry.value) {
+          key = getSerializedKeyForKeypath(value, this._primaryKeyPath) ?? key;
+
+          if (key === entry.key) {
+            this.logger.warn(
+              `getSerializedKeyForKeypath returned undefined key in InMemoryIndex for table: ${this.tableName}, with index: ${this._indexSchema?.name}`
+            );
+          }
+
+          if (!keys.includes(key)) {
+            keys.push(key);
+          }
+        }
+      }
+      // Otherwise, we can just use the key as is.
+      else {
         keys.push(key);
       }
     }
+
     return keys;
   }
 

src/tests/ObjectStoreProvider.spec.ts

@@ -35,10 +35,26 @@ function openProvider(
 
   switch (providerName) {
     case "memory-rbtree":
-      provider = new InMemoryProvider("red-black-tree", supportsRollback);
+      provider = new InMemoryProvider(
+        "red-black-tree",
+        supportsRollback,
+        undefined,
+        () => ({
+          usePushForGetRange: false,
+          usePrimaryKeyForGetKeysForRange: true,
+        })
+      );
       break;
     case "memory-btree":
-      provider = new InMemoryProvider("b+tree", supportsRollback);
+      provider = new InMemoryProvider(
+        "b+tree",
+        supportsRollback,
+        undefined,
+        () => ({
+          usePushForGetRange: false,
+          usePrimaryKeyForGetKeysForRange: true,
+        })
+      );
       break;
     case "indexeddb":
       provider = new IndexedDbProvider();
@@ -1886,6 +1902,140 @@ describe("ObjectStoreProvider", function () {
             );
         });
 
+        it("Remove range (all removed with index)", (done) => {
+          // Not working with index, need to be fix in the future.
+          if (provName === "indexeddbfakekeys") {
+            done();
+            return;
+          }
+
+          openProvider(
+            provName,
+            {
+              version: 1,
+              stores: [
+                {
+                  name: "test",
+                  primaryKeyPath: "id",
+                  indexes: [
+                    {
+                      name: "index",
+                      keyPath: "a",
+                    },
+                  ],
+                },
+              ],
+            },
+            true
+          )
+            .then((prov) => {
+              return prov
+                .put(
+                  "test",
+                  [1, 2, 3, 4, 5].map((i) => {
+                    return { id: "a" + i, a: "index_value_a_" + i };
+                  })
+                )
+                .then(() => {
+                  return prov.getAll("test", undefined).then((rets) => {
+                    assert(!!rets);
+                    assert.equal(rets.length, 5);
+                    return prov
+                      .removeRange(
+                        "test",
+                        "index",
+                        "index_value_a_1",
+                        "index_value_a_5"
+                      )
+                      .then(() => {
+                        return prov
+                          .getAll("test", undefined)
+                          .then((retVals) => {
+                            const rets = retVals as TestObj[];
+                            assert(!!rets);
+                            assert.equal(rets.length, 0);
+                          });
+                      });
+                  });
+                })
+                .then(() => prov.close())
+                .catch((e) => prov.close().then(() => Promise.reject(e)));
+            })
+            .then(
+              () => done(),
+              (err) => done(err)
+            );
+        });
+
+        it("Remove range (all removed with index not unique)", (done) => {
+          // Not working with index, need to be fix in the future.
+          if (provName === "indexeddbfakekeys") {
+            done();
+            return;
+          }
+
+          openProvider(
+            provName,
+            {
+              version: 1,
+              stores: [
+                {
+                  name: "test",
+                  primaryKeyPath: "id",
+                  indexes: [
+                    {
+                      name: "index",
+                      keyPath: "a",
+                      unique: false,
+                    },
+                  ],
+                },
+              ],
+            },
+            true
+          )
+            .then((prov) => {
+              return prov
+                .put(
+                  "test",
+                  [1, 1, 2, 3, 4, 5].map((i, index) => {
+                    return {
+                      id: index,
+                      a: "index_value_a_" + i,
+                    };
+                  })
+                )
+                .then(() => {
+                  return prov.getAll("test", undefined).then((rets) => {
+                    assert(!!rets);
+                    assert.equal(rets.length, 6);
+                    return prov
+                      .removeRange(
+                        "test",
+                        "index",
+                        "index_value_a_1",
+                        "index_value_a_5"
+                      )
+                      .then(() => {
+                        return prov
+                          .getAll("test", undefined)
+                          .then((retVals) => {
+                            const rets = retVals as TestObj[];
+                            assert(!!rets);
+                            assert.equal(rets.length, 0);
+                          });
+                      });
+                  });
+                })
+                .then(() => prov.close())
+                .catch((e) => prov.close().then(() => Promise.reject(e)));
+            })
+            .then(
+              () => done(),
+              (err) => done(err)
+            );
+        });
+
         it("Invalid Key Type", (done) => {
           openProvider(
             provName,