Refactor: nearly done phase 5
This commit is contained in:
ziesorx
parent
227e696ed6
commit
7a9a149955
12 changed files with 2750 additions and 105 deletions
|
|
@ -1 +1,10 @@
|
|||
# Storage module for Redis and PostgreSQL operations
|
||||
"""
|
||||
Storage module for the Python Detector Worker.
|
||||
|
||||
This module provides Redis and PostgreSQL operations for data persistence
|
||||
and caching in the detection pipeline.
|
||||
"""
|
||||
from .redis import RedisManager
|
||||
from .database import DatabaseManager
|
||||
|
||||
__all__ = ['RedisManager', 'DatabaseManager']
|
||||
357
core/storage/database.py
Normal file
357
core/storage/database.py
Normal file
|
|
@ -0,0 +1,357 @@
|
|||
"""
|
||||
Database Operations Module.
|
||||
Handles PostgreSQL operations for the detection pipeline.
|
||||
"""
|
||||
import psycopg2
|
||||
import psycopg2.extras
|
||||
from typing import Optional, Dict, Any
|
||||
import logging
|
||||
import uuid
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DatabaseManager:
|
||||
"""
|
||||
Manages PostgreSQL connections and operations for the detection pipeline.
|
||||
Handles database operations and schema management.
|
||||
"""
|
||||
|
||||
def __init__(self, config: Dict[str, Any]):
|
||||
"""
|
||||
Initialize database manager with configuration.
|
||||
|
||||
Args:
|
||||
config: Database configuration dictionary
|
||||
"""
|
||||
self.config = config
|
||||
self.connection: Optional[psycopg2.extensions.connection] = None
|
||||
|
||||
def connect(self) -> bool:
|
||||
"""
|
||||
Connect to PostgreSQL database.
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
try:
|
||||
self.connection = psycopg2.connect(
|
||||
host=self.config['host'],
|
||||
port=self.config['port'],
|
||||
database=self.config['database'],
|
||||
user=self.config['username'],
|
||||
password=self.config['password']
|
||||
)
|
||||
logger.info("PostgreSQL connection established successfully")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to connect to PostgreSQL: {e}")
|
||||
return False
|
||||
|
||||
def disconnect(self):
|
||||
"""Disconnect from PostgreSQL database."""
|
||||
if self.connection:
|
||||
self.connection.close()
|
||||
self.connection = None
|
||||
logger.info("PostgreSQL connection closed")
|
||||
|
||||
def is_connected(self) -> bool:
|
||||
"""
|
||||
Check if database connection is active.
|
||||
|
||||
Returns:
|
||||
True if connected, False otherwise
|
||||
"""
|
||||
try:
|
||||
if self.connection and not self.connection.closed:
|
||||
cur = self.connection.cursor()
|
||||
cur.execute("SELECT 1")
|
||||
cur.fetchone()
|
||||
cur.close()
|
||||
return True
|
||||
except:
|
||||
pass
|
||||
return False
|
||||
|
||||
def update_car_info(self, session_id: str, brand: str, model: str, body_type: str) -> bool:
|
||||
"""
|
||||
Update car information in the database.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
brand: Car brand
|
||||
model: Car model
|
||||
body_type: Car body type
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return False
|
||||
|
||||
try:
|
||||
cur = self.connection.cursor()
|
||||
query = """
|
||||
INSERT INTO car_frontal_info (session_id, car_brand, car_model, car_body_type, updated_at)
|
||||
VALUES (%s, %s, %s, %s, NOW())
|
||||
ON CONFLICT (session_id)
|
||||
DO UPDATE SET
|
||||
car_brand = EXCLUDED.car_brand,
|
||||
car_model = EXCLUDED.car_model,
|
||||
car_body_type = EXCLUDED.car_body_type,
|
||||
updated_at = NOW()
|
||||
"""
|
||||
cur.execute(query, (session_id, brand, model, body_type))
|
||||
self.connection.commit()
|
||||
cur.close()
|
||||
logger.info(f"Updated car info for session {session_id}: {brand} {model} ({body_type})")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to update car info: {e}")
|
||||
if self.connection:
|
||||
self.connection.rollback()
|
||||
return False
|
||||
|
||||
def execute_update(self, table: str, key_field: str, key_value: str, fields: Dict[str, str]) -> bool:
|
||||
"""
|
||||
Execute a dynamic update query on the database.
|
||||
|
||||
Args:
|
||||
table: Table name
|
||||
key_field: Primary key field name
|
||||
key_value: Primary key value
|
||||
fields: Dictionary of fields to update
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return False
|
||||
|
||||
try:
|
||||
cur = self.connection.cursor()
|
||||
|
||||
# Build the UPDATE query dynamically
|
||||
set_clauses = []
|
||||
values = []
|
||||
|
||||
for field, value in fields.items():
|
||||
if value == "NOW()":
|
||||
set_clauses.append(f"{field} = NOW()")
|
||||
else:
|
||||
set_clauses.append(f"{field} = %s")
|
||||
values.append(value)
|
||||
|
||||
# Add schema prefix if table doesn't already have it
|
||||
full_table_name = table if '.' in table else f"gas_station_1.{table}"
|
||||
|
||||
query = f"""
|
||||
INSERT INTO {full_table_name} ({key_field}, {', '.join(fields.keys())})
|
||||
VALUES (%s, {', '.join(['%s'] * len(fields))})
|
||||
ON CONFLICT ({key_field})
|
||||
DO UPDATE SET {', '.join(set_clauses)}
|
||||
"""
|
||||
|
||||
# Add key_value to the beginning of values list
|
||||
all_values = [key_value] + list(fields.values()) + values
|
||||
|
||||
cur.execute(query, all_values)
|
||||
self.connection.commit()
|
||||
cur.close()
|
||||
logger.info(f"Updated {table} for {key_field}={key_value}")
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to execute update on {table}: {e}")
|
||||
if self.connection:
|
||||
self.connection.rollback()
|
||||
return False
|
||||
|
||||
def create_car_frontal_info_table(self) -> bool:
|
||||
"""
|
||||
Create the car_frontal_info table in gas_station_1 schema if it doesn't exist.
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return False
|
||||
|
||||
try:
|
||||
# Since the database already exists, just verify connection
|
||||
cur = self.connection.cursor()
|
||||
|
||||
# Simple verification that the table exists
|
||||
cur.execute("""
|
||||
SELECT EXISTS (
|
||||
SELECT FROM information_schema.tables
|
||||
WHERE table_schema = 'gas_station_1'
|
||||
AND table_name = 'car_frontal_info'
|
||||
)
|
||||
""")
|
||||
|
||||
table_exists = cur.fetchone()[0]
|
||||
cur.close()
|
||||
|
||||
if table_exists:
|
||||
logger.info("Verified car_frontal_info table exists")
|
||||
return True
|
||||
else:
|
||||
logger.error("car_frontal_info table does not exist in the database")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to create car_frontal_info table: {e}")
|
||||
if self.connection:
|
||||
self.connection.rollback()
|
||||
return False
|
||||
|
||||
def insert_initial_detection(self, display_id: str, captured_timestamp: str, session_id: str = None) -> str:
|
||||
"""
|
||||
Insert initial detection record and return the session_id.
|
||||
|
||||
Args:
|
||||
display_id: Display identifier
|
||||
captured_timestamp: Timestamp of the detection
|
||||
session_id: Optional session ID, generates one if not provided
|
||||
|
||||
Returns:
|
||||
Session ID string or None on error
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return None
|
||||
|
||||
# Generate session_id if not provided
|
||||
if not session_id:
|
||||
session_id = str(uuid.uuid4())
|
||||
|
||||
try:
|
||||
# Ensure table exists
|
||||
if not self.create_car_frontal_info_table():
|
||||
logger.error("Failed to create/verify table before insertion")
|
||||
return None
|
||||
|
||||
cur = self.connection.cursor()
|
||||
insert_query = """
|
||||
INSERT INTO gas_station_1.car_frontal_info
|
||||
(display_id, captured_timestamp, session_id, license_character, license_type, car_brand, car_model, car_body_type)
|
||||
VALUES (%s, %s, %s, NULL, 'No model available', NULL, NULL, NULL)
|
||||
ON CONFLICT (session_id) DO NOTHING
|
||||
"""
|
||||
|
||||
cur.execute(insert_query, (display_id, captured_timestamp, session_id))
|
||||
self.connection.commit()
|
||||
cur.close()
|
||||
logger.info(f"Inserted initial detection record with session_id: {session_id}")
|
||||
return session_id
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to insert initial detection record: {e}")
|
||||
if self.connection:
|
||||
self.connection.rollback()
|
||||
return None
|
||||
|
||||
def get_session_info(self, session_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get session information from the database.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
Dictionary with session data or None if not found
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return None
|
||||
|
||||
try:
|
||||
cur = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
|
||||
query = "SELECT * FROM gas_station_1.car_frontal_info WHERE session_id = %s"
|
||||
cur.execute(query, (session_id,))
|
||||
result = cur.fetchone()
|
||||
cur.close()
|
||||
|
||||
if result:
|
||||
return dict(result)
|
||||
else:
|
||||
logger.debug(f"No session info found for session_id: {session_id}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to get session info: {e}")
|
||||
return None
|
||||
|
||||
def delete_session(self, session_id: str) -> bool:
|
||||
"""
|
||||
Delete session record from the database.
|
||||
|
||||
Args:
|
||||
session_id: Session identifier
|
||||
|
||||
Returns:
|
||||
True if successful, False otherwise
|
||||
"""
|
||||
if not self.is_connected():
|
||||
if not self.connect():
|
||||
return False
|
||||
|
||||
try:
|
||||
cur = self.connection.cursor()
|
||||
query = "DELETE FROM gas_station_1.car_frontal_info WHERE session_id = %s"
|
||||
cur.execute(query, (session_id,))
|
||||
rows_affected = cur.rowcount
|
||||
self.connection.commit()
|
||||
cur.close()
|
||||
|
||||
if rows_affected > 0:
|
||||
logger.info(f"Deleted session record: {session_id}")
|
||||
return True
|
||||
else:
|
||||
logger.warning(f"No session record found to delete: {session_id}")
|
||||
return False
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to delete session: {e}")
|
||||
if self.connection:
|
||||
self.connection.rollback()
|
||||
return False
|
||||
|
||||
def get_statistics(self) -> Dict[str, Any]:
|
||||
"""
|
||||
Get database statistics.
|
||||
|
||||
Returns:
|
||||
Dictionary with database statistics
|
||||
"""
|
||||
stats = {
|
||||
'connected': self.is_connected(),
|
||||
'host': self.config.get('host', 'unknown'),
|
||||
'port': self.config.get('port', 'unknown'),
|
||||
'database': self.config.get('database', 'unknown')
|
||||
}
|
||||
|
||||
if self.is_connected():
|
||||
try:
|
||||
cur = self.connection.cursor()
|
||||
|
||||
# Get table record count
|
||||
cur.execute("SELECT COUNT(*) FROM gas_station_1.car_frontal_info")
|
||||
stats['total_records'] = cur.fetchone()[0]
|
||||
|
||||
# Get recent records count (last hour)
|
||||
cur.execute("""
|
||||
SELECT COUNT(*) FROM gas_station_1.car_frontal_info
|
||||
WHERE created_at > NOW() - INTERVAL '1 hour'
|
||||
""")
|
||||
stats['recent_records'] = cur.fetchone()[0]
|
||||
|
||||
cur.close()
|
||||
except Exception as e:
|
||||
logger.warning(f"Failed to get database statistics: {e}")
|
||||
stats['error'] = str(e)
|
||||
|
||||
return stats
|
||||
478
core/storage/redis.py
Normal file
478
core/storage/redis.py
Normal file
|
|
@ -0,0 +1,478 @@
|
|||
"""
|
||||
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")
|
||||
Loading…
Add table
Add a link
Reference in a new issue