feat(session): allow configuring session bind behavior

docs/api.rst

@@ -13,6 +13,7 @@ API Docs
 
 .. automodule:: invenio_db.shared
    :members:
+   :exclude-members: Session
 
 .. automodule:: invenio_db.cli
    :members:

docs/conf.py

@@ -9,11 +9,6 @@
 
 """Sphinx configuration."""
 
-import os
-import sys
-
-import sphinx.environment
-
 from invenio_db import __version__
 
 # -- General configuration ------------------------------------------------

docs/configuration.rst

@@ -33,6 +33,19 @@ packages.
    User class used by versioning manager. Defaults to ``'User'`` if
    ``invenio_accounts`` package is installed.
 
+
+.. data:: DB_SESSION_BIND_FUNC
+
+   Function used for dynamically selecting the SQLAlchemy session bind
+   engine/connection. The function accepts the ``db.session`` instance and the original
+   ``flask_sqlalchemy.session.Session.get_bind(*args, **kwargs)`` arguments, i.e. it is
+   like a method call. Defaults to ``None``.
+
+   This function is particularly useful for dynamic connection routing, e.g. in the case
+   of primary-replica database setups, where some read operations could be routed to a
+   read replica database, while write operations would go to the primary as normal.
+
+
 .. data:: ALEMBIC
 
    Dictionary containing general configuration for Flask-Alembic. It contains
@@ -50,5 +63,5 @@ packages.
 Please check following packages for further configuration options:
 
 1. `Flask-SQLAlchemy <https://flask-sqlalchemy.readthedocs.io/en/stable/config/>`_
-2. `Flask-Alembic <https://flask-alembic.readthedocs.io/en/stable/#configuration>`_
+2. `Flask-Alembic <https://flask-alembic.readthedocs.io/en/latest/config/>`_
 3. `SQLAlchemy-Continuum <https://sqlalchemy-continuum.readthedocs.io/en/latest/configuration.html>`_

invenio_db/ext.py

@@ -77,6 +77,9 @@ def init_db(self, app, entry_point_group="invenio_db.models", **kwargs):
         # Needed for before/after_flush/commit/rollback events
         app.config.setdefault("SQLALCHEMY_TRACK_MODIFICATIONS", True)
 
+        # Set the session bind function
+        app.config.setdefault("DB_SESSION_BIND_FUNC", None)
+
         # Initialize Flask-SQLAlchemy extension.
         database = kwargs.get("db", db)
         database.init_app(app)

invenio_db/shared.py

@@ -8,7 +8,9 @@
 
 """Shared database object for Invenio."""
 
+from flask import current_app
 from flask_sqlalchemy import SQLAlchemy as FlaskSQLAlchemy
+from flask_sqlalchemy.session import Session as BaseSession
 from sqlalchemy import MetaData, event, util
 from sqlalchemy.engine import Engine
 from sqlalchemy.sql import text
@@ -106,7 +108,43 @@ def do_sqlite_begin(dbapi_connection):
     dbapi_connection.execute(text("BEGIN"))
 
 
-db = SQLAlchemy(metadata=metadata)
+# NOTE: We are defining this class here, since the Flask-SQLAlchemy extension doesn't
+# follow a traditional "lazy" configuration, to allow overriding `session_options` via
+# the usual application config.
+class Session(BaseSession):
+    """Custom session class to allow configuring dynamic engine/connection binding."""
+
+    def get_bind(self, *args, **kwargs):
+        """Hijacked to allow dynamic binding.
+
+        Example usage for routing anonymous read requests to a read replica configured
+        bind:
+
+        .. code-block::python
+
+            from flask import request, current_app
+            from flask_login import current_user
+
+            def routed_bind(session, *args, **kwargs)
+                if request and request.method in ["GET", "HEAD"]:
+                    read_endpoints = current_app.config.get(
+                        "MY_READ_REPLICA_ENDOINTS",
+                        ["records.detail", "records.file_download"],
+                    )
+                    if current_user.is_anonymous and request.endpoint in read_endpoints:
+                        return session._db.engines["read_replica"]
+        """
+        if current_app:  # We can't be sure that we're in the Flask app context
+            bind_func = current_app.config.get("DB_SESSION_BIND_FUNC")
+            if bind_func and callable(bind_func):
+                bind = bind_func(self, *args, **kwargs)
+                if bind:
+                    return bind
+
+        return super().get_bind(*args, **kwargs)
+
+
+db = SQLAlchemy(metadata=metadata, session_options={"class_": Session})
 """Shared database instance using Flask-SQLAlchemy extension.
 
 This object is initialized during initialization of ``InvenioDB``