Repository: vytas7/falcon-sqla

Issue:
Allow customization of session closure
The default logic, commit on success, rollback on failure is a sensible one
https://github.com/vytas7/falcon-sqla/blob/e4d7865892c1111c49ecb7e06a7b2a425990ec73/falcon_sqla/middleware.py#L53-L56

maybe it could be useful to add the ability of a resource to customize this logic.
I'm not sure how to best implement it. Probably a function in the resource class that returns what to do (commit, rollback, nothing) is the best solution.


Discussion:

Patch:
diff --git a/README.rst b/README.rst
index 051accf..b432112 100644
--- a/README.rst
+++ b/README.rst
@@ -81,7 +81,7 @@ include the provided engines in the runtime bind selection logic:
     manager = falcon_sqla.Manager(engine)
 
     read_replica = create_engine('dialect+driver://my/database.replica')
-    manager.add_engine(read_replica, 'r')
+    manager.add_engine(read_replica, falcon_sqla.EngineRole.READ)
 
 
 The ``Manager.get_bind()`` method can be overridden to implement custom engine
diff --git a/docs/api/constants.rst b/docs/api/constants.rst
new file mode 100644
index 0000000..9b431ef
--- /dev/null
+++ b/docs/api/constants.rst
@@ -0,0 +1,5 @@
+Constants and Enums
+===================
+
+.. automodule:: falcon_sqla.constants
+    :members:
diff --git a/docs/api/index.rst b/docs/api/index.rst
index 7655bd1..282199e 100644
--- a/docs/api/index.rst
+++ b/docs/api/index.rst
@@ -5,6 +5,7 @@ Module Reference
     :maxdepth: 1
 
     manager
+    constants
     options
     middleware
     session
diff --git a/docs/user.rst b/docs/user.rst
index 57e2ffd..1ba53e8 100644
--- a/docs/user.rst
+++ b/docs/user.rst
@@ -68,7 +68,7 @@ the runtime bind selection logic:
     manager = falcon_sqla.Manager(engine)
 
     read_replica = create_engine('dialect+driver://my/database.replica')
-    manager.add_engine(read_replica, 'r')
+    manager.add_engine(read_replica, falcon_sqla.EngineRole.READ)
 
 The :func:`Manager.get_bind() <falcon_sqla.Manager.get_bind>` method can be
 overridden to implement custom engine selection logic for more complex use
diff --git a/falcon_sqla/__init__.py b/falcon_sqla/__init__.py
index a65aa04..a82dd68 100644
--- a/falcon_sqla/__init__.py
+++ b/falcon_sqla/__init__.py
@@ -12,10 +12,14 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+from .constants import EngineRole
+from .constants import SessionCleanup
 from .manager import Manager
 from .version import __version__
 
 __all__ = [
+    'EngineRole',
     'Manager',
+    'SessionCleanup',
     '__version__',
 ]
