Merge pull request #42 from bugout-dev/ro-db-connection

Using ro db session for middlewares

Author: kompotkot

brood/api.py

@@ -35,7 +35,7 @@
     is_token_restricted_or_installation,
     get_current_user_or_installation,
 )
-from .external import yield_db_session_from_env
+from .db import yield_db_session_from_env
 from .version import BROOD_VERSION
 from .settings import (
     group_invite_link_from_env,

brood/cli.py

@@ -11,7 +11,7 @@
 from . import data
 from . import exceptions
 from . import subscriptions
-from .external import SessionLocal
+from .db import SessionLocal
 from .models import (
     User,
     Group,

brood/db.py

@@ -0,0 +1,69 @@
+"""
+Connections to external services
+"""
+from typing import Optional
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm.session import Session, sessionmaker
+
+from .settings import (
+    BROOD_DB_URI,
+    BROOD_DB_URI_READ_ONLY,
+    BROOD_POOL_SIZE,
+    BROOD_DB_STATEMENT_TIMEOUT_MILLIS,
+)
+
+
+def create_brood_engine(url: Optional[str], pool_size: int, statement_timeout: int):
+    # Pooling: https://docs.sqlalchemy.org/en/14/core/pooling.html#sqlalchemy.pool.QueuePool
+    # Statement timeout: https://stackoverflow.com/a/44936982
+    return create_engine(
+        url=url,
+        pool_size=pool_size,
+        connect_args={"options": f"-c statement_timeout={statement_timeout}"},
+    )
+
+
+engine = create_brood_engine(
+    url=BROOD_DB_URI,
+    pool_size=BROOD_POOL_SIZE,
+    statement_timeout=BROOD_DB_STATEMENT_TIMEOUT_MILLIS,
+)
+SessionLocal = sessionmaker(bind=engine)
+
+
+def yield_db_session_from_env() -> Session:
+    """
+    Creates an active database session using configuration from the environment and yields it as
+    per FastAPI docs:
+    https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
+
+    Behaves identically to db_session_from_env in all other respects.
+    """
+    session = SessionLocal()
+    try:
+        yield session
+    finally:
+        session.close()
+
+
+# Read only
+RO_engine = create_brood_engine(
+    url=BROOD_DB_URI_READ_ONLY,
+    pool_size=BROOD_POOL_SIZE,
+    statement_timeout=BROOD_DB_STATEMENT_TIMEOUT_MILLIS,
+)
+RO_SessionLocal = sessionmaker(bind=RO_engine)
+
+
+def yield_db_read_only_session() -> Session:
+    """
+    Yields read only database connection (created using environment variables).
+    As per FastAPI docs:
+    https://fastapi.tiangolo.com/tutorial/sql-databases/#create-a-dependency
+    """
+    session = RO_SessionLocal()
+    try:
+        yield session
+    finally:
+        session.close()

brood/external.py

@@ -1,3 +1,4 @@
+import logging
 from typing import Optional, Union
 from uuid import UUID
 
@@ -11,9 +12,11 @@
 from . import actions
 from . import data
 from . import models
-from .external import yield_db_session_from_env
+from .db import yield_db_read_only_session
 from .settings import BOT_INSTALLATION_TOKEN, BOT_INSTALLATION_TOKEN_HEADER
 
+logger = logging.getLogger(__name__)
+
 # Login implementation follows:
 # https://fastapi.tiangolo.com/tutorial/security/simple-oauth2/
 oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
@@ -22,20 +25,23 @@
 
 async def get_current_user(
     token: UUID = Depends(oauth2_scheme),
-    db_session=Depends(yield_db_session_from_env),
+    db_session=Depends(yield_db_read_only_session),
 ) -> models.User:
     try:
         token_object = actions.get_token(session=db_session, token=token)
     except actions.TokenNotFound:
         raise HTTPException(status_code=404, detail="Access token not found")
+    except Exception:
+        logger.error("Unhandled exception at get_current_user")
+        raise HTTPException(status_code=500)
     if not token_object.active:
         raise HTTPException(status_code=403, detail="Token has expired")
     return token_object.user
 
 
 async def get_current_user_with_groups(
     token: UUID = Depends(oauth2_scheme),
-    db_session=Depends(yield_db_session_from_env),
+    db_session=Depends(yield_db_read_only_session),
 ) -> data.UserWithGroupsResponse:
     try:
         token_active, user_extended = actions.get_current_user_with_groups(
@@ -44,6 +50,7 @@ async def get_current_user_with_groups(
     except actions.TokenNotFound:
         raise HTTPException(status_code=404, detail="Access token not found")
     except Exception:
+        logger.error("Unhandled exception at get_current_user_with_groups")
         raise HTTPException(status_code=500)
     if not token_active:
         raise HTTPException(status_code=403, detail="Token has expired")
@@ -78,7 +85,7 @@ def autogenerated_user_token_check(request: Request) -> bool:
 async def get_current_user_or_installation(
     request: Request,
     token: UUID = Depends(oauth2_scheme_manual),
-    db_session=Depends(yield_db_session_from_env),
+    db_session=Depends(yield_db_read_only_session),
 ) -> Union[models.User, bool]:
     """
     Allow access if Bugout installation token provided, if not
@@ -97,7 +104,7 @@ async def get_current_user_or_installation(
 async def is_token_restricted_or_installation(
     request: Request,
     token: UUID = Depends(oauth2_scheme_manual),
-    db_session=Depends(yield_db_session_from_env),
+    db_session=Depends(yield_db_read_only_session),
 ) -> bool:
     """
     Allow access if Bugout installation provided.
@@ -114,10 +121,13 @@ async def is_token_restricted_or_installation(
 
 async def is_token_restricted(
     token: UUID = Depends(oauth2_scheme),
-    db_session=Depends(yield_db_session_from_env),
+    db_session=Depends(yield_db_read_only_session),
 ) -> bool:
     try:
         token_object = actions.get_token(session=db_session, token=token)
     except actions.TokenNotFound:
         raise HTTPException(status_code=404, detail="Access token not found")
+    except Exception:
+        logger.error("Unhandled exception at is_token_restricted")
+        raise HTTPException(status_code=500)
     return token_object.restricted

brood/middleware.py

@@ -21,7 +21,7 @@
 from .version import BROOD_RESOURCES_VERSION
 from ..data import VersionResponse
 from .. import models as brood_models
-from ..external import yield_db_session_from_env
+from ..db import yield_db_session_from_env
 from ..middleware import get_current_user
 from ..settings import ORIGINS, DOCS_TARGET_PATH, BROOD_OPENAPI_LIST
 

brood/resources/api.py

@@ -5,7 +5,7 @@
 import json
 
 from .models import Resource
-from ..external import SessionLocal
+from ..db import SessionLocal
 
 
 def resources_list_handler(args: argparse.Namespace) -> None:

brood/resources/cli.py

@@ -26,8 +26,34 @@
 )
 MOONSTREAM_APPLICATION_ID = os.environ.get("MOONSTREAM_APPLICATION_ID")
 
-DB_URI = os.environ.get("BROOD_DB_URI")
+# Database
+BROOD_DB_URI = os.environ.get("BROOD_DB_URI")
+if BROOD_DB_URI is None:
+    raise ValueError("BROOD_DB_URI environment variable not set")
+BROOD_DB_URI_READ_ONLY = os.environ.get("BROOD_DB_URI_READ_ONLY")
+if BROOD_DB_URI_READ_ONLY is None:
+    raise ValueError("BROOD_DB_URI_READ_ONLY environment variable not set")
 
+BROOD_POOL_SIZE_RAW = os.environ.get("BROOD_POOL_SIZE", 0)
+try:
+    if BROOD_POOL_SIZE_RAW is not None:
+        BROOD_POOL_SIZE = int(BROOD_POOL_SIZE_RAW)
+except:
+    raise Exception(f"Could not parse BROOD_POOL_SIZE as int: {BROOD_POOL_SIZE_RAW}")
+
+BROOD_DB_STATEMENT_TIMEOUT_MILLIS_RAW = os.environ.get(
+    "BROOD_DB_STATEMENT_TIMEOUT_MILLIS"
+)
+BROOD_DB_STATEMENT_TIMEOUT_MILLIS = 10000
+try:
+    if BROOD_DB_STATEMENT_TIMEOUT_MILLIS_RAW is not None:
+        BROOD_DB_STATEMENT_TIMEOUT_MILLIS = int(BROOD_DB_STATEMENT_TIMEOUT_MILLIS_RAW)
+except:
+    raise ValueError(
+        f"BROOD_DB_STATEMENT_TIMEOUT_MILLIS must be an integer: {BROOD_DB_STATEMENT_TIMEOUT_MILLIS_RAW}"
+    )
+
+# Bots
 BOT_INSTALLATION_TOKEN = os.environ.get("BUGOUT_BOT_INSTALLATION_TOKEN")
 BOT_INSTALLATION_TOKEN_HEADER_RAW = os.environ.get(
     "BUGOUT_BOT_INSTALLATION_TOKEN_HEADER"

brood/settings.py

@@ -1,5 +1,6 @@
 # Required environment variables
 export BROOD_DB_URI="postgresql://<username>:<password>@<db_host>/<db_name>"
+export BROOD_DB_URI_READ_ONLY="postgresql://<username>:<password>@<db_host>/<db_name>"
 export BROOD_CORS_ALLOWED_ORIGINS="http://localhost:3000,https://bugout.dev,https://www.bugout.dev"
 export BUGOUT_WEB_URL="https://bugout.dev"
 export BUGOUT_GROUP_FREE_SEATS=5