Add app.REPLICA_DATABASE (#4289)

skyvern/config.py

@@ -59,6 +59,7 @@ class Settings(BaseSettings):
         if platform.system() == "Windows"
         else "postgresql+psycopg://skyvern@localhost/skyvern"
     )
+    DATABASE_REPLICA_STRING: str | None = None
     DATABASE_STATEMENT_TIMEOUT_MS: int = 60000
     DISABLE_CONNECTION_POOL: bool = False
     PROMPT_ACTION_HISTORY_WINDOW: int = 1

skyvern/forge/forge_app.py

@@ -21,7 +21,7 @@ from skyvern.forge.sdk.artifact.storage.factory import StorageFactory
 from skyvern.forge.sdk.artifact.storage.s3 import S3Storage
 from skyvern.forge.sdk.cache.base import BaseCache
 from skyvern.forge.sdk.cache.factory import CacheFactory
-from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.experimentation.providers import BaseExperimentationProvider, NoOpExperimentationProvider
 from skyvern.forge.sdk.schemas.credentials import CredentialVaultType
 from skyvern.forge.sdk.schemas.organizations import AzureClientSecretCredential, Organization
@@ -44,6 +44,7 @@ class ForgeApp:
 
     SETTINGS_MANAGER: Settings
     DATABASE: AgentDB
+    REPLICA_DATABASE: AgentDB
     STORAGE: BaseStorage
     CACHE: BaseCache
     ARTIFACT_MANAGER: ArtifactManager
@@ -93,6 +94,12 @@ def create_forge_app() -> ForgeApp:
     app.SETTINGS_MANAGER = settings
 
     app.DATABASE = AgentDB(settings.DATABASE_STRING, debug_enabled=settings.DEBUG_MODE)
+
+    if settings.DATABASE_REPLICA_STRING and settings.DATABASE_REPLICA_STRING != settings.DATABASE_STRING:
+        app.REPLICA_DATABASE = AgentDB(settings.DATABASE_REPLICA_STRING, debug_enabled=settings.DEBUG_MODE)
+    else:
+        app.REPLICA_DATABASE = app.DATABASE
+
     if settings.SKYVERN_STORAGE_TYPE == "s3":
         StorageFactory.set_storage(S3Storage())
     app.STORAGE = StorageFactory.get_storage()

skyvern/forge/sdk/db/agent_db_test.py

@@ -4,7 +4,7 @@ import pytest
 import pytest_asyncio
 from sqlalchemy.ext.asyncio import create_async_engine
 
-from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.db.models import Base
 
 

skyvern/forge/sdk/db/polls.py

@@ -2,7 +2,7 @@ import asyncio
 
 from structlog import get_logger
 
-from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import PersistentBrowserSession
 
 LOG = get_logger(__name__)

skyvern/forge/sdk/services/org_auth_service.py

@@ -13,7 +13,7 @@ from pydantic import ValidationError
 from skyvern.config import settings
 from skyvern.forge import app
 from skyvern.forge.sdk.core import skyvern_context
-from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.models import TokenPayload
 from skyvern.forge.sdk.schemas.organizations import Organization, OrganizationAuthToken, OrganizationAuthTokenType
 

skyvern/webeye/persistent_sessions_manager.py

@@ -9,7 +9,7 @@ from playwright._impl._errors import TargetClosedError
 
 from skyvern.config import settings
 from skyvern.exceptions import BrowserSessionNotRenewable, MissingBrowserAddressError
-from skyvern.forge.sdk.db.client import AgentDB
+from skyvern.forge.sdk.db.agent_db import AgentDB
 from skyvern.forge.sdk.db.polls import wait_on_persistent_browser_address
 from skyvern.forge.sdk.schemas.persistent_browser_sessions import (
     PersistentBrowserSession,