howler-api 2.13.0.dev329__py3-none-any.whl

howler/remote/datatypes/hash.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
+
+from howler.common.exceptions import HowlerValueError
+from howler.remote.datatypes import get_client, retry_call
+
+if TYPE_CHECKING:
+    from redis import Redis
+
+
+_conditional_remove_script = """
+local hash_name = KEYS[1]
+local key_in_hash = ARGV[1]
+local expected_value = ARGV[2]
+local result = redis.call('hget', hash_name, key_in_hash)
+if result == expected_value then
+    redis.call('hdel', hash_name, key_in_hash)
+    return 1
+end
+return 0
+"""
+
+
+h_pop_script = """
+local result = redis.call('hget', ARGV[1], ARGV[2])
+if result then redis.call('hdel', ARGV[1], ARGV[2]) end
+return result
+"""
+
+
+_limited_add = """
+local set_name = KEYS[1]
+local key = ARGV[1]
+local value = ARGV[2]
+local limit = tonumber(ARGV[3])
+
+if redis.call('hlen', set_name) < limit then
+    return redis.call('hsetnx', set_name, key, value)
+end
+return nil
+"""
+
+T = TypeVar("T")
+
+
+class HashIterator(Generic[T]):
+    def __init__(self, hash_object: Hash[T]):
+        self.hash_object = hash_object
+        self.cursor = 0
+        self.buffer: list[tuple[str, T]] = []
+        self._load_next()
+
+    def __next__(self) -> tuple[str, T]:
+        while True:
+            if self.buffer:
+                return self.buffer.pop(0)
+            if self.cursor == 0:
+                raise StopIteration()
+            self._load_next()
+
+    def _load_next(self):
+        self.cursor, data = retry_call(self.hash_object.c.hscan, self.hash_object.name, self.cursor)
+        for key, value in data.items():
+            self.buffer.append((key.decode("utf-8"), json.loads(value)))
+
+
+class Hash(Generic[T]):
+    def __init__(
+        self,
+        name: str,
+        host: Optional[Union[str, Redis]] = None,
+        port: Optional[int] = None,
+    ):
+        self.c = get_client(host, port, False)
+        self.name = name
+        self._pop = self.c.register_script(h_pop_script)
+        self._limited_add = self.c.register_script(_limited_add)
+        self._conditional_remove = self.c.register_script(_conditional_remove_script)
+
+    def __iter__(self):
+        return HashIterator(self)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.delete()
+
+    def add(self, key: str, value: T) -> int:
+        """Add the (key, value) pair to the hash for new keys.
+
+        If a key already exists this operation doesn't add it.
+
+        Returns:
+            True if key has been added to the table, False otherwise.
+        """
+        if isinstance(key, bytes):
+            raise HowlerValueError("Cannot use bytes for hashmap keys")
+        return retry_call(self.c.hsetnx, self.name, key, json.dumps(value))
+
+    def increment(self, key, increment: int = 1):
+        return int(retry_call(self.c.hincrby, self.name, key, increment))
+
+    def limited_add(self, key, value, size_limit):
+        """Add a single value to the set, but only if that wouldn't make the set grow past a given size.
+
+        If the hash has hit the size limit returns None
+        Otherwise, returns the result of hsetnx (same as `add`)
+        """
+        return retry_call(
+            self._limited_add,
+            keys=[self.name],
+            args=[key, json.dumps(value), size_limit],
+        )
+
+    def exists(self, key: str) -> bool:
+        return retry_call(self.c.hexists, self.name, key)
+
+    def get(self, key: str) -> Optional[T]:
+        item = retry_call(self.c.hget, self.name, key)
+        if not item:
+            return item
+        return json.loads(item)
+
+    def keys(self) -> list[str]:
+        return [k.decode("utf-8") for k in retry_call(self.c.hkeys, self.name)]
+
+    def length(self):
+        return retry_call(self.c.hlen, self.name)
+
+    def items(self) -> dict[str, T]:
+        items = retry_call(self.c.hgetall, self.name)
+        if not isinstance(items, dict):
+            return {}
+        return {k.decode("utf-8"): json.loads(v) for k, v in items.items()}
+
+    def conditional_remove(self, key: str, value) -> bool:
+        return bool(
+            retry_call(
+                self._conditional_remove,
+                keys=[self.name],
+                args=[key, json.dumps(value)],
+            )
+        )
+
+    def pop(self, key: str):
+        item = retry_call(self._pop, args=[self.name, key])
+        if not item:
+            return item
+        return json.loads(item)
+
+    def set(self, key: str, value: T):
+        if isinstance(key, bytes):
+            raise HowlerValueError("Cannot use bytes for hashmap keys")
+        return retry_call(self.c.hset, self.name, key, json.dumps(value))
+
+    def multi_set(self, data: dict[str, T]):
+        if any(isinstance(key, bytes) for key in data.keys()):
+            raise HowlerValueError("Cannot use bytes for hashmap keys")
+        encoded = {key: json.dumps(value) for key, value in data.items()}
+        return retry_call(self.c.hset, self.name, mapping=encoded)
+
+    def delete(self):
+        retry_call(self.c.delete, self.name)
+
+
+class ExpiringHash(Hash):
+    def __init__(self, name, ttl=86400, host=None, port=None):
+        super(ExpiringHash, self).__init__(name, host, port)
+        self.ttl = ttl
+        self.last_expire_time: float = 0
+
+    def _conditional_expire(self):
+        if self.ttl:
+            ctime = time.time()
+            if ctime > self.last_expire_time + (self.ttl / 2):
+                retry_call(self.c.expire, self.name, self.ttl)
+                self.last_expire_time = ctime
+
+    def add(self, key, value):
+        rval = super(ExpiringHash, self).add(key, value)
+        self._conditional_expire()
+        return rval
+
+    def set(self, key, value):
+        rval = super(ExpiringHash, self).set(key, value)
+        self._conditional_expire()
+        return rval
+
+    def multi_set(self, data):
+        rval = super(ExpiringHash, self).multi_set(data)
+        self._conditional_expire()
+        return rval
+
+    def increment(self, key, increment=1):
+        rval = super(ExpiringHash, self).increment(key, increment)
+        self._conditional_expire()
+        return rval
+
+    def limited_add(self, key, value, size_limit):
+        rval = super(ExpiringHash, self).limited_add(key, value, size_limit)
+        self._conditional_expire()
+        return rval

