Add nodup option to ScanKeys (SYN-7497)

synapse/lib/lmdbslab.py

@@ -1253,18 +1253,18 @@ def stat(self, db=None):
         finally:
             self._relXactForReading()
 
-    def scanKeys(self, db=None):
+    def scanKeys(self, db=None, nodup=False):
 
-        with ScanKeys(self, db) as scan:
+        with ScanKeys(self, db, nodup=nodup) as scan:
 
             if not scan.first():
                 return
 
             yield from scan.iternext()
 
-    def scanKeysByPref(self, byts, db=None):
+    def scanKeysByPref(self, byts, db=None, nodup=False):
 
-        with ScanKeys(self, db) as scan:
+        with ScanKeys(self, db, nodup=nodup) as scan:
 
             if not scan.set_range(byts):
                 return
@@ -1283,7 +1283,7 @@ async def countByPref(self, byts, db=None, maxsize=None):
         '''
         count = 0
         size = len(byts)
-        with ScanKeys(self, db) as scan:
+        with ScanKeys(self, db, nodup=True) as scan:
 
             if not scan.set_range(byts):
                 return 0
@@ -1293,7 +1293,7 @@ async def countByPref(self, byts, db=None, maxsize=None):
                 if lkey[:size] != byts:
                     return count
 
-                count += 1
+                count += scan.curs.count()
                 if maxsize is not None and maxsize == count:
                     return count
 
@@ -1756,14 +1756,21 @@ class ScanKeys(Scan):
     An iterator over the keys of the database.  If the database is dupsort, a key with multiple values with be yielded
     once for each value.
     '''
+    def __init__(self, slab, db, nodup=False):
+        Scan.__init__(self, slab, db)
+        self.nodup = nodup
+
     def iterfunc(self):
         if self.dupsort:
-            return Scan.iterfunc(self)
+            if self.nodup:
+                return self.curs.iternext_nodup(keys=True, values=False)
+            else:
+                return Scan.iterfunc(self)
 
         return self.curs.iternext(keys=True, values=False)
 
     def resume(self):
-        if self.dupsort:
+        if self.dupsort and not self.nodup:
             return Scan.resume(self)
 
         return self.curs.set_range(self.atitem)
@@ -1772,13 +1779,13 @@ def isatitem(self):
         '''
         Returns if the cursor is at the value in atitem
         '''
-        if self.dupsort:
+        if self.dupsort and not self.nodup:
             return Scan.isatitem(self)
 
         return self.atitem == self.curs.key()
 
     def iternext(self):
-        if self.dupsort:
+        if self.dupsort and not self.nodup:
             yield from (item[0] for item in Scan.iternext(self))
             return
 

synapse/tests/test_lib_lmdbslab.py

@@ -57,9 +57,11 @@ async def test_lmdbslab_scankeys(self):
 
                 testgenr = slab.scanKeys(db=testdb)
                 dupsgenr = slab.scanKeys(db=dupsdb)
+                nodupsgenr = slab.scanKeys(db=dupsdb, nodup=True)
 
                 testlist = [next(testgenr)]
                 dupslist = [next(dupsgenr)]
+                nodupslist = [next(nodupsgenr)]
 
                 slab.put(b'derp', b'derp', db=editdb)
 
@@ -68,16 +70,25 @@ async def test_lmdbslab_scankeys(self):
 
                 testlist.extend(testgenr)
                 dupslist.extend(dupsgenr)
+                nodupslist.extend(nodupsgenr)
 
                 self.eq(testlist, (b'hehe', b'hoho'))
                 self.eq(dupslist, (b'hehe', b'hehe', b'hoho'))
+                self.eq(nodupslist, (b'hehe', b'hoho'))
+
+                self.eq([b'hehe', b'hehe', b'hoho'], list(slab.scanKeysByPref(b'h', db=dupsdb)))
+                self.eq([b'hehe', b'hoho'], list(slab.scanKeysByPref(b'h', db=dupsdb, nodup=True)))
+
+                self.eq(3, await slab.countByPref(b'h', db=dupsdb))
 
                 # now lets delete the key we're on
                 testgenr = slab.scanKeys(db=testdb)
                 dupsgenr = slab.scanKeys(db=testdb)
+                nodupsgenr = slab.scanKeys(db=testdb, nodup=True)
 
                 testlist = [next(testgenr)]
                 dupslist = [next(dupsgenr)]
+                nodupslist = [next(nodupsgenr)]
 
                 slab.delete(b'hehe', db=testdb)
                 for lkey, lval in slab.scanByDups(b'hehe', db=dupsdb):
@@ -87,9 +98,11 @@ async def test_lmdbslab_scankeys(self):
 
                 testlist.extend(testgenr)
                 dupslist.extend(dupsgenr)
+                nodupslist.extend(nodupsgenr)
 
                 self.eq(testlist, (b'hehe', b'hoho'))
                 self.eq(dupslist, (b'hehe', b'hoho'))
+                self.eq(nodupslist, (b'hehe', b'hoho'))
 
                 # by pref
                 self.eq([b'hoho'], list(slab.scanKeysByPref(b'h', db=dupsdb)))