clue-api 1.0.0.dev7__py3-none-any.whl

clue/plugin/utils.py

@@ -0,0 +1,78 @@
+import os
+import time
+from datetime import datetime, timezone
+
+from flask import request
+from pydantic import BaseModel
+
+from clue.common.logging import get_logger
+
+# Default settings
+MAX_LIMIT = int(os.environ.get("MAX_LIMIT", 100))
+MAX_TIMEOUT = float(os.environ.get("MAX_TIMEOUT", 3))
+
+logger = get_logger(__file__)
+
+
+class Params(BaseModel):
+    "A model build to parse arguments provided in the request object into a model"
+
+    deadline: float
+    "The epoch the central server wants to have a response by"
+
+    max_timeout: float
+    "The raw timeout value provided by the user"
+
+    annotate: bool
+    "Should the plugin return annotations about the given selector(s)?"
+
+    raw: bool
+    "Should the plugin return the raw data applicable to the given selector(s)?"
+
+    limit: int
+    "What is the maximum number of query entries that should be returned?"
+
+    use_cache: bool
+    "Does the request want to bypass any cached results?"
+
+    @classmethod
+    def from_request(cls):
+        "Create a Params object from flask's request object"
+        max_timeout = request.args.get("max_timeout", MAX_TIMEOUT)
+        try:
+            max_timeout = float(max_timeout)
+        except (ValueError, TypeError):
+            max_timeout = MAX_TIMEOUT
+
+        deadline = request.args.get("deadline", time.time() + max_timeout, type=float)
+        current_time = datetime.now(timezone.utc)
+        if deadline > 0 and current_time.timestamp() > deadline:
+            logger.warning(
+                "Deadline %s was earlier than the current time, %s",
+                str(datetime.fromtimestamp(deadline)),
+                str(current_time),
+            )
+
+            raise RuntimeError("Deadline exceeded")
+
+        logger.debug(
+            "Deadline %s hits in %sms",
+            str(datetime.fromtimestamp(deadline)),
+            round((deadline - current_time.timestamp()) * 1000),
+        )
+
+        annotate = request.args.get("no_annotation", "false").lower() not in ("true", "1", "")
+        use_cache = request.args.get("no_cache", "false").lower() not in ("true", "1", "")
+        raw = request.args.get("include_raw", "false").lower() in ("true", "1", "")
+        limit = request.args.get("limit", 100, type=int)
+        if limit > int(MAX_LIMIT):
+            limit = int(MAX_LIMIT)
+
+        return cls(
+            deadline=deadline, max_timeout=max_timeout, annotate=annotate, raw=raw, limit=limit, use_cache=use_cache
+        )
+
+    def __str__(self):
+        # Make a string representation of the params that can be used for caching purposes
+        # Deadline and max_timout are explicitely ignored otherwise it would never hit the cache
+        return f"a={self.annotate},r={self.raw},l={self.limit}"

clue/remote/datatypes/__init__.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+
+import json
+import os
+import time
+from datetime import datetime
+
+import redis
+from packaging.version import parse
+
+from clue.common.logging import get_logger
+from clue.common.uid import get_random_id
+
+logger = get_logger(__file__)
+
+# Add a version warning if redis python client is < 2.10.0. Older versions
+# have a connection bug that can manifest with the dispatcher.
+if parse(redis.__version__) <= parse("2.10.0"):
+    import warnings
+
+    warnings.warn(
+        "%s works best with redis > 2.10.0. You're running"
+        " redis %s. You should upgrade." % (__name__, redis.__version__)
+    )
+
+
+pool: dict[tuple[str, int], redis.ConnectionPool] = {}
+
+
+def now_as_iso():
+    s = datetime.utcfromtimestamp(time.time()).isoformat()
+    return f"{s}Z"
+
+
+def reply_queue_name(prefix=None, suffix=None):
+    if prefix:
+        components = [prefix]
+    else:
+        components = []
+
+    components.append(get_random_id())
+
+    if suffix:
+        components.append(str(suffix))
+
+    return "-".join(components)
+
+
+def retry_call(func, *args, **kw):
+    max_attempts = 10
+    maximum = 2
+    exponent = -7
+
+    attempts = 0
+    while True:
+        try:
+            ret_val = func(*args, **kw)
+
+            if exponent != -7:
+                logger.info("Reconnected to Redis!")
+
+            return ret_val
+
+        except (redis.ConnectionError, ConnectionResetError):
+            attempts += 1
+
+            if attempts > max_attempts:
+                logger.exception("Redis connection failed.")
+                raise
+            else:
+                logger.exception("No connection to Redis, reconnecting...")
+                time.sleep(2**exponent)
+                exponent = exponent + 1 if exponent < maximum else exponent
+
+
+def _redis_ssl_kwargs(host: str) -> dict:
+    return dict(ssl_ca_certs=os.environ.get(f"{host.upper()}_ROOT_CA_PATH", "/etc/clue/ssl/clue_root-ca.crt"))
+
+
+def get_client(host, port, private):
+    # In case a structure is passed a client as host
+    if isinstance(host, (redis.Redis, redis.StrictRedis)):
+        return host
+
+    if not host or not port:
+        from clue.config import config
+
+        host = host or config.core.redis.host
+        port = int(port or config.core.redis.port)
+
+    ssl_kwargs = {}
+
+    # Automatically detect if encryption was enabled
+    tmp_ssl_kwargs = _redis_ssl_kwargs(host)
+    if os.path.exists(tmp_ssl_kwargs["ssl_ca_certs"]):
+        ssl_kwargs = tmp_ssl_kwargs
+        ssl_kwargs["ssl"] = True
+
+    if private:
+        return redis.StrictRedis(host=host, port=port, socket_keepalive=True, **ssl_kwargs)
+    else:
+        return redis.StrictRedis(
+            connection_pool=get_pool(host, port, ssl=ssl_kwargs.get("ssl", False)), socket_keepalive=True
+        )
+
+
+def get_pool(host, port, ssl=False):
+    key = (host, port)
+    connection_class = redis.connection.Connection
+    connection_kwargs = {}
+    if ssl:
+        connection_class = redis.connection.SSLConnection  # type: ignore[assignment]
+        connection_kwargs = _redis_ssl_kwargs(host)
+
+    connection_pool = pool.get(key, None)
+    if not connection_pool:
+        connection_pool = redis.BlockingConnectionPool(
+            host=host, port=port, max_connections=200, connection_class=connection_class, **connection_kwargs
+        )
+        pool[key] = connection_pool
+
+    return connection_pool
+
+
+def decode(data):
+    try:
+        return json.loads(data)
+    except ValueError:
+        logger.warning("Invalid data on queue: %s", str(data))
+        return None