diff --git a/falcon_sqla/constants.py b/falcon_sqla/constants.py
new file mode 100644
index 0000000..bbf8622
--- /dev/null
+++ b/falcon_sqla/constants.py
@@ -0,0 +1,94 @@
+#  Copyright 2021-2023 Vytautas Liuolia
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import enum
+
+
+class EngineRole(enum.Enum):
+    """Engine role in :class:`~falcon_sqla.Manager`."""
+
+    READ = 'r'
+    """This engine is only suitable for reading (e.g., a read replica)."""
+
+    WRITE = 'w'
+    """This engine is only preferred for writing.
+
+    Note:
+        A :attr:`~EngineRole.WRITE` engine might still receive read queries
+        when, for instance, these are issued from non-idempotent HTTP
+        methods. This role should be seen as merely a hint that this engine
+        should not be picked when a :attr:`~EngineRole.READ` one is sufficient.
+    """
+
+    READ_WRITE = 'rw'
+    """This engine is suitable for all types of queries.
+
+    When :attr:`choosing <falcon_sqla.Manager.get_bind>`, this
+    engine will participate in balancing load in both :attr:`~EngineRole.READ`
+    and :attr:`~EngineRole.WRITE` contexts (unless
+    :attr:`~.SessionOptions.read_from_rw_engines` is set to ``True``).
+    """
+
+
+class SessionCleanup(enum.Enum):
+    """Session cleanup behavior.
+
+    Sessions are automatically cleaned up and returned to the pool when using
+    :class:`~falcon_sqla.Manager`\\'s
+    :func:`~falcon_sqla.Manager.session_scope` or
+    :attr:`~falcon_sqla.Manager.middleware`.
+    In addition to closing the session, to the mode-specific behavior is
+    governed by the below constants.
+
+    Unless configured otherwise, the default behavior throughout this add-on is
+    :attr:`COMMIT_ON_SUCCESS`.
+
+    .. note::
+       Customizable session cleanup complements SQLAlchemy's "reset on return"
+       behaviour, see also:
+       https://docs.sqlalchemy.org/en/20/core/pooling.html#pool-reset-on-return.
+    """
+
+    COMMIT_ON_SUCCESS = 'default'
+    """
+    Commit on success (the default behavior).
+
+    This mode attempts to commit in the case there was no exception raised in
+    the block in question (or in the case of middleware, request-response
+    cycle), otherwise rollback.
+
+    In this mode, in the case the attempt to commit results in an exception
+    itself, it is also followed up with a rollback.
+    """
+
+    COMMIT = 'commit'
+    """Always commit.
+
+    This mode always attempts to commit regardless of any exceptions raised.
+
+    Even if the commit attempt raises an exception, no rollback is performed.
+    """
+
+    ROLLBACK = 'rollback'
+    """Rollback.
+
+    This mode always attempts to roll back regardless of any exceptions raised.
+    """
+
+    CLOSE_ONLY = 'close'
+    """Close only.
+
+    This mode only closes the session. Any commit or rollback should be
+    performed explicitly in the code.
+    """
diff --git a/falcon_sqla/manager.py b/falcon_sqla/manager.py
index 93efa9f..f389734 100644
--- a/falcon_sqla/manager.py
+++ b/falcon_sqla/manager.py
@@ -19,10 +19,19 @@ import uuid
 from sqlalchemy.orm.session import sessionmaker
 from sqlalchemy.sql import Update, Delete
 
+from .constants import EngineRole, SessionCleanup
 from .middleware import Middleware
 from .session import RequestSession
 
 
+__all__ = ['Manager', 'SessionOptions']
+
+CLOSE_ONLY = SessionCleanup.CLOSE_ONLY
+COMMIT = SessionCleanup.COMMIT
+COMMIT_ON_SUCCESS = SessionCleanup.COMMIT_ON_SUCCESS
+ROLLBACK = SessionCleanup.ROLLBACK
+
+
 class Manager:
     """A manager for SQLAlchemy sessions.
 
@@ -43,7 +52,7 @@ class Manager:
     """
     def __init__(self, engine, session_cls=RequestSession, binds=None):
         self._main_engine = engine
-        self._engines = {engine: 'rw'}
+        self._engines = {engine: EngineRole.READ_WRITE}
         self._read_engines = (engine,)
         self._write_engines = (engine,)
         self._session_kwargs = {}
@@ -64,30 +73,35 @@ class Manager:
                          if self._engines.get(engine) == role)
         return filtered or engines
 
-    def add_engine(self, engine, role='r'):
+    def add_engine(self, engine, role=EngineRole.READ):
         """Adds a new engine with the specified role.
 
         Args:
             engine (Engine): An instance of a SQLAlchemy Engine.
-            role ({'r', 'rw', 'w'}, optional): The role of the engine
-                ('r': read-ony, 'rw': read-write, 'w': write-only).
-                Defaults to 'r'.
+            role (EngineRole): The role of the provided engine.
+                Defaults to :attr:`~.EngineRole.READ`.
+
+                Note:
+                    In early versions of this library, `role` used to take
+                    string values: ``'r'``, ``'w'``, ``'rw'``. These values
+                    will continue to be supported in the foreseeable future for
+                    backwards compatibility, but new code should prefer passing
+                    enum constants instead.
         """
