From e89da1f440202ea98e4b7606be32f43370213b12 Mon Sep 17 00:00:00 2001 From: kompotkot Date: Thu, 18 Aug 2022 12:25:23 +0000 Subject: [PATCH 1/3] Using ro db session for middlewares --- brood/api.py | 2 +- brood/cli.py | 2 +- brood/db.py | 67 ++++++++++++++++++++++++++++++++++++++++++ brood/external.py | 28 ------------------ brood/middleware.py | 12 ++++---- brood/resources/api.py | 2 +- brood/resources/cli.py | 2 +- brood/settings.py | 28 +++++++++++++++++- configs/sample.env | 1 + 9 files changed, 105 insertions(+), 39 deletions(-) create mode 100644 brood/db.py delete mode 100644 brood/external.py diff --git a/brood/api.py b/brood/api.py index b6c897a..700bea1 100644 --- a/brood/api.py +++ b/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, diff --git a/brood/cli.py b/brood/cli.py index 1feceed..f01068e 100644 --- a/brood/cli.py +++ b/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, diff --git a/brood/db.py b/brood/db.py new file mode 100644 index 0000000..9f54c37 --- /dev/null +++ b/brood/db.py @@ -0,0 +1,67 @@ +""" +Connections to external services +""" +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: 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() diff --git a/brood/external.py b/brood/external.py deleted file mode 100644 index 5e0e839..0000000 --- a/brood/external.py +++ /dev/null @@ -1,28 +0,0 @@ -""" -Connections to external services -""" -from sqlalchemy import create_engine -from sqlalchemy.orm.session import Session, sessionmaker - -from .settings import DB_URI - - -if DB_URI is None: - raise ValueError("BROOD_DB_URI environment variable not set") -engine = create_engine(DB_URI) -SessionLocal = sessionmaker(autocommit=False, autoflush=False, 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() diff --git a/brood/middleware.py b/brood/middleware.py index 434d5cd..39fa00c 100644 --- a/brood/middleware.py +++ b/brood/middleware.py @@ -11,7 +11,7 @@ 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 # Login implementation follows: @@ -22,7 +22,7 @@ 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) @@ -35,7 +35,7 @@ async def get_current_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( @@ -78,7 +78,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 +97,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,7 +114,7 @@ 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) diff --git a/brood/resources/api.py b/brood/resources/api.py index 74ac030..960dd70 100644 --- a/brood/resources/api.py +++ b/brood/resources/api.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 diff --git a/brood/resources/cli.py b/brood/resources/cli.py index c44281c..dcdb15a 100644 --- a/brood/resources/cli.py +++ b/brood/resources/cli.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: diff --git a/brood/settings.py b/brood/settings.py index 35f3928..01b2fe0 100644 --- a/brood/settings.py +++ b/brood/settings.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" diff --git a/configs/sample.env b/configs/sample.env index 1650cf1..2ddd0b9 100644 --- a/configs/sample.env +++ b/configs/sample.env @@ -1,5 +1,6 @@ # Required environment variables export BROOD_DB_URI="postgresql://:@/" +export BROOD_DB_URI_READ_ONLY="postgresql://:@/" 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 From 16f8bbe5b5c11d09569b0719d928af09bfcaf66d Mon Sep 17 00:00:00 2001 From: kompotkot Date: Thu, 18 Aug 2022 12:29:39 +0000 Subject: [PATCH 2/3] Common exceptions --- brood/middleware.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/brood/middleware.py b/brood/middleware.py index 39fa00c..5ea8cfe 100644 --- a/brood/middleware.py +++ b/brood/middleware.py @@ -28,6 +28,8 @@ async def get_current_user( 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: + raise HTTPException(status_code=500) if not token_object.active: raise HTTPException(status_code=403, detail="Token has expired") return token_object.user @@ -120,4 +122,6 @@ async def is_token_restricted( 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: + raise HTTPException(status_code=500) return token_object.restricted From e38b2dc979359ab562e3c5afa1b03c78a8122435 Mon Sep 17 00:00:00 2001 From: kompotkot Date: Thu, 18 Aug 2022 12:34:14 +0000 Subject: [PATCH 3/3] Fix mypy and logger for common exceptions --- brood/db.py | 4 +++- brood/middleware.py | 6 ++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/brood/db.py b/brood/db.py index 9f54c37..ed895e1 100644 --- a/brood/db.py +++ b/brood/db.py @@ -1,6 +1,8 @@ """ Connections to external services """ +from typing import Optional + from sqlalchemy import create_engine from sqlalchemy.orm.session import Session, sessionmaker @@ -12,7 +14,7 @@ ) -def create_brood_engine(url: str, pool_size: int, statement_timeout: int): +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( diff --git a/brood/middleware.py b/brood/middleware.py index 5ea8cfe..d7a40c2 100644 --- a/brood/middleware.py +++ b/brood/middleware.py @@ -1,3 +1,4 @@ +import logging from typing import Optional, Union from uuid import UUID @@ -14,6 +15,8 @@ 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") @@ -29,6 +32,7 @@ async def get_current_user( 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") @@ -46,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") @@ -123,5 +128,6 @@ async def is_token_restricted( 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