clue/remote/datatypes/cache.py

@@ -0,0 +1,62 @@
+import json
+
+from redis.exceptions import ConnectionError
+
+from clue.common.uid import get_id_from_data
+from clue.remote.datatypes import get_client, retry_call
+
+DEFAULT_TTL = 60 * 60  # 1 Hour
+
+
+class RedisCache(object):
+    def __init__(self, prefix="brl_cache", host=None, port=None, ttl=DEFAULT_TTL):
+        self.c = get_client(host, port, False)
+        self.prefix = prefix
+        self.ttl = ttl
+
+    def __enter__(self):
+        return self
+
+    def _get_key(self, name):
+        return f"{self.prefix}-{name}"
+
+    def clear(self):
+        # Clear all items belonging to this cahce
+        for queue in retry_call(self.c.keys, "%s-*" % self.prefix):
+            retry_call(self.c.delete, queue)
+
+    def create_key(self, *args):
+        key_str = "-".join([str(x) for x in args])
+        return get_id_from_data(key_str)
+
+    def get(self, key, ttl=None):
+        # Get the key name
+        cache_name = self._get_key(key)
+
+        # Get the value from the cache
+        item = retry_call(self.c.get, cache_name)
+        if not item:
+            return item
+
+        # Reset the cache while we're still using it
+        retry_call(self.c.expire, cache_name, ttl or self.ttl)
+        return json.loads(item)
+
+    def ready(self):
+        try:
+            self.c.ping()
+        except ConnectionError:
+            return False
+
+        return True
+
+    def set(self, key, value, ttl=None):
+        # Get the key name
+        cache_name = self._get_key(key)
+
+        # Set the value and the expiry for the name
+        retry_call(self.c.set, cache_name, json.dumps(value))
+        retry_call(self.c.expire, cache_name, ttl or self.ttl)
+
+
+Cache = RedisCache

clue/remote/datatypes/events.py

