refactor: separate read / write sessions.

edwardzjl · edwardzjl · commit 460bb46ab46e · 2024-10-11T18:14:19.000+08:00
diff --git a/api/chatbot/config.py b/api/chatbot/config.py
@@ -29,12 +29,22 @@ class Settings(BaseSettings):
     model_config = SettingsConfigDict(env_nested_delimiter="__")
 
     llm: dict[str, Any] = Field(default_factory=lambda: {"api_key": "NOT_SET"})
-    db_url: PostgresDsn = "postgresql+psycopg://postgres:postgres@localhost:5432/"
-    """Database url. Must be a valid postgresql connection string."""
+    postgres_primary_url: PostgresDsn = (
+        "postgresql+psycopg://postgres:postgres@localhost:5432/"
+    )
+    """Primary database url. Read/Write. Must be a valid postgresql connection string."""
+    postgres_standby_url: PostgresDsn = postgres_primary_url
+    """Standby database url. Read Only. If present must be a valid postgresql connection string.
+    Defaults to `postgres_primary_url`.
+    """
+
+    @property
+    def psycopg_primary_url(self) -> str:
+        return remove_postgresql_variants(str(self.postgres_primary_url))
 
     @property
-    def psycopg_url(self) -> str:
-        return remove_postgresql_variants(str(self.db_url))
+    def psycopg_standby_url(self) -> str:
+        return remove_postgresql_variants(str(self.postgres_standby_url))
 
 
 settings = Settings()
diff --git a/api/chatbot/dependencies.py b/api/chatbot/dependencies.py
@@ -9,7 +9,7 @@
 
 from chatbot.agent import create_agent
 from chatbot.config import settings
-from chatbot.state import chat_model, sqlalchemy_session
+from chatbot.state import chat_model, sqlalchemy_session, sqlalchemy_ro_session
 
 
 def UserIdHeader(alias: str | None = None, **kwargs):
@@ -47,9 +47,17 @@ async def get_sqlalchemy_session() -> AsyncGenerator[AsyncSession, None]:
 SqlalchemySessionDep = Annotated[AsyncSession, Depends(get_sqlalchemy_session)]
 
 
+async def get_sqlalchemy_ro_session() -> AsyncGenerator[AsyncSession, None]:
+    async with sqlalchemy_ro_session() as session:
+        yield session
+
+
+SqlalchemyROSessionDep = Annotated[AsyncSession, Depends(get_sqlalchemy_ro_session)]
+
+
 async def get_agent() -> AsyncGenerator[CompiledGraph, None]:
     async with AsyncPostgresSaver.from_conn_string(
-        settings.psycopg_url
+        settings.psycopg_primary_url
     ) as checkpointer:
         yield create_agent(chat_model, checkpointer)
 
diff --git a/api/chatbot/main.py b/api/chatbot/main.py
@@ -31,7 +31,7 @@ async def lifespan(app: FastAPI):
 
     # Create checkpointer tables
     async with AsyncPostgresSaver.from_conn_string(
-        settings.psycopg_url
+        settings.psycopg_primary_url
     ) as checkpointer:
         await checkpointer.setup()
 
diff --git a/api/chatbot/routers/conversation.py b/api/chatbot/routers/conversation.py
@@ -3,7 +3,12 @@
 from sqlalchemy import select
 
 from chatbot.chains.summarization import create_smry_chain
