ruffed

state/redis.py

@@ -0,0 +1,84 @@
+from bot.config import REDIS_URL
+import redis.asyncio as aredis
+import logging
+
+# Create a logger instance
+logger = logging.getLogger("state.redis")
+logging.basicConfig(level=logging.WARNING)
+
+
+class RedisService:
+    def __init__(self, uri=REDIS_URL):
+        self._uri: str = uri
+        self.pubsub_channels = []
+        self._client = None
+
+    async def connect(self):
+        self._client = aredis.Redis.from_url(self._uri, decode_responses=True)
+
+    async def disconnect(self):
+        if self._client:
+            await self._client.close()
+
+    async def execute(self, command, *args, **kwargs):
+        if self._client:
+            try:
+                logger.debug(f"{command}")  # {args[0]}") # {args} {kwargs}")
+                for arg in args:
+                    if isinstance(arg, dict):
+                        if arg.get("_sa_instance_state"):
+                            del arg["_sa_instance_state"]
+                r = await self._client.execute_command(command, *args, **kwargs)
+                # logger.debug(type(r))
+                # logger.debug(r)
+                return r
+            except Exception as e:
+                logger.error(e)
+
+    async def subscribe(self, *channels):
+        if self._client:
+            async with self._client.pubsub() as pubsub:
+                for channel in channels:
+                    await pubsub.subscribe(channel)
+                    self.pubsub_channels.append(channel)
+
+    async def unsubscribe(self, *channels):
+        if not self._client:
+            return
+        async with self._client.pubsub() as pubsub:
+            for channel in channels:
+                await pubsub.unsubscribe(channel)
+                self.pubsub_channels.remove(channel)
+
+    async def publish(self, channel, data):
+        if not self._client:
+            return
+        await self._client.publish(channel, data)
+
+    async def set(self, key, data, ex=None):
+        # Prepare the command arguments
+        args = [key, data]
+
+        # If an expiration time is provided, add it to the arguments
+        if ex is not None:
+            args.append("EX")
+            args.append(ex)
+
+        # Execute the command with the provided arguments
+        await self.execute("set", *args)
+
+    async def scan_iter(self, pattern="*"):
+        """Asynchronously iterate over keys matching the given pattern."""
+        cursor = "0"
+        while cursor != 0:
+            cursor, keys = await self._client.scan(cursor=cursor, match=pattern)
+            for key in keys:
+                yield key
+
+    async def get(self, key):
+        return await self.execute("get", key)
+
+
+redis = RedisService()
+
+__all__ = ["redis"]

state/scan.py

@@ -0,0 +1,34 @@
+from state.redis import redis
+import logging
+import math
+
+# Create a logger instance
+logger = logging.getLogger("state.scan")
+logging.basicConfig(level=logging.DEBUG)
+
+
+async def get_all_pattern(uid):
+    pattern = f"removed:{uid}:*"
+
+    # Create a dictionary to hold the keys and values
+    texts = []
+
+    # Use scan_iter to find all keys matching the pattern
+    async for key in redis.scan_iter(pattern):
+        # Fetch the value for each key
+        value = await redis.get(key)
+        if value:
+            texts.append(value.decode("utf-8"))
+
+    return texts
+
+
+async def get_average_pattern(pattern):
+    scores = []
+    async for key in redis.scan_iter(pattern):
+        scr = await redis.get(key)
+        if isinstance(scr, int):
+            scores.append(scr)
+    logger.debug(f"found {len(scores)} messages")
+    toxic_score = math.floor(sum(scores) / len(scores)) if scores else 0
+    return toxic_score