-        if role not in {'r', 'rw', 'w'}:
-            raise ValueError("role must be one of ('r', 'rw', 'w')")
+        role = EngineRole(role)
 
         self._engines[engine] = role
-        if 'r' in role:
+        if role in {EngineRole.READ, EngineRole.READ_WRITE}:
             self._read_engines += (engine,)
-        if 'w' in role:
+        if role in {EngineRole.WRITE, EngineRole.READ_WRITE}:
             self._write_engines += (engine,)
 
         if not self.session_options.read_from_rw_engines:
             self._read_engines = self._filter_by_role(
-                self._read_engines, 'r')
+                self._read_engines, EngineRole.READ)
         if not self.session_options.write_to_rw_engines:
             self._write_engines = self._filter_by_role(
-                self._write_engines, 'w')
+                self._write_engines, EngineRole.WRITE)
 
         # NOTE(vytas): Do not tamper with custom binds.
         # NOTE(vytas): We can only rely on RequestSession and its subclasses to
@@ -122,6 +136,33 @@ class Manager:
 
         return self._Session()
 
+    def close_session(self, session, succeeded, req=None, resp=None):
+        """Close a session obtained via :func:`get_session`.
+
+        .. note:: There is no need to invoke this method manually if you are
+                  using the :func:`session_scope` context manager, or if you
+                  are using middleware.
+        """
+        session_cleanup = self.session_options.session_cleanup
+        attempt_commit = session_cleanup == COMMIT_ON_SUCCESS and succeeded
+
+        try:
+            if attempt_commit or session_cleanup == COMMIT:
+                session.commit()
+            elif session_cleanup != CLOSE_ONLY:
+                session.rollback()
+        except Exception:
+            if attempt_commit:
+                session.rollback()
+            raise
+        finally:
+            if req and resp:
+                # NOTE(vytas): Break circular references between the request
+                #   and session in case the latter was stored in req.context.
+                del session.info['req']
+                del session.info['resp']
+            session.close()
+
     @property
     def read_engines(self):
         """A tuple of read capable engines."""
@@ -136,19 +177,22 @@ class Manager:
     def session_scope(self, req=None, resp=None):
         """Provide a transactional scope around a series of operations.
 
+        The session is obtained via :func:`get_sesion`, and finalized using
+        :func:`close_session` upon exiting the context manager.
+
         Based on the ``session_scope()`` recipe from
         https://docs.sqlalchemy.org/orm/session_basics.html.
         """
         session = self.get_session(req, resp)
+        succeeded = True
 
         try:
             yield session
-            session.commit()
         except Exception:
-            session.rollback()
+            succeeded = False
             raise
         finally:
-            session.close()
+            self.close_session(session, succeeded, req, resp)
 
     @property
     def middleware(self):
@@ -164,6 +208,9 @@ class SessionOptions:
     An instance of this class is exposed via :attr:`Manager.session_options`.
 
     Attributes:
+        session_cleanup (SessionCleanup): Session cleanup mode; one of the
+            :class:`~.SessionCleanup` constants.
+            Defaults to :attr:`~.SessionCleanup.COMMIT_ON_SUCCESS`.
         no_session_methods (frozenset): HTTP methods that by default do not
             require a DB session. Defaults to
             :attr:`SessionOptions.NO_SESSION_METHODS`.
