203 lines
6.5 KiB
Python
203 lines
6.5 KiB
Python
"""Command mixin for emulating `redis-py`'s BF functionality."""
|
|
import io
|
|
|
|
import pybloom_live
|
|
|
|
from fakeredis import _msgs as msgs
|
|
from fakeredis._command_args_parsing import extract_args
|
|
from fakeredis._commands import command, Key, CommandItem, Float, Int
|
|
from fakeredis._helpers import SimpleError, OK, casematch
|
|
|
|
|
|
class ScalableBloomFilter(pybloom_live.ScalableBloomFilter):
|
|
NO_GROWTH = 0
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.filters.append(
|
|
pybloom_live.BloomFilter(
|
|
capacity=self.initial_capacity,
|
|
error_rate=self.error_rate * self.ratio))
|
|
|
|
def add(self, key):
|
|
if key in self:
|
|
return True
|
|
if self.scale == self.NO_GROWTH and self.filters and self.filters[-1].count >= self.filters[-1].capacity:
|
|
raise SimpleError(msgs.FILTER_FULL_MSG)
|
|
return super(ScalableBloomFilter, self).add(key)
|
|
|
|
|
|
class BFCommandsMixin:
|
|
|
|
@staticmethod
|
|
def _bf_add(key: CommandItem, item: bytes) -> int:
|
|
res = key.value.add(item)
|
|
key.updated()
|
|
return 0 if res else 1
|
|
|
|
@staticmethod
|
|
def _bf_exist(key: CommandItem, item: bytes) -> int:
|
|
return 1 if (item in key.value) else 0
|
|
|
|
@command(
|
|
name="BF.ADD",
|
|
fixed=(Key(ScalableBloomFilter), bytes),
|
|
repeat=(),
|
|
)
|
|
def bf_add(self, key, value: bytes):
|
|
return BFCommandsMixin._bf_add(key, value)
|
|
|
|
@command(
|
|
name="BF.MADD",
|
|
fixed=(Key(ScalableBloomFilter), bytes),
|
|
repeat=(bytes,),
|
|
)
|
|
def bf_madd(self, key, *values):
|
|
res = list()
|
|
for value in values:
|
|
res.append(BFCommandsMixin._bf_add(key, value))
|
|
return res
|
|
|
|
@command(
|
|
name="BF.CARD",
|
|
fixed=(Key(ScalableBloomFilter),),
|
|
repeat=(),
|
|
)
|
|
def bf_card(self, key):
|
|
return len(key.value)
|
|
|
|
@command(
|
|
name="BF.EXISTS",
|
|
fixed=(Key(ScalableBloomFilter), bytes),
|
|
repeat=(),
|
|
)
|
|
def bf_exist(self, key, value: bytes):
|
|
return BFCommandsMixin._bf_exist(key, value)
|
|
|
|
@command(
|
|
name="BF.MEXISTS",
|
|
fixed=(Key(ScalableBloomFilter), bytes),
|
|
repeat=(bytes,),
|
|
)
|
|
def bf_mexists(self, key, *values: bytes):
|
|
res = list()
|
|
for value in values:
|
|
res.append(BFCommandsMixin._bf_exist(key, value))
|
|
return res
|
|
|
|
@command(
|
|
name="BF.RESERVE",
|
|
fixed=(Key(), Float, Int,),
|
|
repeat=(bytes,),
|
|
flags=msgs.FLAG_LEAVE_EMPTY_VAL,
|
|
)
|
|
def bf_reserve(self, key: CommandItem, error_rate, capacity, *args: bytes):
|
|
if key.value is not None:
|
|
raise SimpleError(msgs.ITEM_EXISTS_MSG)
|
|
(expansion, non_scaling), _ = extract_args(args, ("+expansion", "nonscaling"))
|
|
if expansion is not None and non_scaling:
|
|
raise SimpleError(msgs.NONSCALING_FILTERS_CANNOT_EXPAND_MSG)
|
|
if expansion is None:
|
|
expansion = 2
|
|
scale = ScalableBloomFilter.NO_GROWTH if non_scaling else expansion
|
|
key.update(ScalableBloomFilter(capacity, error_rate, scale))
|
|
return OK
|
|
|
|
@command(
|
|
name="BF.INSERT",
|
|
fixed=(Key(),),
|
|
repeat=(bytes,),
|
|
)
|
|
def bf_insert(self, key: CommandItem, *args: bytes):
|
|
(capacity, error_rate, expansion, non_scaling, no_create), left_args = extract_args(
|
|
args, ("+capacity", ".error", "+expansion", "nonscaling", "nocreate"),
|
|
error_on_unexpected=False, left_from_first_unexpected=True)
|
|
# if no_create and (capacity is not None or error_rate is not None):
|
|
# raise SimpleError("...")
|
|
if len(left_args) < 2 or not casematch(left_args[0], b'items'):
|
|
raise SimpleError("...")
|
|
items = left_args[1:]
|
|
|
|
error_rate = error_rate or 0.001
|
|
capacity = capacity or 100
|
|
if key.value is None and no_create:
|
|
raise SimpleError(msgs.NOT_FOUND_MSG)
|
|
if expansion is not None and non_scaling:
|
|
raise SimpleError(msgs.NONSCALING_FILTERS_CANNOT_EXPAND_MSG)
|
|
if expansion is None:
|
|
expansion = 2
|
|
scale = ScalableBloomFilter.NO_GROWTH if non_scaling else expansion
|
|
if key.value is None:
|
|
key.value = ScalableBloomFilter(capacity, error_rate, scale)
|
|
res = list()
|
|
for item in items:
|
|
res.append(self._bf_add(key, item))
|
|
key.updated()
|
|
return res
|
|
|
|
@command(
|
|
name="BF.INFO",
|
|
fixed=(Key(),),
|
|
repeat=(bytes,),
|
|
)
|
|
def bf_info(self, key: CommandItem, *args: bytes):
|
|
if key.value is None or type(key.value) is not ScalableBloomFilter:
|
|
raise SimpleError('...')
|
|
if len(args) > 1:
|
|
raise SimpleError(msgs.SYNTAX_ERROR_MSG)
|
|
if len(args) == 0:
|
|
return [
|
|
b'Capacity', key.value.capacity,
|
|
b'Size', key.value.capacity,
|
|
b'Number of filters', len(key.value.filters),
|
|
b'Number of items inserted', key.value.count,
|
|
b'Expansion rate', key.value.scale if key.value.scale > 0 else None,
|
|
]
|
|
if casematch(args[0], b'CAPACITY'):
|
|
return key.value.capacity
|
|
elif casematch(args[0], b'SIZE'):
|
|
return key.value.capacity
|
|
elif casematch(args[0], b'FILTERS'):
|
|
return len(key.value.filters)
|
|
elif casematch(args[0], b'ITEMS'):
|
|
return key.value.count
|
|
elif casematch(args[0], b'EXPANSION'):
|
|
return key.value.scale if key.value.scale > 0 else None
|
|
else:
|
|
raise SimpleError(msgs.SYNTAX_ERROR_MSG)
|
|
|
|
@command(
|
|
name="BF.SCANDUMP",
|
|
fixed=(Key(), Int,),
|
|
repeat=(),
|
|
flags=msgs.FLAG_LEAVE_EMPTY_VAL,
|
|
)
|
|
def bf_scandump(self, key: CommandItem, iterator: int):
|
|
if key.value is None:
|
|
raise SimpleError(msgs.NOT_FOUND_MSG)
|
|
f = io.BytesIO()
|
|
|
|
if iterator == 0:
|
|
key.value.tofile(f)
|
|
f.seek(0)
|
|
s = f.read()
|
|
f.close()
|
|
return [1, s]
|
|
else:
|
|
return [0, None]
|
|
|
|
@command(
|
|
name="BF.LOADCHUNK",
|
|
fixed=(Key(), Int, bytes),
|
|
repeat=(),
|
|
flags=msgs.FLAG_LEAVE_EMPTY_VAL,
|
|
)
|
|
def bf_loadchunk(self, key: CommandItem, iterator: int, data: bytes):
|
|
if key.value is not None and type(key.value) is not ScalableBloomFilter:
|
|
raise SimpleError(msgs.NOT_FOUND_MSG)
|
|
f = io.BytesIO(data)
|
|
key.value = ScalableBloomFilter.fromfile(f)
|
|
f.close()
|
|
key.updated()
|
|
return OK
|