114 lines
3.4 KiB
Python
114 lines
3.4 KiB
Python
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 |