howler/remote/datatypes/lock.py

@@ -0,0 +1,42 @@
+from howler.remote.datatypes import get_client, retry_call
+from howler.utils.uid import get_random_id
+
+lock_acquire_script = """
+local lock_holder = ARGV[1]
+local uuid = ARGV[2]
+local timeout = ARGV[3]
+if redis.call('setnx', lock_holder, uuid) == 1 then
+    redis.call('expire', lock_holder, timeout)
+    return true
+end
+return false
+"""
+
+lock_release_script = """
+local lock_holder = ARGV[1]
+local lock_release = ARGV[2]
+local uuid = ARGV[3]
+if redis.call('get', lock_holder) == uuid then
+    redis.call('del', lock_holder)
+    redis.call('rpush', lock_release, uuid)
+    redis.call('expire', lock_release, 1)
+end
+"""
+
+
+class Lock(object):
+    def __init__(self, name, timeout, host=None, port=None):
+        self.uuid = get_random_id()
+        self.c = get_client(host, port, False)
+        self.lock_release = "-".join(("lock", str(timeout), name, "released"))
+        self.lock_holder = "-".join(("lock", str(timeout), name, "holder"))
+        self.timeout = timeout
+        self._acquire = self.c.register_script(lock_acquire_script)
+        self._release = self.c.register_script(lock_release_script)
+
+    def __enter__(self):
+        while not retry_call(self._acquire, args=[self.lock_holder, self.uuid, self.timeout]):
+            retry_call(self.c.blpop, self.lock_release, 1)
+
+    def __exit__(self, unused1, unused2, unused3):
+        retry_call(self._release, args=[self.lock_holder, self.lock_release, self.uuid])

howler/remote/datatypes/queues/comms.py

