welcomecenterbot/utils/store.py

from bot.config import REDIS_URL
import asyncio
import redis.asyncio as aredis
import logging

# Create a logger instance
logger = logging.getLogger('store')
logging.basicConfig(level=logging.DEBUG)


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"]



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 = []
    scoring_msg_id = 0
    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