-from chatbot.dependencies import AgentStateDep, SqlalchemySessionDep, UserIdHeaderDep
+from chatbot.dependencies import (
+    AgentStateDep,
+    SqlalchemyROSessionDep,
+    SqlalchemySessionDep,
+    UserIdHeaderDep,
+)
 from chatbot.models import Conversation as ORMConversation
 from chatbot.schemas import (
     ChatMessage,
@@ -23,7 +28,7 @@
 @router.get("")
 async def get_conversations(
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
 ) -> list[Conversation]:
     # TODO: support pagination
     stmt = (
@@ -38,7 +43,7 @@ async def get_conversations(
 async def get_conversation(
     conversation_id: str,
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
     agent_state: AgentStateDep,
 ) -> ConversationDetail:
     conv: ORMConversation = await session.get(ORMConversation, conversation_id)
@@ -110,7 +115,7 @@ async def delete_conversation(
 async def summarize(
     conversation_id: str,
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
     agent_state: AgentStateDep,
 ) -> dict[str, str]:
     conv: ORMConversation = await session.get(ORMConversation, conversation_id)
diff --git a/api/chatbot/routers/message.py b/api/chatbot/routers/message.py
@@ -1,7 +1,7 @@
 from fastapi import APIRouter, HTTPException
 from langchain_core.messages import BaseMessage
 
-from chatbot.dependencies import AgentDep, SqlalchemySessionDep, UserIdHeaderDep
+from chatbot.dependencies import AgentDep, SqlalchemyROSessionDep, UserIdHeaderDep
 from chatbot.models import Conversation as ORMConversation
 
 router = APIRouter(
@@ -16,7 +16,7 @@ async def thumbup(
     conversation_id: str,
     message_id: str,
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
     agent: AgentDep,
 ) -> None:
     """Using message index as the uuid is in the message body which is json dumped into redis,
@@ -48,7 +48,7 @@ async def thumbdown(
     conversation_id: str,
     message_id: str,
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
     agent: AgentDep,
 ) -> None:
     """Using message index as the uuid is in the message body which is json dumped into redis,
diff --git a/api/chatbot/routers/share.py b/api/chatbot/routers/share.py
@@ -7,7 +7,12 @@
 from sqlalchemy import select
 from starlette.requests import Request
 
-from chatbot.dependencies import AgentDep, SqlalchemySessionDep, UserIdHeaderDep
+from chatbot.dependencies import (
+    AgentDep,
+    SqlalchemyROSessionDep,
+    SqlalchemySessionDep,
+    UserIdHeaderDep,
+)
 from chatbot.models import Conversation as ORMConv, Share as ORMShare
 from chatbot.schemas import ChatMessage, CreateShare, Share
 
@@ -22,7 +27,7 @@
 @router.get("")
 async def get_shares(
     userid: UserIdHeaderDep,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
 ) -> list[Share]:
     """Get shares by userid"""
     # TODO: support pagination
@@ -37,7 +42,7 @@ async def get_shares(
 @router.get("/{share_id}")
 async def get_share(
     share_id: str,
-    session: SqlalchemySessionDep,
+    session: SqlalchemyROSessionDep,
     agent: AgentDep,
 ) -> Share:
     """Get a share by id"""
diff --git a/api/chatbot/state.py b/api/chatbot/state.py
@@ -9,7 +9,7 @@
 
 
 sqlalchemy_engine = create_async_engine(
-    str(settings.db_url),
+    str(settings.postgres_primary_url),
     poolclass=NullPool,
 )
 sqlalchemy_session = sessionmaker(
@@ -19,4 +19,15 @@
     autoflush=False,
     class_=AsyncSession,
 )
+sqlalchemy_ro_engine = create_async_engine(
+    str(settings.postgres_standby_url),
+    poolclass=NullPool,
+)
+sqlalchemy_ro_session = sessionmaker(
+    sqlalchemy_ro_engine,
+    autocommit=False,
+    expire_on_commit=False,
+    autoflush=False,
+    class_=AsyncSession,
+)
 chat_model = ChatOpenAI(**settings.llm)
diff --git a/api/tests/test_config.py b/api/tests/test_config.py
@@ -37,20 +37,63 @@ def test_llm_custom(self):
         settings = Settings(llm=custom_llm)
         self.assertEqual(settings.llm, custom_llm)
 
-    def test_psycopg_url_default(self):
+    def test_postgres_primary_url_default(self):
         settings = Settings()
-        expected = "postgresql://postgres:postgres@localhost:5432/"
-        self.assertEqual(settings.psycopg_url, expected)
+        expected = "postgresql+psycopg://postgres:postgres@localhost:5432/"
+        self.assertEqual(str(settings.postgres_primary_url), expected)
+
+    def test_postgres_primary_url_custom(self):
+        custom_primary_url = (
+            "postgresql+psycopg://primary_user:primary_pass@localhost/primary_db"
+        )
+        settings = Settings(postgres_primary_url=custom_primary_url)
+        expected = "postgresql+psycopg://primary_user:primary_pass@localhost/primary_db"
+        self.assertEqual(str(settings.postgres_primary_url), expected)
+
+    def test_postgres_standby_url_default(self):
+        settings = Settings()
+        expected = "postgresql+psycopg://postgres:postgres@localhost:5432/"
+        self.assertEqual(str(settings.postgres_standby_url), expected)
 
-    def test_psycopg_url_custom(self):
-        custom_url = "postgresql+psycopg2://custom_user:custom_pass@localhost/custom_db"
-        settings = Settings(db_url=custom_url)
-        expected = "postgresql://custom_user:custom_pass@localhost/custom_db"
-        self.assertEqual(settings.psycopg_url, expected)
+    def test_postgres_standby_url_custom(self):
+        custom_standby_url = (
+            "postgresql+psycopg://standby_user:standby_pass@localhost/standby_db"
+        )
+        settings = Settings(postgres_standby_url=custom_standby_url)
+        expected = "postgresql+psycopg://standby_user:standby_pass@localhost/standby_db"
+        self.assertEqual(str(settings.postgres_standby_url), expected)
 
     def test_invalid_db_url(self):
         with self.assertRaises(ValidationError):
-            Settings(db_url="invalid_url")
+            Settings(postgres_primary_url="invalid_url")
+        with self.assertRaises(ValidationError):
+            Settings(postgres_standby_url="invalid_url")
+
+    def test_psycopg_primary_url_default(self):
+        settings = Settings()
+        expected = "postgresql://postgres:postgres@localhost:5432/"
+        self.assertEqual(settings.psycopg_primary_url, expected)
+
+    def test_psycopg_primary_url_custom(self):
+        custom_primary_url = (
+            "postgresql+psycopg://primary_user:primary_pass@localhost/primary_db"
+        )
+        settings = Settings(postgres_primary_url=custom_primary_url)
+        expected = "postgresql://primary_user:primary_pass@localhost/primary_db"
+        self.assertEqual(settings.psycopg_primary_url, expected)
+
+    def test_psycopg_standby_url_default(self):
+        settings = Settings()
+        expected = "postgresql://postgres:postgres@localhost:5432/"
+        self.assertEqual(settings.psycopg_standby_url, expected)
+
+    def test_psycopg_standby_url_custom(self):
+        custom_standby_url = (
+            "postgresql+psycopg://standby_user:standby_pass@localhost/standby_db"
+        )
+        settings = Settings(postgres_standby_url=custom_standby_url)
+        expected = "postgresql://standby_user:standby_pass@localhost/standby_db"
+        self.assertEqual(settings.psycopg_standby_url, expected)
 
 
 if __name__ == "__main__":
