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

falkordb/asyncio/cluster.py

@@ -9,32 +9,98 @@
 
 # detect if a connection is a cluster
 def Is_Cluster(conn: redis.Redis):
+    try:
+        pool = conn.connection_pool
+        kwargs = pool.connection_kwargs.copy()
 
-    pool = conn.connection_pool
-    kwargs = pool.connection_kwargs.copy()
+        # Check if the connection is using SSL and add it
+        # this propery is not kept in the connection_kwargs
+        kwargs["ssl"] = pool.connection_class is redis.SSLConnection
 
-    # Check if the connection is using SSL and add it
-    # this propery is not kept in the connection_kwargs
-    kwargs["ssl"] = pool.connection_class is redis.SSLConnection
+        # The async Unix-domain-socket pool stores the socket path under "path",
+        # but the synchronous redis.Redis constructor expects "unix_socket_path".
+        # Translate the key so the sync probe can be built for unix:// connections.
+        if pool.connection_class is redis.UnixDomainSocketConnection:
+            kwargs["unix_socket_path"] = kwargs.pop("path")
 
-    # The async Unix-domain-socket pool stores the socket path under "path",
-    # but the synchronous redis.Redis constructor expects "unix_socket_path".
-    # Translate the key so the sync probe can be built for unix:// connections.
-    if pool.connection_class is redis.UnixDomainSocketConnection:
-        kwargs["unix_socket_path"] = kwargs.pop("path")
+        # Keep only the parameters the synchronous constructor actually accepts.
+        # redis-py stores internal state in ``connection_kwargs`` that is not part
+        # of the ``Redis.__init__`` signature — redis 8.1.0 added ``himport_registry``
+        # there — and forwarding those raises TypeError before any I/O happens.
+        accepted = inspect.signature(sync_redis.Redis.__init__).parameters
+        kwargs = {k: v for k, v in kwargs.items() if k in accepted}
 
-    # Keep only the parameters the synchronous constructor actually accepts.
-    # redis-py stores internal state in ``connection_kwargs`` that is not part
-    # of the ``Redis.__init__`` signature — redis 8.1.0 added ``himport_registry``
-    # there — and forwarding those raises TypeError before any I/O happens.
-    accepted = inspect.signature(sync_redis.Redis.__init__).parameters
-    kwargs = {k: v for k, v in kwargs.items() if k in accepted}
+        # Create a synchronous Redis client with the same parameters
+        # as the connection pool just to keep Is_Cluster synchronous
+        info = sync_redis.Redis(**kwargs).info(section="server")
 
-    # Create a synchronous Redis client with the same parameters
-    # as the connection pool just to keep Is_Cluster synchronous
-    info = sync_redis.Redis(**kwargs).info(section="server")
+        return "redis_mode" in info and info["redis_mode"] == "cluster"
+    except (
+        ConnectionError,
+        ConnectionRefusedError,
+        OSError,
+        sync_redis.exceptions.ConnectionError,
+    ):
+        raise
+    except Exception:
+        return False
 
-    return "redis_mode" in info and info["redis_mode"] == "cluster"
+
+def _str_val(v):
+    if isinstance(v, bytes):
+        return v.decode("utf-8")
+    return str(v)
+
+
+def parse_cluster_slots(raw_slots):
+    """
+    Parses CLUSTER SLOTS output into a list of primary and replica shard mappings.
+    """
+    shards_map = {}
+
+    for item in raw_slots:
+        if not item or len(item) < 3:
+            continue
+        start_slot = int(item[0])
+        end_slot = int(item[1])
+
+        p_info = item[2]
+        p_host = _str_val(p_info[0])
+        p_port = int(p_info[1])
+        p_id = _str_val(p_info[2]) if len(p_info) > 2 else f"{p_host}:{p_port}"
+
+        primary_key = (p_host, p_port)
+        if primary_key not in shards_map:
+            replicas = []
+            for r_info in item[3:]:
+                if not r_info or len(r_info) < 2:
+                    continue
+                r_host = _str_val(r_info[0])
+                r_port = int(r_info[1])
+                r_id = _str_val(r_info[2]) if len(r_info) > 2 else f"{r_host}:{r_port}"
+                replicas.append(
+                    {
+                        "id": r_id,
+                        "host": r_host,
+                        "port": r_port,
+                        "endpoint": f"{r_host}:{r_port}",
+                    }
+                )
+
+            shards_map[primary_key] = {
+                "primary": {
+                    "id": p_id,
+                    "host": p_host,
+                    "port": p_port,
+                    "endpoint": f"{p_host}:{p_port}",
+                },
+                "replicas": replicas,
+                "slots": [(start_slot, end_slot)],
+            }
+        else:
+            shards_map[primary_key]["slots"].append((start_slot, end_slot))
+
+    return list(shards_map.values())
 
 
 # create a cluster connection from a Redis connection

falkordb/asyncio/falkordb.py

@@ -1,12 +1,14 @@
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Union
 
 import redis.asyncio as redis  # type: ignore[import-not-found]
+from redis.asyncio.cluster import RedisCluster  # type: ignore[import-not-found]
 from redis.driver_info import DriverInfo
 from redis.exceptions import RedisError
 
 from .._version import get_package_version
-from .cluster import Cluster_Conn, Is_Cluster
+from .cluster import Cluster_Conn, Is_Cluster, parse_cluster_slots
 from .graph import AsyncGraph
+from .sentinel import Is_Sentinel, Sentinel_Conn
 
 # config command
 UDF_CMD = "GRAPH.UDF"
@@ -115,6 +117,17 @@ def __init__(
             protocol=protocol,
         )
 
+        self._raw_conn = conn
+        self._ssl = ssl
+        self._replica_connection = 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,64 @@ 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):
+                if getattr(self.connection, "read_from_replicas", False):
+                    return self.connection
+                if self._replica_connection is None:
+                    self._replica_connection = Cluster_Conn(
+                        self._raw_conn, ssl=self._ssl, read_from_replicas=True
+                    )
+                return self._replica_connection
+            return Cluster_Conn(self.connection, ssl=False, read_from_replicas=True)
+
+        return self.connection
+
+    async def get_cluster_shards(self) -> List[Dict[str, Any]]:
+        """
+        Deduces and returns FalkorDB Cluster shards asynchronously,
+        mapping primary nodes to their replicas.
+        """
+        if Is_Cluster(self.connection):
+            raw_slots = await self.connection.execute_command("CLUSTER", "SLOTS")
+            return parse_cluster_slots(raw_slots)
+
+        kwargs = self.connection.connection_pool.connection_kwargs
+        host = kwargs.get("host", "localhost")
+        port = kwargs.get("port", 6379)
+        return [
+            {
+                "primary": {
+                    "id": "standalone",
+                    "host": host,
+                    "port": port,
+                    "endpoint": f"{host}:{port}",
+                },
+                "replicas": [],
+                "slots": [(0, 16383)],
+            }
+        ]
+
+    async def _disconnect_connection(self):
+        """
+        Disconnects the underlying connection or pool to clear dirty socket state.
+        """
+        try:
+            await self.connection.aclose()
+        except (RedisError, OSError):
+            pass
+
     @classmethod
     def from_url(cls, url: str, **kwargs) -> "FalkorDB":
         """
@@ -230,8 +301,8 @@ async def aclose(self) -> None:
 
         try:
             await self.connection.aclose()
-        except RedisError:
-            # best-effort close — don't raise on Redis errors
+        except (RedisError, OSError):
+            # best-effort close — don't raise on Redis or socket errors
             pass
 
     async def __aenter__(self) -> "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,17 @@ 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]