@@ -0,0 +1,59 @@
+import json
+
+import redis
+
+from howler.remote.datatypes import decode, get_client, log, retry_call
+
+
+class CommsQueue(object):
+    def __init__(self, names, host=None, port=None, private=False):
+        self.c = get_client(host, port, private)
+        self.p = retry_call(self.c.pubsub)
+        if not isinstance(names, list):
+            names = [names]
+        self.names = names
+        self._connected = False
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        retry_call(self.p.unsubscribe)
+
+    def _connect(self):
+        if not self._connected:
+            retry_call(self.p.subscribe, self.names)
+            self._connected = True
+
+    def close(self):
+        retry_call(self.p.close)
+
+    def listen(self, blocking=True):
+        retried = False
+        while True:
+            self._connect()
+            try:
+                if blocking:
+                    i = self.p.listen()
+                    v = next(i)
+                else:
+                    v = self.p.get_message()
+                    if v is None:
+                        yield None
+                        continue
+
+                if isinstance(v, dict) and v.get("type", None) == "message":
+                    data = decode(v.get("data", "null"))
+                    yield data
+            except redis.ConnectionError:
+                log.warning("No connection to Redis, reconnecting...")
+                self._connected = False
+                retried = True
+            finally:
+                if self._connected and retried:
+                    log.info("Reconnected to Redis!")
+                    retried = False
+
+    def publish(self, message):
+        for name in self.names:
+            retry_call(self.c.publish, name, json.dumps(message))

howler/remote/datatypes/queues/multi.py

@@ -0,0 +1,32 @@
+import json
+
+from howler.remote.datatypes import get_client, retry_call
+
+
+class MultiQueue(object):
+    def __init__(self, host=None, port=None, private=False):
+        self.c = get_client(host, port, private)
+
+    def delete(self, name):
+        retry_call(self.c.delete, name)
+
+    def length(self, name):
+        return retry_call(self.c.llen, name)
+
+    def pop(self, name, blocking=True, timeout=0):
+        if blocking:
+            response = retry_call(self.c.blpop, name, timeout)
+        else:
+            response = retry_call(self.c.lpop, name)
+
+        if not response:
+            return response
+
+        if blocking:
+            return json.loads(response[1])
+        else:
+            return json.loads(response)
+
+    def push(self, name, *messages):
+        for message in messages:
+            retry_call(self.c.rpush, name, json.dumps(message))

howler/remote/datatypes/queues/named.py

@@ -0,0 +1,93 @@
+import json
+import time
+from typing import Generic, Optional, TypeVar
+
+from howler.common.exceptions import HowlerTypeError
+from howler.remote.datatypes import get_client, retry_call
+
+T = TypeVar("T")
+
+
+class NamedQueue(Generic[T]):
+    def __init__(self, name: str, host=None, port=None, private: bool = False, ttl: int = 0):
+        self.c = get_client(host, port, private)
+        self.name: str = name
+        self.ttl: int = ttl
+        self.last_expire_time: float = 0
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.delete()
+
+    def _conditional_expire(self):
+        if self.ttl:
+            ctime = time.time()
+            if ctime > self.last_expire_time + (self.ttl / 2):
+                retry_call(self.c.expire, self.name, self.ttl)
+                self.last_expire_time = ctime
+
+    def delete(self):
+        retry_call(self.c.delete, self.name)
+
+    def __len__(self):
+        return self.length()
+
+    def length(self):
+        return retry_call(self.c.llen, self.name)
+
+    def peek_next(self) -> Optional[T]:
+        response = retry_call(self.c.lrange, self.name, 0, 0)
+
+        if response:
+            return json.loads(response[0])
+        return None
+
+    def pop_batch(self, size) -> list[T]:
+        response = retry_call(self.c.lpop, self.name, size)
+
+        if not response:
+            return []
+        return [json.loads(r) for r in response]
+
+    def pop(self, blocking: bool = True, timeout: int = 0) -> Optional[T]:
+        if blocking:
+            response = retry_call(self.c.blpop, self.name, timeout)
+        else:
+            response = retry_call(self.c.lpop, self.name)
+
+        if not response:
+            return response
+
+        if blocking:
+            return json.loads(response[1])
+        else:
+            return json.loads(response)
+
+    def push(self, *messages: T):
+        for message in messages:
+            retry_call(self.c.rpush, self.name, json.dumps(message))
+        self._conditional_expire()
+
+    def unpop(self, *messages: T):
+        """Put all messages passed back at the head of the FIFO queue."""
+        for message in messages:
+            retry_call(self.c.lpush, self.name, json.dumps(message))
+        self._conditional_expire()
+
+
+def select(*queues, **kw):
+    timeout = kw.get("timeout", 0)
+    if len(queues) < 1:
+        raise HowlerTypeError("At least one queue must be specified")
+    if any([not isinstance(q, NamedQueue) for q in queues]):
+        raise HowlerTypeError("Only NamedQueues supported")
+
+    c = queues[0].c
+    response = retry_call(c.blpop, [q.name for q in queues], timeout)
+
+    if not response:
+        return response
+
+    return response[0].decode("utf-8"), json.loads(response[1])

