249 lines
8.4 KiB
Python
249 lines
8.4 KiB
Python
import json
|
|
from typing import Optional, Dict, Any, List, Callable, Union
|
|
from threading import Thread
|
|
|
|
from Controllers.Redis.connection import redis_cli
|
|
from Controllers.Redis.response import RedisResponse
|
|
|
|
|
|
class RedisPublisher:
|
|
"""Redis Publisher class for broadcasting messages to channels."""
|
|
|
|
def __init__(self, redis_client=redis_cli):
|
|
self.redis_client = redis_client
|
|
|
|
def publish(self, channel: str, message: Union[Dict, List, str]) -> RedisResponse:
|
|
"""Publish a message to a Redis channel.
|
|
|
|
Args:
|
|
channel: The channel to publish to
|
|
message: The message to publish (will be JSON serialized if dict or list)
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
# Convert dict/list to JSON string if needed
|
|
if isinstance(message, (dict, list)):
|
|
message = json.dumps(message)
|
|
|
|
# Publish the message
|
|
recipient_count = self.redis_client.publish(channel, message)
|
|
|
|
return RedisResponse(
|
|
status=True,
|
|
message=f"Message published successfully to {channel}.",
|
|
data={"recipients": recipient_count},
|
|
)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False,
|
|
message=f"Failed to publish message to {channel}.",
|
|
error=str(e),
|
|
)
|
|
|
|
|
|
class RedisSubscriber:
|
|
"""Redis Subscriber class for listening to channels."""
|
|
|
|
def __init__(self, redis_client=redis_cli):
|
|
self.redis_client = redis_client
|
|
self.pubsub = self.redis_client.pubsub()
|
|
self.active_threads = {}
|
|
|
|
def subscribe(
|
|
self, channel: str, callback: Callable[[Dict], None]
|
|
) -> RedisResponse:
|
|
"""Subscribe to a Redis channel with a callback function.
|
|
|
|
Args:
|
|
channel: The channel to subscribe to
|
|
callback: Function to call when a message is received
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
# Subscribe to the channel
|
|
self.pubsub.subscribe(**{channel: self._message_handler(callback)})
|
|
|
|
return RedisResponse(
|
|
status=True, message=f"Successfully subscribed to {channel}."
|
|
)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False, message=f"Failed to subscribe to {channel}.", error=str(e)
|
|
)
|
|
|
|
def psubscribe(
|
|
self, pattern: str, callback: Callable[[Dict], None]
|
|
) -> RedisResponse:
|
|
"""Subscribe to Redis channels matching a pattern.
|
|
|
|
Args:
|
|
pattern: The pattern to subscribe to (e.g., 'user.*')
|
|
callback: Function to call when a message is received
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
# Subscribe to the pattern
|
|
self.pubsub.psubscribe(**{pattern: self._message_handler(callback)})
|
|
|
|
return RedisResponse(
|
|
status=True, message=f"Successfully pattern-subscribed to {pattern}."
|
|
)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False,
|
|
message=f"Failed to pattern-subscribe to {pattern}.",
|
|
error=str(e),
|
|
)
|
|
|
|
def _message_handler(self, callback: Callable[[Dict], None]):
|
|
"""Create a message handler function for the subscription."""
|
|
|
|
def handler(message):
|
|
# Skip subscription confirmation messages
|
|
if message["type"] in ("subscribe", "psubscribe"):
|
|
return
|
|
|
|
# Parse JSON if the message is a JSON string
|
|
data = message["data"]
|
|
if isinstance(data, bytes):
|
|
data = data.decode("utf-8")
|
|
try:
|
|
data = json.loads(data)
|
|
except json.JSONDecodeError:
|
|
# Not JSON, keep as is
|
|
pass
|
|
|
|
# Call the callback with the message data
|
|
callback(
|
|
{
|
|
"channel": (
|
|
message.get("channel", b"").decode("utf-8")
|
|
if isinstance(message.get("channel", b""), bytes)
|
|
else message.get("channel", "")
|
|
),
|
|
"pattern": (
|
|
message.get("pattern", b"").decode("utf-8")
|
|
if isinstance(message.get("pattern", b""), bytes)
|
|
else message.get("pattern", "")
|
|
),
|
|
"data": data,
|
|
}
|
|
)
|
|
|
|
return handler
|
|
|
|
def start_listening(self, in_thread: bool = True) -> RedisResponse:
|
|
"""Start listening for messages on subscribed channels.
|
|
|
|
Args:
|
|
in_thread: If True, start listening in a separate thread
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
if in_thread:
|
|
thread = Thread(target=self._listen_thread, daemon=True)
|
|
thread.start()
|
|
self.active_threads["listener"] = thread
|
|
return RedisResponse(
|
|
status=True, message="Listening thread started successfully."
|
|
)
|
|
else:
|
|
# This will block the current thread
|
|
self._listen_thread()
|
|
return RedisResponse(
|
|
status=True, message="Listening started successfully (blocking)."
|
|
)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False, message="Failed to start listening.", error=str(e)
|
|
)
|
|
|
|
def _listen_thread(self):
|
|
"""Thread function for listening to messages."""
|
|
self.pubsub.run_in_thread(sleep_time=0.01)
|
|
|
|
def stop_listening(self) -> RedisResponse:
|
|
"""Stop listening for messages."""
|
|
try:
|
|
self.pubsub.close()
|
|
return RedisResponse(status=True, message="Successfully stopped listening.")
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False, message="Failed to stop listening.", error=str(e)
|
|
)
|
|
|
|
def unsubscribe(self, channel: Optional[str] = None) -> RedisResponse:
|
|
"""Unsubscribe from a channel or all channels.
|
|
|
|
Args:
|
|
channel: The channel to unsubscribe from, or None for all channels
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
if channel:
|
|
self.pubsub.unsubscribe(channel)
|
|
message = f"Successfully unsubscribed from {channel}."
|
|
else:
|
|
self.pubsub.unsubscribe()
|
|
message = "Successfully unsubscribed from all channels."
|
|
|
|
return RedisResponse(status=True, message=message)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False,
|
|
message=f"Failed to unsubscribe from {'channel' if channel else 'all channels'}.",
|
|
error=str(e),
|
|
)
|
|
|
|
def punsubscribe(self, pattern: Optional[str] = None) -> RedisResponse:
|
|
"""Unsubscribe from a pattern or all patterns.
|
|
|
|
Args:
|
|
pattern: The pattern to unsubscribe from, or None for all patterns
|
|
|
|
Returns:
|
|
RedisResponse with status and message
|
|
"""
|
|
try:
|
|
if pattern:
|
|
self.pubsub.punsubscribe(pattern)
|
|
message = f"Successfully unsubscribed from pattern {pattern}."
|
|
else:
|
|
self.pubsub.punsubscribe()
|
|
message = "Successfully unsubscribed from all patterns."
|
|
|
|
return RedisResponse(status=True, message=message)
|
|
except Exception as e:
|
|
return RedisResponse(
|
|
status=False,
|
|
message=f"Failed to unsubscribe from {'pattern' if pattern else 'all patterns'}.",
|
|
error=str(e),
|
|
)
|
|
|
|
|
|
class RedisPubSub:
|
|
"""Singleton class that provides both publisher and subscriber functionality."""
|
|
|
|
_instance = None
|
|
|
|
def __new__(cls):
|
|
if cls._instance is None:
|
|
cls._instance = super(RedisPubSub, cls).__new__(cls)
|
|
cls._instance.publisher = RedisPublisher()
|
|
cls._instance.subscriber = RedisSubscriber()
|
|
return cls._instance
|
|
|
|
|
|
# Create a singleton instance
|
|
redis_pubsub = RedisPubSub()
|