478 lines
No EOL
16 KiB
Python
478 lines
No EOL
16 KiB
Python
"""
|
|
Redis Operations Module.
|
|
Handles Redis connections, image storage, and pub/sub messaging.
|
|
"""
|
|
import logging
|
|
import json
|
|
import time
|
|
from typing import Optional, Dict, Any, Union
|
|
import asyncio
|
|
import cv2
|
|
import numpy as np
|
|
import redis.asyncio as redis
|
|
from redis.exceptions import ConnectionError, TimeoutError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class RedisManager:
|
|
"""
|
|
Manages Redis connections and operations for the detection pipeline.
|
|
Handles image storage with region cropping and pub/sub messaging.
|
|
"""
|
|
|
|
def __init__(self, redis_config: Dict[str, Any]):
|
|
"""
|
|
Initialize Redis manager with configuration.
|
|
|
|
Args:
|
|
redis_config: Redis configuration dictionary
|
|
"""
|
|
self.config = redis_config
|
|
self.redis_client: Optional[redis.Redis] = None
|
|
|
|
# Connection parameters
|
|
self.host = redis_config.get('host', 'localhost')
|
|
self.port = redis_config.get('port', 6379)
|
|
self.password = redis_config.get('password')
|
|
self.db = redis_config.get('db', 0)
|
|
self.decode_responses = redis_config.get('decode_responses', True)
|
|
|
|
# Connection pool settings
|
|
self.max_connections = redis_config.get('max_connections', 10)
|
|
self.socket_timeout = redis_config.get('socket_timeout', 5)
|
|
self.socket_connect_timeout = redis_config.get('socket_connect_timeout', 5)
|
|
self.health_check_interval = redis_config.get('health_check_interval', 30)
|
|
|
|
# Statistics
|
|
self.stats = {
|
|
'images_stored': 0,
|
|
'messages_published': 0,
|
|
'connection_errors': 0,
|
|
'operations_successful': 0,
|
|
'operations_failed': 0
|
|
}
|
|
|
|
logger.info(f"RedisManager initialized for {self.host}:{self.port}")
|
|
|
|
async def initialize(self) -> bool:
|
|
"""
|
|
Initialize Redis connection and test connectivity.
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
# Validate configuration
|
|
if not self._validate_config():
|
|
return False
|
|
|
|
# Create Redis connection
|
|
self.redis_client = redis.Redis(
|
|
host=self.host,
|
|
port=self.port,
|
|
password=self.password,
|
|
db=self.db,
|
|
decode_responses=self.decode_responses,
|
|
max_connections=self.max_connections,
|
|
socket_timeout=self.socket_timeout,
|
|
socket_connect_timeout=self.socket_connect_timeout,
|
|
health_check_interval=self.health_check_interval
|
|
)
|
|
|
|
# Test connection
|
|
await self.redis_client.ping()
|
|
logger.info(f"Successfully connected to Redis at {self.host}:{self.port}")
|
|
return True
|
|
|
|
except ConnectionError as e:
|
|
logger.error(f"Failed to connect to Redis: {e}")
|
|
self.stats['connection_errors'] += 1
|
|
return False
|
|
except Exception as e:
|
|
logger.error(f"Error initializing Redis connection: {e}", exc_info=True)
|
|
self.stats['connection_errors'] += 1
|
|
return False
|
|
|
|
def _validate_config(self) -> bool:
|
|
"""
|
|
Validate Redis configuration parameters.
|
|
|
|
Returns:
|
|
True if valid, False otherwise
|
|
"""
|
|
required_fields = ['host', 'port']
|
|
for field in required_fields:
|
|
if field not in self.config:
|
|
logger.error(f"Missing required Redis config field: {field}")
|
|
return False
|
|
|
|
if not isinstance(self.port, int) or self.port <= 0:
|
|
logger.error(f"Invalid Redis port: {self.port}")
|
|
return False
|
|
|
|
return True
|
|
|
|
async def is_connected(self) -> bool:
|
|
"""
|
|
Check if Redis connection is active.
|
|
|
|
Returns:
|
|
True if connected, False otherwise
|
|
"""
|
|
try:
|
|
if self.redis_client:
|
|
await self.redis_client.ping()
|
|
return True
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
async def save_image(self,
|
|
key: str,
|
|
image: np.ndarray,
|
|
expire_seconds: Optional[int] = None,
|
|
image_format: str = 'jpeg',
|
|
quality: int = 90) -> bool:
|
|
"""
|
|
Save image to Redis with optional expiration.
|
|
|
|
Args:
|
|
key: Redis key for the image
|
|
image: Image array to save
|
|
expire_seconds: Optional expiration time in seconds
|
|
image_format: Image format ('jpeg' or 'png')
|
|
quality: JPEG quality (1-100)
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
# Encode image
|
|
encoded_image = self._encode_image(image, image_format, quality)
|
|
if encoded_image is None:
|
|
logger.error("Failed to encode image")
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
# Save to Redis
|
|
if expire_seconds:
|
|
await self.redis_client.setex(key, expire_seconds, encoded_image)
|
|
logger.debug(f"Saved image to Redis with key: {key} (expires in {expire_seconds}s)")
|
|
else:
|
|
await self.redis_client.set(key, encoded_image)
|
|
logger.debug(f"Saved image to Redis with key: {key}")
|
|
|
|
self.stats['images_stored'] += 1
|
|
self.stats['operations_successful'] += 1
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error saving image to Redis: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
async def get_image(self, key: str) -> Optional[np.ndarray]:
|
|
"""
|
|
Retrieve image from Redis.
|
|
|
|
Args:
|
|
key: Redis key for the image
|
|
|
|
Returns:
|
|
Image array or None if not found
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return None
|
|
|
|
# Get image data from Redis
|
|
image_data = await self.redis_client.get(key)
|
|
if image_data is None:
|
|
logger.debug(f"Image not found for key: {key}")
|
|
return None
|
|
|
|
# Decode image
|
|
image_array = np.frombuffer(image_data, np.uint8)
|
|
image = cv2.imdecode(image_array, cv2.IMREAD_COLOR)
|
|
|
|
if image is not None:
|
|
logger.debug(f"Retrieved image from Redis with key: {key}")
|
|
self.stats['operations_successful'] += 1
|
|
return image
|
|
else:
|
|
logger.error(f"Failed to decode image for key: {key}")
|
|
self.stats['operations_failed'] += 1
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error retrieving image from Redis: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return None
|
|
|
|
async def delete_image(self, key: str) -> bool:
|
|
"""
|
|
Delete image from Redis.
|
|
|
|
Args:
|
|
key: Redis key for the image
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
result = await self.redis_client.delete(key)
|
|
if result > 0:
|
|
logger.debug(f"Deleted image from Redis with key: {key}")
|
|
self.stats['operations_successful'] += 1
|
|
return True
|
|
else:
|
|
logger.debug(f"Image not found for deletion: {key}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error deleting image from Redis: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
async def publish_message(self, channel: str, message: Union[str, Dict]) -> int:
|
|
"""
|
|
Publish message to Redis channel.
|
|
|
|
Args:
|
|
channel: Redis channel name
|
|
message: Message to publish (string or dict)
|
|
|
|
Returns:
|
|
Number of subscribers that received the message, -1 on error
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return -1
|
|
|
|
# Convert dict to JSON string if needed
|
|
if isinstance(message, dict):
|
|
message_str = json.dumps(message)
|
|
else:
|
|
message_str = str(message)
|
|
|
|
# Test connection before publishing
|
|
await self.redis_client.ping()
|
|
|
|
# Publish message
|
|
result = await self.redis_client.publish(channel, message_str)
|
|
|
|
logger.info(f"Published message to Redis channel '{channel}': {message_str}")
|
|
logger.info(f"Redis publish result (subscribers count): {result}")
|
|
|
|
if result == 0:
|
|
logger.warning(f"No subscribers listening to channel '{channel}'")
|
|
else:
|
|
logger.info(f"Message delivered to {result} subscriber(s)")
|
|
|
|
self.stats['messages_published'] += 1
|
|
self.stats['operations_successful'] += 1
|
|
return result
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error publishing message to Redis: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return -1
|
|
|
|
async def subscribe_to_channel(self, channel: str, callback=None):
|
|
"""
|
|
Subscribe to Redis channel (for future use).
|
|
|
|
Args:
|
|
channel: Redis channel name
|
|
callback: Optional callback function for messages
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
return
|
|
|
|
pubsub = self.redis_client.pubsub()
|
|
await pubsub.subscribe(channel)
|
|
|
|
logger.info(f"Subscribed to Redis channel: {channel}")
|
|
|
|
if callback:
|
|
async for message in pubsub.listen():
|
|
if message['type'] == 'message':
|
|
try:
|
|
await callback(message['data'])
|
|
except Exception as e:
|
|
logger.error(f"Error in message callback: {e}")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error subscribing to Redis channel: {e}", exc_info=True)
|
|
|
|
async def set_key(self, key: str, value: Union[str, bytes], expire_seconds: Optional[int] = None) -> bool:
|
|
"""
|
|
Set a key-value pair in Redis.
|
|
|
|
Args:
|
|
key: Redis key
|
|
value: Value to store
|
|
expire_seconds: Optional expiration time in seconds
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
if expire_seconds:
|
|
await self.redis_client.setex(key, expire_seconds, value)
|
|
else:
|
|
await self.redis_client.set(key, value)
|
|
|
|
logger.debug(f"Set Redis key: {key}")
|
|
self.stats['operations_successful'] += 1
|
|
return True
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error setting Redis key: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
async def get_key(self, key: str) -> Optional[Union[str, bytes]]:
|
|
"""
|
|
Get value for a Redis key.
|
|
|
|
Args:
|
|
key: Redis key
|
|
|
|
Returns:
|
|
Value or None if not found
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return None
|
|
|
|
value = await self.redis_client.get(key)
|
|
if value is not None:
|
|
logger.debug(f"Retrieved Redis key: {key}")
|
|
self.stats['operations_successful'] += 1
|
|
|
|
return value
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error getting Redis key: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return None
|
|
|
|
async def delete_key(self, key: str) -> bool:
|
|
"""
|
|
Delete a Redis key.
|
|
|
|
Args:
|
|
key: Redis key
|
|
|
|
Returns:
|
|
True if successful, False otherwise
|
|
"""
|
|
try:
|
|
if not self.redis_client:
|
|
logger.error("Redis client not initialized")
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
result = await self.redis_client.delete(key)
|
|
if result > 0:
|
|
logger.debug(f"Deleted Redis key: {key}")
|
|
self.stats['operations_successful'] += 1
|
|
return True
|
|
else:
|
|
logger.debug(f"Redis key not found: {key}")
|
|
return False
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error deleting Redis key: {e}", exc_info=True)
|
|
self.stats['operations_failed'] += 1
|
|
return False
|
|
|
|
def _encode_image(self, image: np.ndarray, image_format: str, quality: int) -> Optional[bytes]:
|
|
"""
|
|
Encode image to bytes for Redis storage.
|
|
|
|
Args:
|
|
image: Image array
|
|
image_format: Image format ('jpeg' or 'png')
|
|
quality: JPEG quality (1-100)
|
|
|
|
Returns:
|
|
Encoded image bytes or None on error
|
|
"""
|
|
try:
|
|
format_lower = image_format.lower()
|
|
|
|
if format_lower == 'jpeg' or format_lower == 'jpg':
|
|
encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
|
|
success, buffer = cv2.imencode('.jpg', image, encode_params)
|
|
elif format_lower == 'png':
|
|
success, buffer = cv2.imencode('.png', image)
|
|
else:
|
|
logger.warning(f"Unknown image format '{image_format}', using JPEG")
|
|
encode_params = [cv2.IMWRITE_JPEG_QUALITY, quality]
|
|
success, buffer = cv2.imencode('.jpg', image, encode_params)
|
|
|
|
if success:
|
|
return buffer.tobytes()
|
|
else:
|
|
logger.error(f"Failed to encode image as {image_format}")
|
|
return None
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error encoding image: {e}", exc_info=True)
|
|
return None
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""
|
|
Get Redis manager statistics.
|
|
|
|
Returns:
|
|
Dictionary with statistics
|
|
"""
|
|
return {
|
|
**self.stats,
|
|
'connected': self.redis_client is not None,
|
|
'host': self.host,
|
|
'port': self.port,
|
|
'db': self.db
|
|
}
|
|
|
|
def cleanup(self):
|
|
"""Cleanup Redis connection."""
|
|
if self.redis_client:
|
|
# Note: redis.asyncio doesn't have a synchronous close method
|
|
# The connection will be closed when the event loop shuts down
|
|
self.redis_client = None
|
|
logger.info("Redis connection cleaned up")
|
|
|
|
async def aclose(self):
|
|
"""Async cleanup for Redis connection."""
|
|
if self.redis_client:
|
|
await self.redis_client.aclose()
|
|
self.redis_client = None
|
|
logger.info("Redis connection closed") |