@@ -0,0 +1,118 @@
+from __future__ import annotations
+
+import json
+import logging
+import threading
+from typing import TYPE_CHECKING, Any, Callable, Generic, Optional, TypeVar
+
+from clue.remote.datatypes import get_client, retry_call
+
+if TYPE_CHECKING:
+    from redis import Redis
+    from redis.client import PubSub
+
+
+logger = logging.getLogger(__name__)
+
+MessageType = TypeVar("MessageType")
+
+
+class EventSender(Generic[MessageType]):
+    def __init__(
+        self, prefix: str, host=None, port=None, private=None, serializer: Callable[[MessageType], str] = json.dumps
+    ):
+        self.client: Redis[Any] = get_client(host, port, private)
+        self.prefix = prefix.lower()
+        if not self.prefix.endswith("."):
+            self.prefix += "."
+        self.serializer = serializer
+
+    def send(self, name: str, data: MessageType):
+        path = self.prefix + name.lower().lstrip(".")
+        retry_call(self.client.publish, path, self.serializer(data))
+
+
+class PubSubWorkerThread(threading.Thread):
+    """Worker thread that continually reads messages from pubsub.
+
+    We reimplement the worker thread here rather than use the one in the redis
+    package because we want to use the subscribe messages for disconnect/reconnect
+    detection.
+    """
+
+    def __init__(self, watcher, exception_handler=None, skip_first_refresh=True):
+        super().__init__(daemon=True)
+        self.watcher = watcher
+        self.exception_handler = exception_handler
+        self._running = threading.Event()
+        self.skip_first_refresh = skip_first_refresh
+
+    def run(self):
+        if self._running.is_set():
+            return
+        self._running.set()
+
+        pubsub: PubSub = self.watcher.pubsub
+        ping = self.watcher.client.ping
+        sleep_time = 1
+        initialized = set()
+
+        while self._running.is_set():
+            try:
+                message = pubsub.get_message(ignore_subscribe_messages=False, timeout=sleep_time)
+                if message is not None and message["type"] == "psubscribe":
+                    channel = message.get("channel")
+                    if channel is None:
+                        continue
+
+                    if self.skip_first_refresh and channel not in initialized:
+                        initialized.add(channel)
+                        continue
+
+                    handler = pubsub.patterns.get(channel, None)
+                    if handler:
+                        handler(None)
+
+            except BaseException:
+                # Present the error
+                logger.exception("Exception in pubsub watcher:")
+
+                # Wait until we can reach the server
+                retry_call(ping)
+
+        pubsub.close()
+
+    def stop(self):
+        # trip the flag so the run loop exits. the run loop will
+        # close the pubsub connection, which disconnects the socket
+        # and returns the connection to the pool.
+        self._running.clear()
+
+
+class EventWatcher(Generic[MessageType]):
+    def __init__(self, host=None, port=None, deserializer: Callable[[str], MessageType] = json.loads):
+        self.client: Redis[Any] = get_client(host, port, False)
+        self.pubsub = retry_call(self.client.pubsub)
+        self.pubsub.ignore_subscribe_messages = False
+        self.thread: Optional[PubSubWorkerThread] = None
+        self.deserializer = deserializer
+        self.skip_first_refresh = True
+
+    def register(self, path: str, callback: Callable[[Optional[MessageType]], None]):
+        def _callback(message: Optional[dict[str, Any]]):
+            if message is None:
+                callback(None)
+            elif message["type"] == "pmessage":
+                data = self.deserializer(message.get("data", ""))
+                callback(data)
+
+        self.pubsub.psubscribe(**{path.lower(): _callback})
+
+    def start(self):
+        self.thread = PubSubWorkerThread(self, skip_first_refresh=self.skip_first_refresh)
+        self.thread.start()
+        return self.thread
+
+    def stop(self):
+        if self.thread is not None:
+            self.thread.stop()

clue/remote/datatypes/hash.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import json
+import time
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
+
+from clue.common.exceptions import ClueTypeError
+from clue.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[T] = []
+        self._load_next()
+
+    def __next__(self) -> 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.client.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: Union[str, Redis] = None, port: int = None):
+        self.client = get_client(host, port, False)
+        self.name = name
+        self._pop = self.client.register_script(h_pop_script)
+        self._limited_add = self.client.register_script(_limited_add)
+        self._conditional_remove = self.client.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 ClueTypeError("Cannot use bytes for hashmap keys")
+
+        return retry_call(self.client.hsetnx, self.name, key, json.dumps(value))
+
+    def increment(self, key, increment: int = 1):
+        return int(retry_call(self.client.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.client.hexists, self.name, key)
+
+    def get(self, key: str) -> Optional[T]:
+        item = retry_call(self.client.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.client.hkeys, self.name)]
+
+    def length(self):
+        return retry_call(self.client.hlen, self.name)
+
+    def items(self) -> dict[str, T]:
+        items = retry_call(self.client.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 ClueTypeError("Cannot use bytes for hashmap keys")
+
+        return retry_call(self.client.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 ValueError("Cannot use bytes for hashmap keys")
+        encoded = {key: json.dumps(value) for key, value in data.items()}
+        return retry_call(self.client.hset, self.name, mapping=encoded)
+
+    def delete(self):
+        retry_call(self.client.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 = 0
+
+    def _conditional_expire(self):
+        if self.ttl:
+            ctime = time.time()
+            if ctime > self.last_expire_time + (self.ttl / 2):
+                retry_call(self.client.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

clue/remote/datatypes/queues/comms.py

@@ -0,0 +1,62 @@
+import json
+
+import redis
+
+from clue.common.logging import get_logger
+from clue.remote.datatypes import decode, get_client, retry_call
+
+logger = get_logger(__file__)
+
+
+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:
+                logger.warning("No connection to Redis, reconnecting...")
+                self._connected = False
+                retried = True
+            finally:
+                if self._connected and retried:
+                    logger.info("Reconnected to Redis!")
+                    retried = False
+
+    def publish(self, message):
+        for name in self.names:
+            retry_call(self.c.publish, name, json.dumps(message))