black shift

This commit is contained in:
2025-04-22 11:10:29 +03:00
parent d7f1da8de8
commit e5f88f2eb4
30 changed files with 671 additions and 521 deletions

View File

@@ -8,17 +8,17 @@ 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
"""
@@ -26,113 +26,124 @@ class RedisPublisher:
# 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}
data={"recipients": recipient_count},
)
except Exception as e:
return RedisResponse(
status=False,
message=f"Failed to publish message to {channel}.",
error=str(e)
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:
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}."
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)
status=False, message=f"Failed to subscribe to {channel}.", error=str(e)
)
def psubscribe(self, pattern: str, callback: Callable[[Dict], None]) -> RedisResponse:
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}."
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)
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'):
if message["type"] in ("subscribe", "psubscribe"):
return
# Parse JSON if the message is a JSON string
data = message['data']
data = message["data"]
if isinstance(data, bytes):
data = data.decode('utf-8')
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
})
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
"""
@@ -140,50 +151,41 @@ class RedisSubscriber:
if in_thread:
thread = Thread(target=self._listen_thread, daemon=True)
thread.start()
self.active_threads['listener'] = thread
self.active_threads["listener"] = thread
return RedisResponse(
status=True,
message="Listening thread started successfully."
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)."
status=True, message="Listening started successfully (blocking)."
)
except Exception as e:
return RedisResponse(
status=False,
message="Failed to start listening.",
error=str(e)
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."
)
return RedisResponse(status=True, message="Successfully stopped listening.")
except Exception as e:
return RedisResponse(
status=False,
message="Failed to stop listening.",
error=str(e)
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
"""
@@ -194,24 +196,21 @@ class RedisSubscriber:
else:
self.pubsub.unsubscribe()
message = "Successfully unsubscribed from all channels."
return RedisResponse(
status=True,
message=message
)
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)
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
"""
@@ -222,24 +221,21 @@ class RedisSubscriber:
else:
self.pubsub.punsubscribe()
message = "Successfully unsubscribed from all patterns."
return RedisResponse(
status=True,
message=message
)
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)
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)

View File

@@ -15,6 +15,7 @@ CHANNEL_WRITER = "chain:writer"
# Flag to control the demo
running = True
def generate_mock_data():
"""Generate a mock message with UUID, timestamp, and sample data."""
return {
@@ -24,40 +25,43 @@ def generate_mock_data():
"data": {
"value": f"Sample data {int(time.time())}",
"status": "new",
"counter": 0
}
"counter": 0,
},
}
def reader_function():
"""
First function in the chain.
Generates mock data and publishes to the reader channel.
"""
print("[READER] Function started")
while running:
# Generate mock data
message = generate_mock_data()
start_time = time.time()
message["start_time"] = start_time
# Publish to reader channel
result = redis_pubsub.publisher.publish(CHANNEL_READER, message)
if result.status:
print(f"[READER] {time.time():.6f} | Published UUID: {message['uuid']}")
else:
print(f"[READER] Publish error: {result.error}")
# Wait before generating next message
time.sleep(2)
def processor_function():
"""
Second function in the chain.
Subscribes to reader channel, processes messages, and publishes to processor channel.
"""
print("[PROCESSOR] Function started")
def on_reader_message(message):
# The message structure from the subscriber has 'data' containing our actual message
# If data is a string, parse it as JSON
@@ -68,47 +72,51 @@ def processor_function():
except json.JSONDecodeError as e:
print(f"[PROCESSOR] Error parsing message data: {e}")
return
# Check if stage is 'red' before processing
if data.get("stage") == "red":
# Process the message
data["processor_timestamp"] = datetime.now().isoformat()
data["data"]["status"] = "processed"
data["data"]["counter"] += 1
# Update stage to 'processed'
data["stage"] = "processed"
# Add some processing metadata
data["processing"] = {
"duration_ms": 150, # Mock processing time
"processor_id": "main-processor"
"processor_id": "main-processor",
}
# Publish to processor channel
result = redis_pubsub.publisher.publish(CHANNEL_PROCESSOR, data)
if result.status:
print(f"[PROCESSOR] {time.time():.6f} | Received UUID: {data['uuid']} | Published UUID: {data['uuid']}")
print(
f"[PROCESSOR] {time.time():.6f} | Received UUID: {data['uuid']} | Published UUID: {data['uuid']}"
)
else:
print(f"[PROCESSOR] Publish error: {result.error}")
else:
print(f"[PROCESSOR] Skipped message: {data['uuid']} (stage is not 'red')")
# Subscribe to reader channel
result = redis_pubsub.subscriber.subscribe(CHANNEL_READER, on_reader_message)
if result.status:
print(f"[PROCESSOR] Subscribed to channel: {CHANNEL_READER}")
else:
print(f"[PROCESSOR] Subscribe error: {result.error}")
def writer_function():
"""
Third function in the chain.
Subscribes to processor channel and performs final processing.
"""
print("[WRITER] Function started")
def on_processor_message(message):
# The message structure from the subscriber has 'data' containing our actual message
# If data is a string, parse it as JSON
@@ -119,42 +127,45 @@ def writer_function():
except json.JSONDecodeError as e:
print(f"[WRITER] Error parsing message data: {e}")
return
# Check if stage is 'processed' before processing
if data.get("stage") == "processed":
# Process the message
data["writer_timestamp"] = datetime.now().isoformat()
data["data"]["status"] = "completed"
data["data"]["counter"] += 1
# Update stage to 'completed'
data["stage"] = "completed"
# Add some writer metadata
data["storage"] = {
"location": "main-db",
"partition": "events-2025-04"
}
data["storage"] = {"location": "main-db", "partition": "events-2025-04"}
# Calculate elapsed time if start_time is available
current_time = time.time()
elapsed_ms = ""
if "start_time" in data:
elapsed_ms = f" | Elapsed: {(current_time - data['start_time']) * 1000:.2f}ms"
elapsed_ms = (
f" | Elapsed: {(current_time - data['start_time']) * 1000:.2f}ms"
)
# Optionally publish to writer channel for any downstream listeners
result = redis_pubsub.publisher.publish(CHANNEL_WRITER, data)
if result.status:
print(f"[WRITER] {current_time:.6f} | Received UUID: {data['uuid']} | Published UUID: {data['uuid']}{elapsed_ms}")
print(
f"[WRITER] {current_time:.6f} | Received UUID: {data['uuid']} | Published UUID: {data['uuid']}{elapsed_ms}"
)
else:
print(f"[WRITER] Publish error: {result.error}")
else:
print(f"[WRITER] Skipped message: {data['uuid']} (stage is not 'processed')")
print(
f"[WRITER] Skipped message: {data['uuid']} (stage is not 'processed')"
)
# Subscribe to processor channel
result = redis_pubsub.subscriber.subscribe(CHANNEL_PROCESSOR, on_processor_message)
if result.status:
print(f"[WRITER] Subscribed to channel: {CHANNEL_PROCESSOR}")
else:
@@ -167,18 +178,18 @@ def run_demo():
print("Chain: READER → PROCESSOR → WRITER")
print(f"Channels: {CHANNEL_READER}{CHANNEL_PROCESSOR}{CHANNEL_WRITER}")
print("Format: [SERVICE] TIMESTAMP | Received/Published UUID | [Elapsed time]")
# Start the Redis subscriber listening thread
redis_pubsub.subscriber.start_listening()
# Start processor and writer functions (these subscribe to channels)
processor_function()
writer_function()
# Create a thread for the reader function (this publishes messages)
reader_thread = Thread(target=reader_function, daemon=True)
reader_thread.start()
# Keep the main thread alive
try:
while True: