feat: support sqlite async for async DB engine (#2384)

This commit is contained in:
Sarah Wooders
2025-05-23 13:49:07 -07:00
committed by GitHub
parent edbddd746d
commit bca6d5f794
3 changed files with 536 additions and 129 deletions

View File

@@ -118,20 +118,21 @@ class DatabaseRegistry:
else:
async_pg_uri = f"postgresql+asyncpg://{pg_uri.split('://', 1)[1]}" if "://" in pg_uri else pg_uri
async_pg_uri = async_pg_uri.replace("sslmode=", "ssl=")
async_engine = create_async_engine(async_pg_uri, **self._build_sqlalchemy_engine_args(is_async=True))
self._async_engines["default"] = async_engine
# Create async session factory
self._async_session_factories["default"] = async_sessionmaker(
autocommit=False, autoflush=False, bind=self._async_engines["default"], class_=AsyncSession
)
self._initialized["async"] = True
else:
self.logger.warning("Async SQLite is currently not supported. Please use PostgreSQL for async database operations.")
# TODO (cliandy): unclear around async sqlite support in sqlalchemy, we will not currently support this
# create sqlite async engine
self._initialized["async"] = False
# TODO: remove self.config
engine_path = "sqlite+aiosqlite:///" + os.path.join(self.config.recall_storage_path, "sqlite.db")
self.logger.info("Creating sqlite engine " + engine_path)
async_engine = create_async_engine(engine_path, **self._build_sqlalchemy_engine_args(is_async=True))
# Create async session factory
self._async_engines["default"] = async_engine
self._async_session_factories["default"] = async_sessionmaker(
autocommit=False, autoflush=False, bind=self._async_engines["default"], class_=AsyncSession
)
self._initialized["async"] = True
def _build_sqlalchemy_engine_args(self, *, is_async: bool) -> dict:
"""Prepare keyword arguments for create_engine / create_async_engine."""

636
poetry.lock generated

File diff suppressed because it is too large Load Diff

View File

@@ -36,7 +36,7 @@ fastapi = { version = "^0.115.6", optional = true}
uvicorn = {version = "^0.24.0.post1", optional = true}
pydantic = "^2.10.6"
html2text = "^2020.1.16"
sqlalchemy = "^2.0.25"
sqlalchemy = {extras = ["asyncio"], version = "^2.0.41"}
pexpect = {version = "^4.9.0", optional = true}
pyright = {version = "^1.1.347", optional = true}
qdrant-client = {version="^1.9.1", optional = true}
@@ -90,12 +90,12 @@ firecrawl-py = "^1.15.0"
apscheduler = "^3.11.0"
aiomultiprocess = "^0.9.1"
matplotlib = "^3.10.1"
asyncpg = "^0.30.0"
asyncpg = {version = "^0.30.0", optional = true}
tavily-python = "^0.7.2"
[tool.poetry.extras]
postgres = ["pgvector", "pg8000", "psycopg2-binary", "psycopg2"]
postgres = ["pgvector", "pg8000", "psycopg2-binary", "psycopg2", "asyncpg"]
dev = ["pytest", "pytest-asyncio", "pexpect", "black", "pre-commit", "pyright", "pytest-order", "autoflake", "isort", "locust"]
server = ["websockets", "fastapi", "uvicorn"]
qdrant = ["qdrant-client"]