Refactor app.main for clearer database setup

back/README.md

@@ -399,15 +399,15 @@ to simulate a god user with ID 8 (for a regular user, set something like `id=1,
 > [!IMPORTANT]
 > To keep the front-end side up-to-date with the GraphQL schema, make sure that the pre-commit command for `*.graphql` files (`id: generate-graphql-ts-types`) is running properly.
 >
-> It should generate both `schema.graphql` (the introspected unified schema) and `graphql-env.d.ts` (the generated types to be ìnferred and consumed in the FE with `gql.tada`) inside `/graphql/generated/`.
+> It should generate both `schema.graphql` (the introspected unified schema) and `graphql-env.d.ts` (the generated types to be inferred and consumed in the FE with `gql.tada`) inside `/graphql/generated/`.
 
 ## Project structure
 
 The back-end codebase is organized as a Python package called `boxtribute_server`. On the top-most level the most relevant modules are
 
 - `main.py` and `api_main.py`: entry-points to start the Flask app
 - `app.py` and `blueprints.py`: Definition and configuration of Flask app
-- `db.py`: Definition of MySQL interface
+- `db.py`: Utility functions to e.g. create MySQL interface and work with database replica, definition of DatabaseManager class
 - `routes.py`: Definition of web endpoints; invocation of ariadne GraphQL server
 - `auth.py` and `authz.py`: Authentication and authorization utilities
 - `exceptions.py` and `errors.py`: Utility classes for error handling

back/boxtribute_server/app.py

@@ -9,35 +9,17 @@
 from sentry_sdk.integrations.flask import FlaskIntegration
 
 from .db import create_db_interface, db
-from .models.definitions import Model
+from .models import MODELS
 
 
 def create_app():
     return Flask(__name__, static_folder=None)
 
 
-def configure_app(
-    app, *blueprints, database_interface=None, replica_socket=None, **mysql_kwargs
-):
-    """Register blueprints. Configure the app's database interface. `mysql_kwargs` are
-    forwarded. Make data models operate on the primary database.
-    """
+def register_blueprints(app, *blueprints):
     for blueprint in blueprints:
         app.register_blueprint(blueprint)
 
-    app.config["DATABASE"] = database_interface or create_db_interface(**mysql_kwargs)
-
-    if replica_socket or mysql_kwargs:
-        # In deployed environment: replica_socket is set
-        # In integration tests: connect to same host/port as primary database
-        # In endpoint tests, no replica connection is used
-        mysql_kwargs["unix_socket"] = replica_socket
-        app.config["DATABASE_REPLICA"] = create_db_interface(**mysql_kwargs)
-
-    db.init_app(app)
-    # With a complete list of models no need to recursively bind dependencies
-    db.database.bind(Model.__subclasses__(), bind_refs=False, bind_backrefs=False)
-
 
 def main(*blueprints):
     """Integrate Sentry SDK. Create and configure Flask app."""
@@ -79,18 +61,34 @@ def before_sentry_send(event, hint):  # pragma: no cover
     )
 
     app = create_app()
-    configure_app(
-        app,
-        *blueprints,
+    register_blueprints(app, *blueprints)
+
+    # Establish DB connection(s) and initialize DatabaseManager
+    db_connection_parameters = dict(
         # always used
         user=os.environ["MYSQL_USER"],
         password=os.environ["MYSQL_PASSWORD"],
         database=os.environ["MYSQL_DB"],
         # used for connecting to development / CI testing DB
         host=os.getenv("MYSQL_HOST"),
         port=int(os.getenv("MYSQL_PORT", 0)),
-        # used for connecting to Google Cloud from GAE
-        unix_socket=os.getenv("MYSQL_SOCKET"),
-        replica_socket=os.getenv("MYSQL_REPLICA_SOCKET"),
     )
+    # used for connecting to Google Cloud from GAE
+    database_socket = os.getenv("MYSQL_SOCKET")
+    replica_socket = os.getenv("MYSQL_REPLICA_SOCKET")
+
+    db.database = create_db_interface(
+        unix_socket=database_socket, **db_connection_parameters
+    )
+    # In deployed environment: replica_socket is set
+    # In integration tests: connect to same host/port as primary database
+    # In endpoint tests, no replica connection is used
+    db.replica = create_db_interface(
+        unix_socket=replica_socket, **db_connection_parameters
+    )
+
+    # Enable opening/closing DB connection before/after request
+    db.register_handlers(app)
+    # With a complete list of models no need to recursively bind dependencies
+    db.database.bind(MODELS, bind_refs=False, bind_backrefs=False)
     return app

back/boxtribute_server/business_logic/statistics/crud.py

@@ -7,7 +7,7 @@
 
 from peewee import JOIN, SQL, fn
 
-from ...db import db
+from ...db import execute_sql
 from ...enums import BoxState, TaggableObjectType, TargetType
 from ...errors import InvalidDate
 from ...models.definitions.base import Base
@@ -24,7 +24,7 @@
 from ...models.definitions.tag import Tag
 from ...models.definitions.tags_relation import TagsRelation
 from ...models.definitions.transaction import Transaction
-from ...models.utils import compute_age, convert_ids, execute_sql, utcnow
+from ...models.utils import compute_age, convert_ids, utcnow
 from ...utils import in_ci_environment, in_production_environment
 from ..metrics.crud import exclude_test_organisation
 from .sql import MOVED_BOXES_QUERY, STOCK_OVERVIEW_QUERY
@@ -412,7 +412,6 @@ def compute_moved_boxes(*base_ids):
         base_ids,
         TargetType.BoxState.name,
         base_ids,
-        database=db.replica or db.database,
         query=MOVED_BOXES_QUERY,
     )
     for fact in facts:
@@ -461,7 +460,6 @@ def compute_stock_overview(base_id, *, tag_ids=None, excluded_tag_ids=None):
         excluded_tag_ids,
         include_filter_active,
         exclude_filter_active,
-        database=db.replica or db.database,
         query=STOCK_OVERVIEW_QUERY,
     )
     for fact in facts:

back/boxtribute_server/business_logic/warehouse/box/mutations.py

@@ -11,6 +11,7 @@
     authorized_bases_filter,
     handle_unauthorized,
 )
+from ....db import execute_sql
 from ....enums import TaggableObjectType, TagType
 from ....errors import (
     DeletedLocation,
@@ -24,7 +25,6 @@
 from ....models.definitions.product import Product
 from ....models.definitions.tag import Tag
 from ....models.definitions.tags_relation import TagsRelation
-from ....models.utils import execute_sql
 from .crud import (
     WAREHOUSE_BOX_STATES,
     BoxesResult,

back/boxtribute_server/cli/main.py

@@ -18,7 +18,7 @@
 import logging
 
 from ..db import create_db_interface
-from ..models.definitions import Model
+from ..models import MODELS
 from ..models.definitions.base import Base
 from ..models.definitions.organisation import Organisation
 from .remove_base_access import LOGGER as RBA_LOGGER
@@ -146,7 +146,7 @@ def main(args=None):
     database = _create_db_interface(
         **{n: options.pop(n) for n in ["host", "port", "password", "database", "user"]}
     )
-    database.bind(Model.__subclasses__())
+    database.bind(MODELS, False, False)
 
     command = options.pop("command")
     try:

back/boxtribute_server/cron/data_faking.py

@@ -45,7 +45,6 @@
     enable_standard_product,
 )
 from ..business_logic.warehouse.qr_code.crud import create_qr_code
-from ..db import db
 from ..enums import (
     BoxState,
     HumanGender,
@@ -235,7 +234,7 @@ def _fetch_bases_and_users(self):
 
         org_ids = {b.id: b.organisation_id for b in bases}
 
-        cursor = db.database.execute_sql(
+        cursor = Base._meta.database.execute_sql(
             """\
     SELECT cuc.camp_id, group_concat(u.id ORDER BY u.id) FROM cms_users u
     INNER JOIN cms_usergroups_camps cuc

back/boxtribute_server/cron/housekeeping.py

@@ -1,14 +1,13 @@
 from peewee import fn
 
-from ..db import db
 from ..models.definitions.user import User
 
 
 def clean_up_user_email_addresses():
     """Remove excessive '.deleted.X' suffix from user email addresses. Return number of
     updated rows.
     """
-    with db.database.atomic():
+    with User._meta.database.atomic():
         return (
             User.update(email=fn.SUBSTRING_INDEX(User.email, ".deleted.", 2))
             .where(User.email.iregexp(r"\.deleted\.\d+\.deleted\.\d+$"))

back/boxtribute_server/cron/reseed_db.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 
-from ..db import db
+from ..db import current_database, execute_sql
 from ..utils import in_demo_environment, in_staging_environment
 
 
@@ -15,7 +15,7 @@ def reseed_db():
     # For testing locally, run
     # dotenv run flask --debug --app boxtribute_server.dev_main:app run -p 5005
     # curl 'http://localhost:5005/cron/reseed-db' -H 'x-appengine-cron: true'
-    with db.database.cursor() as cursor, open(seed_filepath) as seed:
+    with current_database().cursor() as cursor, open(seed_filepath) as seed:
         execute_sql_statements_from_file(cursor, seed)
 
     if in_staging_environment() or in_demo_environment():
@@ -157,8 +157,8 @@ def update_auth0_role_ids():
         {when_then_statements}
     END
 ; """
-    db.database.execute_sql(
-        command,
+    execute_sql(
         # convert mapping of role names and IDs into flat list
-        [item for role_info in role_ids.items() for item in role_info],
+        *[item for role_info in role_ids.items() for item in role_info],
+        query=command,
     )

back/boxtribute_server/db.py

@@ -2,7 +2,6 @@
 
 from flask import request
 from peewee import MySQLDatabase
-from playhouse.flask_utils import FlaskDB  # type: ignore
 
 from .blueprints import (
     API_GRAPHQL_PATH,
@@ -14,24 +13,29 @@
     shared_bp,
 )
 from .business_logic.statistics import statistics_queries
-from .models.definitions import Model
+from .models import MODELS
 
 
-class DatabaseManager(FlaskDB):
-    """Custom class to glue Flask and Peewee together.
-    If configured accordingly, connect to a database replica for statistics-related
-    GraphQL queries. To use the replica for database queries, wrap the calling code in
-    the `use_db_replica` decorator, and make sure the replica connection is set up in
-    the connect_db() method.
+class DatabaseManager:
+    """Custom class to glue Flask and Peewee together, borrowed from peewee's
+    playhouse.flask_utils.FlaskDB, with irrelevant parts stripped.
+    It holds the references to the primary and the replica database.
+
+    Most importantly, this class handles opening/closing database connection(s)
+    before/after handling incoming requests.
+
+    It connects to the database replica for statistics-related GraphQL queries. To use
+    the replica for database queries, wrap the calling code in the `use_db_replica`
+    decorator.
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-        self.replica = None
+    def __init__(self) -> None:
+        self.database: MySQLDatabase | None = None
+        self.replica: MySQLDatabase | None = None
 
-    def init_app(self, app):
-        self.replica = app.config.get("DATABASE_REPLICA")  # expecting peewee.Database
-        super().init_app(app)
+    def register_handlers(self, app):
+        app.before_request(self.connect_db)
+        app.teardown_request(self.close_db)
 
     def connect_db(self):
         # GraphQL queries are sent as POST requests. Don't open database connection on
@@ -52,6 +56,9 @@ def connect_db(self):
         ):
             return
 
+        if not self.database:
+            raise RuntimeError("DatabaseManager.database not set")
+
         self.database.connect()
 
         # Provide fallback for non-JSON and non-GraphQL requests
@@ -62,8 +69,8 @@ def connect_db(self):
         ):
             self.replica.connect()
 
-    def close_db(self, exc):
-        if not self.database.is_closed():
+    def close_db(self, _):
+        if self.database and not self.database.is_closed():
             self.database.close()
 
         if self.replica and not self.replica.is_closed():
@@ -80,9 +87,7 @@ def use_db_replica(f):
     def decorated(*args, **kwargs):
         if db.replica is not None:
             # With a complete list of models no need to recursively bind dependencies
-            with db.replica.bind_ctx(
-                Model.__subclasses__(), bind_refs=False, bind_backrefs=False
-            ):
+            with db.replica.bind_ctx(MODELS, bind_refs=False, bind_backrefs=False):
                 return f(*args, **kwargs)
 
         return f(*args, **kwargs)
@@ -104,3 +109,29 @@ def create_db_interface(**mysql_kwargs):
     return MySQLDatabase(
         **mysql_kwargs, field_types={"AUTO": "INTEGER UNSIGNED AUTO_INCREMENT"}
     )
+
+
+def current_database() -> MySQLDatabase:
+    """Return the database object that the data models currently are bound to.
+    Return None if run without prior Database.bind() or Database.bind_ctx() call.
+    """
+    database = MODELS[0]._meta.database
+    if database is None:
+        raise RuntimeError("Data models not bound to database.")
+    return database
+
+
+def execute_sql(*params, query):
+    """Utility function to execute a raw SQL query, returning the result rows as dicts.
+    Use the database that the data models currently are bound to. If execute_sql() is
+    called wrapped in use_db_replica(), the replica database is used. Any `params` are
+    passed into peewee's `execute_sql` method as values for query parameters.
+    """
+    database = MODELS[0]._meta.database
+    cursor = database.execute_sql(query, params=params)
+    if cursor.description is None:
+        # For e.g. UPDATE statements no description is available
+        return
+    # Turn cursor result into dict (https://stackoverflow.com/a/56219996/3865876)
+    column_names = [x[0] for x in cursor.description]
+    return [dict(zip(column_names, row)) for row in cursor.fetchall()]

back/boxtribute_server/graph_ql/loaders.py

@@ -6,7 +6,6 @@
 from peewee import SQL, Case, NodeList, fn
 
 from ..authz import authorize, authorized_bases_filter
-from ..db import db
 from ..enums import BoxState as BoxStateEnum
 from ..enums import TaggableObjectType
 from ..exceptions import Forbidden
@@ -264,7 +263,7 @@ async def batch_load_fn(self, box_ids):
 
         # Increase the default of 1024 (would be exceeded for concat'ing the change_date
         # column of a box with 54 or more history entries).
-        db.database.execute_sql("SET SESSION group_concat_max_len = 10000;")
+        History._meta.database.execute_sql("SET SESSION group_concat_max_len = 10000;")
         # Return formatted history entries of boxes with given IDs, sorted by most
         # recent first.
         # Group history entry IDs, change dates, user IDs, and formatted messages for