@@ -204,6 +251,7 @@ class SessionOptions:
     """
 
     __slots__ = [
+        'session_cleanup',
         'no_session_methods',
         'safe_methods',
         'read_from_rw_engines',
@@ -215,6 +263,8 @@ class SessionOptions:
     ]
 
     def __init__(self):
+        self.session_cleanup = SessionCleanup.COMMIT_ON_SUCCESS
+
         self.no_session_methods = self.NO_SESSION_METHODS
         self.safe_methods = self.SAFE_METHODS
 
diff --git a/falcon_sqla/middleware.py b/falcon_sqla/middleware.py
index 6996416..074705d 100644
--- a/falcon_sqla/middleware.py
+++ b/falcon_sqla/middleware.py
@@ -12,6 +12,8 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
+import functools
+
 from .util import ClosingStreamWrapper
 
 
@@ -48,30 +50,17 @@ class Middleware:
         """
         Clean up the session, if one was provided.
 
-        This response hook finalizes the session by calling its ``.commit()``
-        if `req_succeeded` is ``True``, and ``.rollback()`` otherwise. Finally,
-        it will close the session.
+        This response hook finalizes the session by calling the manager's
+        :func:`~falcon_sqla.Manager.close_session` method.
         """
-        def cleanup():
-            # NOTE(vytas): Break circular references between the request and
-            #   the session.
-            req.context.session = None
-            del session.info['req']
-            del session.info['resp']
-
-            session.close()
-
         session = getattr(req.context, 'session', None)
 
         if session:
-            try:
-                if req_succeeded:
-                    session.commit()
-                else:
-                    session.rollback()
-            finally:
-                if (req_succeeded and resp.stream is not None and
-                        self._options.wrap_response_stream):
-                    resp.stream = ClosingStreamWrapper(resp.stream, cleanup)
-                else:
-                    cleanup()
+            if resp.stream is not None and self._options.wrap_response_stream:
+                resp.stream = ClosingStreamWrapper(
+                    resp.stream,
+                    functools.partial(
+                        self._manager.close_session,
+                        session, req_succeeded, req, resp))
+            else:
+                self._manager.close_session(session, req_succeeded, req, resp)


Test patch:
diff --git a/tests/test_manager_engines.py b/tests/test_manager_engines.py
index c578759..3a98c43 100644
--- a/tests/test_manager_engines.py
+++ b/tests/test_manager_engines.py
@@ -1,7 +1,7 @@
 import pytest
 from sqlalchemy import create_engine
 
-from falcon_sqla import Manager
+from falcon_sqla import Manager, EngineRole
 
 
 def test_unsupported_role():
@@ -21,15 +21,15 @@ def test_use_rw_engines():
     assert manager.read_engines == (primary,)
     assert manager.write_engines == (primary,)
 
-    manager.add_engine(one, 'r')
+    manager.add_engine(one, EngineRole.READ)
     assert manager.read_engines == (primary, one)
     assert manager.write_engines == (primary,)
 
-    manager.add_engine(two, 'rw')
+    manager.add_engine(two, EngineRole.READ_WRITE)
     assert manager.read_engines == (primary, one, two)
     assert manager.write_engines == (primary, two)
 
-    manager.add_engine(three, 'w')
+    manager.add_engine(three, EngineRole.WRITE)
     assert manager.read_engines == (primary, one, two)
     assert manager.write_engines == (primary, two, three)
 
diff --git a/tests/test_middleware.py b/tests/test_middleware.py
index ef5123e..153a3da 100644
--- a/tests/test_middleware.py
+++ b/tests/test_middleware.py
@@ -145,20 +145,26 @@ def test_options(client):
 @pytest.mark.parametrize('wrap_stream', [True, False])
 @pytest.mark.parametrize('use_file_wrapper', [True, False])
 def test_wrap_response_stream(tunable_client, wrap_stream, use_file_wrapper):
-    def file_wrapper(obj, ignored_size=None):
-        size = 3
-        while True:
-            chunk = obj.read(size)
-            if not chunk:
-                break
-            history.append(chunk)
-            yield chunk
-            size += 1
+    class FileWrapper:
+        def __init__(self, filelike, ignored_size=8192):
+            self._filelike = filelike
+            if hasattr(filelike, 'close'):
+                self.close = filelike.close
+
+        def __iter__(self):
+            size = 3
+            while True:
+                chunk = self._filelike.read(size)
+                if not chunk:
+                    break
+                history.append(chunk)
+                yield chunk
+                size += 1
 
     history = []
 
     client = tunable_client({'wrap_response_stream': wrap_stream})
-    wrapper = file_wrapper if use_file_wrapper else None
+    wrapper = FileWrapper if use_file_wrapper else None
 
     client.simulate_post('/languages',
                          json={'name': 'Python', 'created': 1991})
diff --git a/tests/test_session_scope.py b/tests/test_session_scope.py
index c58fac5..a51ae0c 100644
--- a/tests/test_session_scope.py
+++ b/tests/test_session_scope.py
@@ -1,8 +1,10 @@
 import falcon
 import falcon.testing
 import pytest
+from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm.session import Session
 
+from falcon_sqla.manager import SessionCleanup
 from falcon_sqla.manager import Manager
 from falcon_sqla.session import RequestSession
 
@@ -71,3 +73,70 @@ def test_generic_scope(database):
         malbolge = session.query(database.Language).first()
         assert malbolge.name == 'Malbolge'
         assert malbolge.created == 1998
+
+
+@pytest.mark.parametrize('cleanup,expected', [
+    (SessionCleanup.CLOSE_ONLY, 1999),
+    (SessionCleanup.COMMIT, 1998),
+    (SessionCleanup.COMMIT_ON_SUCCESS, 1998),
+    (SessionCleanup.ROLLBACK, 1999),
+])
+def test_session_cleanup(database, cleanup, expected):
+    manager = Manager(database.write_engine)
+    manager.session_options.session_cleanup = cleanup
+
+    with manager.session_scope() as session:
+        session.add(database.Language(name='Malbolge', created=1998+1))
+        session.commit()
+
+    with manager.session_scope() as session:
+        malbolge = session.query(database.Language).first()
+        malbolge.created = 1998
+
+    with manager.session_scope() as session:
+        malbolge = session.query(database.Language).first()
+        assert malbolge.created == expected
+
+
+@pytest.mark.parametrize('cleanup,expected', [
+    (SessionCleanup.CLOSE_ONLY, 1999),
+    (SessionCleanup.COMMIT, 1998),
+    (SessionCleanup.COMMIT_ON_SUCCESS, 1999),
+    (SessionCleanup.ROLLBACK, 1999),
+])
+def test_close_on_error(database, cleanup, expected):
+    manager = Manager(database.write_engine)
+    manager.session_options.session_cleanup = cleanup
+
+    with manager.session_scope() as session:
+        session.add(database.Language(name='Malbolge', created=1998+1))
+        session.commit()
+
+    with pytest.raises(ZeroDivisionError):
+        with manager.session_scope() as session:
+            malbolge = session.query(database.Language).first()
+            malbolge.created = 1998
+            malbolge.statement = 0/0
+
+    with manager.session_scope() as session:
+        malbolge = session.query(database.Language).first()
+        assert malbolge.created == expected
+
+
+@pytest.mark.parametrize('cleanup,expected', [
+    (SessionCleanup.COMMIT, 1998),
+    (SessionCleanup.COMMIT_ON_SUCCESS, 1999),
+])
+def test_close_error_on_commit(database, cleanup, expected):
+    manager = Manager(database.write_engine)
+    manager.session_options.session_cleanup = cleanup
+
+    with manager.session_scope() as session:
+        malbolge = database.Language(name='Malbolge', created=1998)
+        session.add(malbolge)
+        session.commit()
+        idn = malbolge.id
+
+    with pytest.raises(IntegrityError):
+        with manager.session_scope() as session:
+            session.add(database.Language(id=idn, name='Error', created=2525))