diff --git a/falkordb/asyncio/falkordb.py b/falkordb/asyncio/falkordb.py index f2dd997..3e129ec 100644 --- a/falkordb/asyncio/falkordb.py +++ b/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() diff --git a/falkordb/falkordb.py b/falkordb/falkordb.py index 5bc5f58..58f0d54 100644 --- a/falkordb/falkordb.py +++ b/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() diff --git a/tests/test_close.py b/tests/test_close.py new file mode 100644 index 0000000..99af2a0 --- /dev/null +++ b/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()