gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: adds connection locks

Browse source
This commit is contained in:
hajdul88 2025-10-07 19:19:57 +02:00
parent b288e5d7a0
commit 68a8977301
1 changed files with 6 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
cognee/infrastructure/databases/graph/kuzu/adapter.py
View file

@ -55,7 +55,7 @@ class KuzuAdapter(GraphDBInterface):
else: else:
self._initialize_connection() self._initialize_connection()
self.KUZU_ASYNC_LOCK = asyncio.Lock() self.KUZU_ASYNC_LOCK = asyncio.Lock()
self._counter_lock = asyncio.Lock() self._connection_change_lock = asyncio.Lock()
def _initialize_connection(self) -> None: def _initialize_connection(self) -> None:
"""Initialize the Kuzu database connection and schema.""" """Initialize the Kuzu database connection and schema."""
@ -224,17 +224,17 @@ class KuzuAdapter(GraphDBInterface):
raise raise
if cache_config.caching: if cache_config.caching:
if self._is_closed: async with self._connection_change_lock:
self.reopen()
async with self._counter_lock:
self.open_connections += 1 self.open_connections += 1
if self._is_closed:
self.reopen()
logger.info(f"Open connections after open: {self.open_connections}") logger.info(f"Open connections after open: {self.open_connections}")
result = await loop.run_in_executor(self.executor, blocking_query) result = await loop.run_in_executor(self.executor, blocking_query)
if cache_config.caching: if cache_config.caching:
async with self._counter_lock: async with self._connection_change_lock:
self.open_connections -= 1 self.open_connections -= 1
logger.info(f"Opened connections after closing {self.open_connections}") logger.info(f"Opened connections after closing {self.open_connections}")
if self.open_connections == 0: if self.open_connections == 0: