Move async SQLite mkdir off the event loop (#1921)
Co-authored-by: DanielWalnut <45447813+hetaoBackend@users.noreply.github.com>
This commit is contained in:
parent
3acdf79beb
commit
c4da0e8ca9
|
|
@ -17,6 +17,7 @@ For sync usage see :mod:`deerflow.agents.checkpointer.provider`.
|
||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
import contextlib
|
import contextlib
|
||||||
import logging
|
import logging
|
||||||
from collections.abc import AsyncIterator
|
from collections.abc import AsyncIterator
|
||||||
|
|
@ -54,7 +55,7 @@ async def _async_checkpointer(config) -> AsyncIterator[Checkpointer]:
|
||||||
raise ImportError(SQLITE_INSTALL) from exc
|
raise ImportError(SQLITE_INSTALL) from exc
|
||||||
|
|
||||||
conn_str = resolve_sqlite_conn_str(config.connection_string or "store.db")
|
conn_str = resolve_sqlite_conn_str(config.connection_string or "store.db")
|
||||||
ensure_sqlite_parent_dir(conn_str)
|
await asyncio.to_thread(ensure_sqlite_parent_dir, conn_str)
|
||||||
async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
|
async with AsyncSqliteSaver.from_conn_string(conn_str) as saver:
|
||||||
await saver.setup()
|
await saver.setup()
|
||||||
yield saver
|
yield saver
|
||||||
|
|
|
||||||
|
|
@ -1,7 +1,7 @@
|
||||||
"""Unit tests for checkpointer config and singleton factory."""
|
"""Unit tests for checkpointer config and singleton factory."""
|
||||||
|
|
||||||
import sys
|
import sys
|
||||||
from unittest.mock import MagicMock, patch
|
from unittest.mock import AsyncMock, MagicMock, patch
|
||||||
|
|
||||||
import pytest
|
import pytest
|
||||||
|
|
||||||
|
|
@ -174,6 +174,46 @@ class TestGetCheckpointer:
|
||||||
mock_saver_instance.setup.assert_called_once()
|
mock_saver_instance.setup.assert_called_once()
|
||||||
|
|
||||||
|
|
||||||
|
class TestAsyncCheckpointer:
|
||||||
|
@pytest.mark.anyio
|
||||||
|
async def test_sqlite_creates_parent_dir_via_to_thread(self):
|
||||||
|
"""Async SQLite setup should move mkdir off the event loop."""
|
||||||
|
from deerflow.agents.checkpointer.async_provider import make_checkpointer
|
||||||
|
|
||||||
|
mock_config = MagicMock()
|
||||||
|
mock_config.checkpointer = CheckpointerConfig(type="sqlite", connection_string="relative/test.db")
|
||||||
|
|
||||||
|
mock_saver = AsyncMock()
|
||||||
|
mock_cm = AsyncMock()
|
||||||
|
mock_cm.__aenter__.return_value = mock_saver
|
||||||
|
mock_cm.__aexit__.return_value = False
|
||||||
|
|
||||||
|
mock_saver_cls = MagicMock()
|
||||||
|
mock_saver_cls.from_conn_string.return_value = mock_cm
|
||||||
|
|
||||||
|
mock_module = MagicMock()
|
||||||
|
mock_module.AsyncSqliteSaver = mock_saver_cls
|
||||||
|
|
||||||
|
with (
|
||||||
|
patch("deerflow.agents.checkpointer.async_provider.get_app_config", return_value=mock_config),
|
||||||
|
patch.dict(sys.modules, {"langgraph.checkpoint.sqlite.aio": mock_module}),
|
||||||
|
patch("deerflow.agents.checkpointer.async_provider.asyncio.to_thread", new_callable=AsyncMock) as mock_to_thread,
|
||||||
|
patch(
|
||||||
|
"deerflow.agents.checkpointer.async_provider.resolve_sqlite_conn_str",
|
||||||
|
return_value="/tmp/resolved/test.db",
|
||||||
|
),
|
||||||
|
):
|
||||||
|
async with make_checkpointer() as saver:
|
||||||
|
assert saver is mock_saver
|
||||||
|
|
||||||
|
mock_to_thread.assert_awaited_once()
|
||||||
|
called_fn, called_path = mock_to_thread.await_args.args
|
||||||
|
assert called_fn.__name__ == "ensure_sqlite_parent_dir"
|
||||||
|
assert called_path == "/tmp/resolved/test.db"
|
||||||
|
mock_saver_cls.from_conn_string.assert_called_once_with("/tmp/resolved/test.db")
|
||||||
|
mock_saver.setup.assert_awaited_once()
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
# app_config.py integration
|
# app_config.py integration
|
||||||
# ---------------------------------------------------------------------------
|
# ---------------------------------------------------------------------------
|
||||||
|
|
|
||||||
Loading…
Reference in New Issue