Add role-based filtering and imporve code

services/notification/app/services/sse_service.py

@@ -5,11 +5,11 @@ Integrated within the notification service for alerts and recommendations
 """
 
 import asyncio
-from redis.asyncio import Redis
 import json
 from typing import Dict, Set, Any
 from datetime import datetime
 import structlog
+from shared.redis_utils import initialize_redis, get_redis_client, close_redis
 
 logger = structlog.get_logger()
 
@@ -18,18 +18,21 @@ class SSEService:
     Server-Sent Events service for real-time notifications
     Handles both alerts and recommendations through unified SSE streams
     """
-    
-    def __init__(self, redis_url: str):
-        self.redis_url = redis_url
+
+    def __init__(self):
         self.redis = None
+        self.redis_url = None
         self.active_connections: Dict[str, Set[asyncio.Queue]] = {}
         self.pubsub_tasks: Dict[str, asyncio.Task] = {}
-        
-    async def initialize(self):
+
+    async def initialize(self, redis_url: str):
         """Initialize Redis connection"""
         try:
-            self.redis = Redis.from_url(self.redis_url)
-            logger.info("SSE Service initialized with Redis connection")
+            self.redis_url = redis_url
+            # Initialize shared Redis connection for SSE
+            await initialize_redis(redis_url, db=0, max_connections=30)
+            self.redis = await get_redis_client()
+            logger.info("SSE Service initialized with shared Redis connection")
         except Exception as e:
             logger.error("Failed to initialize SSE service", error=str(e))
             raise
@@ -45,7 +48,7 @@ class SSEService:
                         await task
                     except asyncio.CancelledError:
                         pass
-            
+
             # Close all client connections
             for tenant_id, connections in self.active_connections.items():
                 for queue in connections.copy():
@@ -53,13 +56,12 @@ class SSEService:
                         await queue.put({"event": "shutdown", "data": json.dumps({"status": "server_shutdown"})})
                     except:
                         pass
-            
-            # Close Redis connection
-            if self.redis:
-                await self.redis.close()
-                
+
+            # Close shared Redis connection
+            await close_redis()
+
             logger.info("SSE Service shutdown completed")
-            
+
         except Exception as e:
             logger.error("Error during SSE shutdown", error=str(e))
             
@@ -124,32 +126,33 @@ class SSEService:
     
     async def _listen_to_tenant_channel(self, tenant_id: str):
         """Listen to Redis channel for tenant-specific items"""
+        pubsub = None
         try:
-            # Create a separate Redis connection for pubsub
-            pubsub_redis = Redis.from_url(self.redis_url)
-            pubsub = pubsub_redis.pubsub()
+            # Use the shared Redis client for pubsub
+            pubsub = self.redis.pubsub()
             channel = f"alerts:{tenant_id}"
             await pubsub.subscribe(channel)
-            
-            logger.info("Started listening to tenant channel", 
-                       tenant_id=tenant_id, 
+
+            logger.info("Started listening to tenant channel",
+                       tenant_id=tenant_id,
                        channel=channel)
-            
+
             async for message in pubsub.listen():
                 if message["type"] == "message":
                     # Broadcast to all connected clients for this tenant
                     await self.broadcast_to_tenant(tenant_id, message["data"])
-                    
+
         except asyncio.CancelledError:
             logger.info("Stopped listening to tenant channel", tenant_id=tenant_id)
         except Exception as e:
             logger.error("Error in pubsub listener", tenant_id=tenant_id, error=str(e))
         finally:
-            try:
-                await pubsub.unsubscribe(channel)
-                await pubsub_redis.close()
-            except:
-                pass
+            if pubsub:
+                try:
+                    await pubsub.unsubscribe(channel)
+                    await pubsub.close()
+                except:
+                    pass
     
     async def broadcast_to_tenant(self, tenant_id: str, message: str):
         """Broadcast message to all connected clients of a tenant"""