howler/remote/datatypes/queues/priority.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import Generic, Optional, TypeVar, Union
+
+from howler.common.exceptions import HowlerTypeError
+from howler.remote.datatypes import decode, get_client, retry_call
+
+SORTING_KEY_LEN = 21
+
+# Work around for inconsistency between ZRANGEBYSCORE and ZREMRANGEBYSCORE
+#   (No limit option available or we would just be using that directly)
+#
+# args:
+#   minimum score to pop
+#   maximum score to pop
+#   number of elements to skip before popping any
+#   max element count to pop
+pq_dequeue_range_script = """
+local unpack = table.unpack or unpack
+local min_score = tonumber(ARGV[1]);
+if min_score == nil then min_score = -math.huge end
+local max_score = tonumber(ARGV[2]);
+if max_score == nil then max_score = math.huge end
+local rem_offset = tonumber(ARGV[3]);
+local rem_limit = tonumber(ARGV[4]);
+
+local entries = redis.call("zrangebyscore", KEYS[1], -max_score, -min_score, "limit", rem_offset, rem_limit);
+if #entries > 0 then redis.call("zrem", KEYS[1], unpack(entries)) end
+return entries
+"""
+
+T = TypeVar("T")
+
+
+class PriorityQueue(Generic[T]):
+    def __init__(self, name, host=None, port=None, private=False):
+        self.c = get_client(host, port, private)
+        self._deque_range = self.c.register_script(pq_dequeue_range_script)
+        self.name = name
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.delete()
+
+    def count(self, lowest, highest):
+        return retry_call(self.c.zcount, self.name, -highest, -lowest)
+
+    def delete(self):
+        retry_call(self.c.delete, self.name)
+
+    def length(self):
+        return retry_call(self.c.zcard, self.name)
+
+    def pop(self, num=None):
+        if num is not None and num <= 0:
+            return []
+
+        if num:
+            return [decode(s[0][SORTING_KEY_LEN:]) for s in retry_call(self.c.zpopmin, self.name, num)]
+        else:
+            ret_val = retry_call(self.c.zpopmin, self.name, 1)
+            if ret_val:
+                return decode(ret_val[0][0][SORTING_KEY_LEN:])
+            return None
+
+    def blocking_pop(self, timeout=0, low_priority=False):
+        """When only one item is requested, blocking is is possible."""
+        if low_priority:
+            result = retry_call(self.c.bzpopmax, self.name, timeout)
+        else:
+            result = retry_call(self.c.bzpopmin, self.name, timeout)
+        if result:
+            return decode(result[1][SORTING_KEY_LEN:])
+        return None
+
+    def dequeue_range(self, lower_limit="", upper_limit="", skip=0, num=1) -> list[T]:
+        """Dequeue a number of elements, within a specified range of scores.
+
+        Limits given are inclusive, can be made exclusive, see redis docs on how to format limits for that.
+
+        NOTE: lower/upper limit is negated+swapped in the lua script, no need to do it here
+
+        :param lower_limit: The score of all dequeued elements must be higher or equal to this.
+        :param upper_limit: The score of all dequeued elements must be lower or equal to this.
+        :param skip: In the range of available items to dequeue skip over this many.
+        :param num: Maximum number of elements to dequeue.
+        :return: list
+        """
+        results = retry_call(
+            self._deque_range,
+            keys=[self.name],
+            args=[lower_limit, upper_limit, skip, num],
+        )
+        return [decode(res[SORTING_KEY_LEN:]) for res in results]
+
+    def push(self, priority: int, data: T, vip=None):
+        vip = 0 if vip else 9
+        value = f"{vip}{f'{int(time.time()*1000000):020}'}{json.dumps(data)}"
+        retry_call(self.c.zadd, self.name, {value: -priority})
+        return value
+
+    def rank(self, raw_value):
+        return retry_call(self.c.zrank, self.name, raw_value)
+
+    def remove(self, raw_value):
+        return retry_call(self.c.zrem, self.name, raw_value)
+
+    def unpush(self, num=None) -> Union[list[T], Optional[T]]:
+        if num is not None and num <= 0:
+            return []
+
+        if num:
+            return [decode(s[0][SORTING_KEY_LEN:]) for s in retry_call(self.c.zpopmax, self.name, num)]
+        else:
+            ret_val = retry_call(self.c.zpopmax, self.name, 1)
+            if ret_val:
+                return decode(ret_val[0][0][SORTING_KEY_LEN:])
+            return None
+
+
+class UniquePriorityQueue(PriorityQueue):
+    """A priority queue where duplicate entries will be merged."""
+
+    def __init__(self, name, host=None, port=None, private=False):
+        super().__init__(name, host, port, private)
+
+    def remove(self, data: str):
+        """Remove a value from the priority  queue."""
+        retry_call(self.c.zrem, self.name, json.dumps(data))
+
+    def push(self, priority: int, data, vip=None) -> int:
+        """Add or update elements in the priority queue.
+
+        Existing elements will have their priority updated.
+
+        Returns:
+            Number of _NEW_ elements in the queue after the operation.
+        """
+        return retry_call(self.c.zadd, self.name, {json.dumps(data): -priority})
+
+    def pop(self, num=None):
+        if num is not None and num <= 0:
+            return []
+
+        if num:
+            return [decode(s[0]) for s in retry_call(self.c.zpopmin, self.name, num)]
+        else:
+            ret_val = retry_call(self.c.zpopmin, self.name, 1)
+            if ret_val:
+                return decode(ret_val[0][0])
+            return None
+
+    def unpush(self, num=None):
+        if num is not None and num <= 0:
+            return []
+
+        if num:
+            return [decode(s[0]) for s in retry_call(self.c.zpopmax, self.name, num)]
+        else:
+            ret_val = retry_call(self.c.zpopmax, self.name, 1)
+            if ret_val:
+                return decode(ret_val[0][0])
+            return None
+
+    def dequeue_range(self, lower_limit="", upper_limit="", skip=0, num=1):
+        """Dequeue a number of elements, within a specified range of scores.
+
+        Limits given are inclusive, can be made exclusive, see redis docs on how to format limits for that.
+
+        NOTE: lower/upper limit is negated+swapped in the lua script, no need to do it here
+
+        :param lower_limit: The score of all dequeued elements must be higher or equal to this.
+        :param upper_limit: The score of all dequeued elements must be lower or equal to this.
+        :param skip: In the range of available items to dequeue skip over this many.
+        :param num: Maximum number of elements to dequeue.
+        :return: list
+        """
+        results = retry_call(
+            self._deque_range,
+            keys=[self.name],
+            args=[lower_limit, upper_limit, skip, num],
+        )
+        return [decode(res) for res in results]
+
+
+def select(*queues, **kw):
+    timeout = kw.get("timeout", 0)
+    if len(queues) < 1:
+        raise HowlerTypeError("At least one queue must be specified")
+    if any([not isinstance(q, PriorityQueue) for q in queues]):
+        raise HowlerTypeError("Only NamedQueues supported")
+
+    c = queues[0].c
+    response = retry_call(c.bzpopmin, [q.name for q in queues], timeout)
+
+    if not response:
+        return response
+
+    return response[0].decode("utf-8"), json.loads(response[1][SORTING_KEY_LEN:])
+
+
+def length(*queues: PriorityQueue) -> list[int]:
+    """Utility function for batch reading queue lengths."""
+    if not queues:
+        return []
+    pipeline = queues[0].c.pipeline(transaction=False)
+
+    for que in queues:
+        pipeline.zcard(que.name)
+
+    return retry_call(pipeline.execute)