File size: 2,694 Bytes
1b7e88c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
from typing import Any, Optional
import threading
import os
import atexit
from omagent_core.utils.registry import registry
from pydantic import Field, PrivateAttr
from redis import ConnectionPool, Redis
from redislite import Redis as RedisLite
from .base import ConnectorBase
class SharedRedisLite:
"""Shared Redis implementation using RedisLite"""
_instance = None
_lock = threading.Lock()
_DB_FILE = os.path.join(os.getcwd(), "shared_redis.db")
@classmethod
def get_instance(cls):
with cls._lock:
if cls._instance is None:
# Create shared RedisLite instance using fixed database file
cls._instance = RedisLite(
dbfilename=cls._DB_FILE,
serverconfig={
'save': '', # Disable persistence
'appendonly': 'no' # Disable AOF
}
)
# Register cleanup function
atexit.register(cls._cleanup)
return cls._instance
@classmethod
def _cleanup(cls):
"""Clean up database files when the program exits"""
if cls._instance is not None:
cls._instance.shutdown()
# Remove database file and settings file
for suffix in ['', '.settings']:
file_path = cls._DB_FILE + suffix
if os.path.exists(file_path):
try:
os.remove(file_path)
except OSError:
pass
@registry.register_connector()
class RedisConnector(ConnectorBase):
host: str = Field(default="localhost")
port: int = Field(default=6379)
password: Optional[str] = Field(default=None)
username: Optional[str] = Field(default=None)
db: int = Field(default=0)
use_lite: bool = Field(default=True)
_client: Any = PrivateAttr(default=None)
def model_post_init(self, __context: Any) -> None:
if self.use_lite:
self._client = SharedRedisLite.get_instance()
else:
pool = ConnectionPool(
host=self.host,
port=self.port,
password=self.password,
username=self.username,
db=self.db,
decode_responses=False,
)
self._client = Redis(connection_pool=pool)
def check_connection(self) -> bool:
"""Check if Redis connection is valid by executing a simple ping command"""
try:
self._client.ping()
return True
except Exception as e:
raise ConnectionError(f"Redis connection failed: {str(e)}")
|