vodorod/Dorod-Sky
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Allow AgentDB init to accept a db_engine (for unit tests) (#2626)

Browse Source
This commit is contained in:
Asher Foa
2025-06-09 16:12:23 -04:00
committed by GitHub
parent 4de7746ebc
commit 0503f403bd
1 changed files with 11 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
skyvern/forge/sdk/db/client.py
View File

@@ -5,7 +5,7 @@ from typing import Any, List, Sequence
import structlog import structlog
from sqlalchemy import and_, delete, distinct, func, pool, select, tuple_, update from sqlalchemy import and_, delete, distinct, func, pool, select, tuple_, update
from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.exc import SQLAlchemyError
from sqlalchemy.ext.asyncio import async_sessionmaker, create_async_engine from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker, create_async_engine
from skyvern.config import settings from skyvern.config import settings
from skyvern.exceptions import WorkflowParameterNotFound, WorkflowRunNotFound from skyvern.exceptions import WorkflowParameterNotFound, WorkflowRunNotFound
@@ -106,14 +106,18 @@ elif "postgresql+asyncpg" in settings.DATABASE_STRING:
class AgentDB: class AgentDB:
def __init__(self, database_string: str, debug_enabled: bool = False) -> None: def __init__(self, database_string: str, debug_enabled: bool = False, db_engine: AsyncEngine | None = None) -> None:
super().__init__() super().__init__()
self.debug_enabled = debug_enabled self.debug_enabled = debug_enabled
self.engine = create_async_engine( self.engine = (
database_string, create_async_engine(
json_serializer=_custom_json_serializer, database_string,
connect_args=DB_CONNECT_ARGS, json_serializer=_custom_json_serializer,
poolclass=pool.NullPool if settings.DISABLE_CONNECTION_POOL else None, connect_args=DB_CONNECT_ARGS,
poolclass=pool.NullPool if settings.DISABLE_CONNECTION_POOL else None,
)
if db_engine is None
else db_engine
) )
self.Session = async_sessionmaker(bind=self.engine) self.Session = async_sessionmaker(bind=self.engine)