feat: add replica connection support and fix silent crash on query timeout

falkordb/asyncio/falkordb.py

@@ -4,9 +4,12 @@
 from redis.driver_info import DriverInfo
 from redis.exceptions import RedisError
 
+from redis.asyncio.cluster import RedisCluster  # type: ignore[import-not-found]
+
 from .._version import get_package_version
 from .cluster import Cluster_Conn, Is_Cluster
 from .graph import AsyncGraph
+from .sentinel import Is_Sentinel, Sentinel_Conn
 
 # config command
 UDF_CMD = "GRAPH.UDF"
@@ -115,6 +118,16 @@ def __init__(
             protocol=protocol,
         )
 
+        self.sentinel = None
+        self.service_name = None
+
+        if Is_Sentinel(conn):
+            self.sentinel, self.service_name = Sentinel_Conn(conn, ssl)
+            if read_from_replicas:
+                conn = self.sentinel.slave_for(self.service_name, ssl=ssl)
+            else:
+                conn = self.sentinel.master_for(self.service_name, ssl=ssl)
+
         if Is_Cluster(conn):
             conn = Cluster_Conn(
                 conn,
@@ -131,6 +144,30 @@ def __init__(
         self.flushdb = conn.flushdb
         self.execute_command = conn.execute_command
 
+    def get_replica_connection(self) -> Union[redis.Redis, RedisCluster]:
+        """
+        Returns a connection instance configured to read from replicas.
+
+        In Sentinel mode: Returns a connection to a Sentinel slave/replica.
+        In Cluster mode: Returns a cluster connection with read_from_replicas=True.
+        In Standalone mode: Returns the underlying connection.
+        """
+        if self.sentinel is not None:
+            return self.sentinel.slave_for(self.service_name)
+
+        if Is_Cluster(self.connection):
+            if isinstance(self.connection, RedisCluster):
+                return self.connection
+            return Cluster_Conn(self.connection, ssl=False, read_from_replicas=True)
+
+        return self.connection
+
+    async def _disconnect_connection(self):
+        """
+        Disconnects the underlying connection or pool to clear dirty socket state.
+        """
+        await self.connection.aclose()
+
     @classmethod
     def from_url(cls, url: str, **kwargs) -> "FalkorDB":
         """

falkordb/asyncio/graph.py

@@ -81,6 +81,10 @@ async def _query(  # type: ignore[override]
 
         # issue query
         try:
+            if not read_only:
+                self.schema._dirty_labels = True
+                self.schema._dirty_properties = True
+                self.schema._dirty_relations = True
             response = await self.execute_command(*command)
             query_result = QueryResult(self)
             await query_result.parse(response)
@@ -90,6 +94,16 @@ async def _query(  # type: ignore[override]
             # set client version and refresh local schema
             await self.schema.refresh(e.version)
             raise e
+        except Exception as e:
+            if "timed out" in str(e).lower() or "timeout" in str(e).lower():
+                await self._disconnect_connection()
+            raise e
+
+    async def _disconnect_connection(self):
+        """
+        Disconnects the underlying client connection to purge dirty socket state after a timeout.
+        """
+        await self.client._disconnect_connection()
 
     async def query(  # type: ignore[override]
         self,

falkordb/asyncio/graph_schema.py

@@ -37,6 +37,9 @@ def clear(self):
         self.labels = []
         self.properties = []
         self.relationships = []
+        self._dirty_labels = True
+        self._dirty_properties = True
+        self._dirty_relations = True
 
     async def refresh_labels(self) -> None:
         """
@@ -49,6 +52,7 @@ async def refresh_labels(self) -> None:
 
         result_set = (await self.graph.call_procedure(DB_LABELS)).result_set
         self.labels = [label[0] for label in result_set]
+        self._dirty_labels = False
 
     async def refresh_relations(self) -> None:
         """
@@ -61,6 +65,7 @@ async def refresh_relations(self) -> None:
 
         result_set = (await self.graph.call_procedure(DB_RELATIONSHIPTYPES)).result_set
         self.relationships = [r[0] for r in result_set]
+        self._dirty_relations = False
 
     async def refresh_properties(self) -> None:
         """
@@ -73,6 +78,7 @@ async def refresh_properties(self) -> None:
 
         result_set = (await self.graph.call_procedure(DB_PROPERTYKEYS)).result_set
         self.properties = [p[0] for p in result_set]
+        self._dirty_properties = False
 
     async def refresh(self, version: int) -> None:
         """
@@ -104,13 +110,9 @@ async def get_label(self, idx: int) -> str:
 
         """
 
-        try:
-            label = self.labels[idx]
-        except IndexError:
-            # refresh labels
+        if self._dirty_labels or not self.labels or idx >= len(self.labels):
             await self.refresh_labels()
-            label = self.labels[idx]
-        return label
+        return self.labels[idx]
 
     async def get_relation(self, idx: int) -> str:
         """
@@ -124,13 +126,13 @@ async def get_relation(self, idx: int) -> str:
 
         """
 
-        try:
-            r = self.relationships[idx]
-        except IndexError:
-            # refresh relationship types
+        if (
+            self._dirty_relations
+            or not self.relationships
+            or idx >= len(self.relationships)
+        ):
             await self.refresh_relations()
-            r = self.relationships[idx]
-        return r
+        return self.relationships[idx]
 
     async def get_property(self, idx: int) -> str:
         """
@@ -144,10 +146,6 @@ async def get_property(self, idx: int) -> str:
 
         """
 
-        try:
-            p = self.properties[idx]
-        except IndexError:
-            # refresh properties
+        if self._dirty_properties or not self.properties or idx >= len(self.properties):
             await self.refresh_properties()
-            p = self.properties[idx]
-        return p
+        return self.properties[idx]

falkordb/asyncio/sentinel.py

@@ -0,0 +1,61 @@
+import inspect
+
+import redis as sync_redis  # type: ignore[import-not-found]
+import redis.asyncio as redis  # type: ignore[import-not-found]
+from redis.asyncio.sentinel import Sentinel  # type: ignore[import-not-found]
+
+
+# detect if a connection is a sentinel
+def Is_Sentinel(conn: redis.Redis) -> bool:
+    pool = conn.connection_pool
+    kwargs = pool.connection_kwargs.copy()
+
+    kwargs["ssl"] = pool.connection_class is redis.SSLConnection
+
+    if pool.connection_class is redis.UnixDomainSocketConnection:
+        kwargs["unix_socket_path"] = kwargs.pop("path")
+
+    accepted = inspect.signature(sync_redis.Redis.__init__).parameters
+    kwargs = {k: v for k, v in kwargs.items() if k in accepted}
+
+    info = sync_redis.Redis(**kwargs).info(section="server")
+    return "redis_mode" in info and info["redis_mode"] == "sentinel"
+
+
+# create an async sentinel connection from a Redis connection
+def Sentinel_Conn(conn: redis.Redis, ssl: bool):
+    pool = conn.connection_pool
+    kwargs = pool.connection_kwargs.copy()
+
+    kwargs["ssl"] = pool.connection_class is redis.SSLConnection
+
+    if pool.connection_class is redis.UnixDomainSocketConnection:
+        kwargs["unix_socket_path"] = kwargs.pop("path")
+
+    accepted = inspect.signature(sync_redis.Redis.__init__).parameters
+    probe_kwargs = {k: v for k, v in kwargs.items() if k in accepted}
+
+    sync_conn = sync_redis.Redis(**probe_kwargs)
+    masters = sync_conn.sentinel_masters()
+
+    if len(masters) != 1:
+        raise Exception("Multiple masters, require service name")
+
+    service_name = list(masters.keys())[0]
+
+    host = kwargs.get("host", "localhost")
+    port = kwargs.get("port", 6379)
+    sentinels_conns = [(host, port)]
+
+    sentinel_kwargs = {}
+    if "username" in kwargs:
+        sentinel_kwargs["username"] = kwargs["username"]
+    if "password" in kwargs:
+        sentinel_kwargs["password"] = kwargs["password"]
+    if ssl:
+        sentinel_kwargs["ssl"] = True
+
+    return (
+        Sentinel(sentinels_conns, sentinel_kwargs=sentinel_kwargs, **kwargs),
+        service_name,
+    )

falkordb/falkordb.py

@@ -1,6 +1,7 @@
 from typing import List, Optional, Union
 
 import redis  # type: ignore[import-not-found]
+from redis.cluster import RedisCluster  # type: ignore[import-not-found]
 from redis.driver_info import DriverInfo
 from redis.exceptions import RedisError
 
@@ -129,9 +130,15 @@ def __init__(
             protocol=protocol,
         )
 
+        self.sentinel = None
+        self.service_name = None
+
         if Is_Sentinel(conn):
             self.sentinel, self.service_name = Sentinel_Conn(conn, ssl)
-            conn = self.sentinel.master_for(self.service_name, ssl=ssl)
+            if read_from_replicas:
+                conn = self.sentinel.slave_for(self.service_name, ssl=ssl)
+            else:
+                conn = self.sentinel.master_for(self.service_name, ssl=ssl)
 
         if Is_Cluster(conn):
             conn = Cluster_Conn(
@@ -151,6 +158,30 @@ def __init__(
         self.flushdb = conn.flushdb
         self.execute_command = conn.execute_command
 
+    def get_replica_connection(self) -> Union[redis.Redis, RedisCluster]:
+        """
+        Returns a connection instance configured to read from replicas.
+
+        In Sentinel mode: Returns a connection to a Sentinel slave/replica.
+        In Cluster mode: Returns a cluster connection with read_from_replicas=True.
+        In Standalone mode: Returns the underlying connection.
+        """
+        if self.sentinel is not None:
+            return self.sentinel.slave_for(self.service_name)
+
+        if Is_Cluster(self.connection):
+            if isinstance(self.connection, RedisCluster):
+                return self.connection
+            return Cluster_Conn(self.connection, ssl=False, read_from_replicas=True)
+
+        return self.connection
+
+    def _disconnect_connection(self):
+        """
+        Disconnects the underlying connection or pool to clear dirty socket state.
+        """
+        self.connection.close()
+
     @classmethod
     def from_url(cls, url: str, **kwargs) -> "FalkorDB":
         """

falkordb/graph.py

@@ -94,13 +94,27 @@ def _query(
 
         # issue query
         try:
+            if not read_only:
+                self.schema._dirty_labels = True
+                self.schema._dirty_properties = True
+                self.schema._dirty_relations = True
             response = self.execute_command(*command)
             return QueryResult(self, response)
         except SchemaVersionMismatchException as e:
             # client view over the graph schema is out of sync
             # set client version and refresh local schema
             self.schema.refresh(e.version)
             raise e
+        except Exception as e:
+            if "timed out" in str(e).lower() or "timeout" in str(e).lower():
+                self._disconnect_connection()
+            raise e
+
+    def _disconnect_connection(self):
+        """
+        Disconnects the underlying client connection to purge dirty socket state after a timeout.
+        """
+        self.client._disconnect_connection()
 
     def query(
         self,