add close and aclose on FalkorDB connections

falkordb/asyncio/falkordb.py

@@ -1,4 +1,5 @@
 import redis.asyncio as redis
+from redis.exceptions import RedisError
 from .cluster import *
 from .graph import AsyncGraph
 from typing import List, Union
@@ -197,3 +198,25 @@ async def config_set(self, name: str, value=None) -> None:
         """
 
         return await self.connection.execute_command(CONFIG_CMD, "SET", name, value)
+
+    async def aclose(self) -> None:
+        """
+        Close the underlying connection(s).
+        """
+
+        try:
+            await self.connection.aclose()
+        except RedisError:
+            # best-effort close — don't raise on Redis errors
+            pass
+
+
+    async def __aenter__(self) -> "FalkorDB":
+        """Return self to support async with-statement usage."""
+
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Close the connection when exiting an async with-statement."""
+
+        await self.aclose()

falkordb/falkordb.py

@@ -1,4 +1,5 @@
 import redis
+from redis.exceptions import RedisError
 from .cluster import *
 from .sentinel import *
 from .graph import Graph
@@ -237,3 +238,23 @@ def config_set(self, name: str, value=None) -> None:
         """
 
         return self.connection.execute_command(CONFIG_CMD, "SET", name, value)
+
+    def close(self) -> None:
+        """
+        Close the underlying connection(s).
+        """
+
+        try:
+            self.connection.close()
+        except RedisError:
+            # best-effort close — don't raise on Redis errors
+            pass
+
+    def __enter__(self) -> "FalkorDB":
+        """Return self to support usage in a with-statement."""
+
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Close the connection when exiting a with-statement."""
+        self.close()

tests/test_close.py

@@ -0,0 +1,55 @@
+from types import SimpleNamespace
+from unittest.mock import Mock, AsyncMock
+
+import pytest
+
+from falkordb.falkordb import FalkorDB as SyncFalkorDB
+from falkordb.asyncio.falkordb import FalkorDB as AsyncFalkorDB
+
+
+def test_sync_context_manager_calls_close():
+    db = object.__new__(SyncFalkorDB)
+
+    mock_conn = SimpleNamespace(close=Mock())
+    mock_conn.connection_pool = SimpleNamespace(disconnect=Mock())
+    db.connection = mock_conn
+
+    # using context manager should call close on the underlying connection
+    with db as d:
+        assert d is db
+
+    mock_conn.close.assert_called_once()
+
+
+def test_sync_close_fallback_disconnect():
+    db = object.__new__(SyncFalkorDB)
+
+    # connection provides close(); close() should be preferred over disconnect
+    mock_conn = SimpleNamespace(close=Mock())
+    mock_conn.connection_pool = SimpleNamespace(disconnect=Mock())
+    db.connection = mock_conn
+
+    db.close()
+
+    mock_conn.close.assert_called_once()
+    # ensure fallback wasn't used
+    mock_conn.connection_pool.disconnect.assert_not_called()
+
+
+@pytest.mark.asyncio
+async def test_async_aclose_and_context_manager():
+    db = object.__new__(AsyncFalkorDB)
+
+    mock_conn = SimpleNamespace(aclose=AsyncMock())
+    db.connection = mock_conn
+
+    # explicit aclose
+    await db.aclose()
+    mock_conn.aclose.assert_awaited_once()
+
+    # async context manager should also await aclose
+    mock_conn.aclose.reset_mock()
+    async with db as d:
+        assert d is db
+
+    mock_conn.aclose.assert_awaited_once()