[BUG FIX][DB]: Restore transaction control to get_db() for middleware sessions (#3731) (#3813)

* fix(db): restore transaction control to get_db() for middleware sessions

PR #3600 introduced a transaction management violation where
ObservabilityMiddleware commits the shared database session instead
of get_db(), breaking the established contract where get_db() controls
transaction boundaries. This creates data integrity risks where failed
validations can be committed to the database.

This fix restores the correct behavior:
- Middleware manages session lifecycle (create/close)
- get_db() manages transactions (commit/rollback)

Changes:
- Remove commit logic from ObservabilityMiddleware (observability_middleware.py:210-216)
- Add commit/rollback handling to get_db() for middleware sessions (main.py:3137-3164)
- Update get_db() docstring to document transaction control responsibility
- Update 2 existing tests to reflect new behavior
- Add 7 comprehensive tests for transaction semantics

Security implications:
- Fixes data integrity bug where invalid data could be committed
- Maintains proper transaction isolation per request
- Preserves connection invalidation on broken connections
- No impact on auth/RBAC (middleware runs before route handlers)

Trade-offs:
- Observability data (traces/spans) is rolled back on errors (acceptable - best-effort tracing)

Closes #3731

Signed-off-by: Mohan Lakshmaiah <mohan.economist@gmail.com>

* test(db): add coverage for double-failure edge case in get_db()

Signed-off-by: Mohan Lakshmaiah <mohan.economist@gmail.com>

* fix(tests): clean up lint violations in transaction control tests

Remove unused AsyncMock import and unused variable assignments
flagged by ruff (F401, F841). Apply isort/black formatting.

Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>

---------

Signed-off-by: Mohan Lakshmaiah <mohan.economist@gmail.com>
Signed-off-by: Mihai Criveti <crivetimihai@gmail.com>
Co-authored-by: Mihai Criveti <crivetimihai@gmail.com>

Author: 2 people

mcpgateway/main.py

@@ -3092,8 +3092,17 @@ def get_db(request: Request = None):
 
     When observability is enabled, this reuses the session created by
     ObservabilityMiddleware (stored in request.state.db) to avoid duplicate
-    session creation. When observability is disabled or the
-    middleware hasn't created a session, this creates its own session.
+    session creation. When observability is disabled or the middleware hasn't
+    created a session, this creates its own session.
+
+    **Transaction Control**: This function ALWAYS controls transaction boundaries
+    (commit/rollback) regardless of whether it creates the session or reuses one
+    from middleware. This ensures predictable transaction semantics for route
+    handlers and maintains data integrity.
+
+    **Session Lifecycle**: Middleware manages session lifecycle (create/close)
+    while this function manages transactions (commit/rollback). This separation
+    of concerns prevents the transaction management violation described in #3731.
 
     Commits the transaction on successful completion to avoid implicit rollbacks
     for read-only operations. Rolls back explicitly on exception.
@@ -3114,7 +3123,10 @@ def get_db(request: Request = None):
         Exception: Re-raises any exception after rolling back the transaction.
 
     Ensures:
-        The database session is closed after the request completes, even in the case of an exception.
+        - Transaction is committed on success (for both owned and reused sessions)
+        - Transaction is rolled back on error (for both owned and reused sessions)
+        - Session is closed only if created by this function (not if reused from middleware)
+        - Broken connections are invalidated to prevent pool corruption
 
     Examples:
         >>> # Test that get_db returns a generator
@@ -3138,9 +3150,32 @@ def get_db(request: Request = None):
         db = request.state.db
         if db is not None:
             logger.debug(f"[GET_DB] Reusing session from middleware: {id(db)}")
-            # Yield the middleware's session without closing it
-            # The middleware will handle commit/rollback/close
-            yield db
+            # Yield the middleware's session. We control transactions, middleware controls lifecycle.
+            try:
+                yield db
+                # Commit on successful completion (only if transaction still active)
+                # The transaction can become inactive if an exception occurred during
+                # async context manager cleanup (e.g., CancelledError during MCP session teardown).
+                if db.is_active:
+                    db.commit()
+            except Exception:
+                try:
+                    # Always call rollback() in exception handler.
+                    # rollback() is safe to call even when is_active=False - it succeeds and
+                    # restores the session to a usable state. When is_active=False (e.g., after
+                    # IntegrityError), rollback() is actually REQUIRED to clear the failed state.
+                    # Skipping rollback when is_active=False would leave the session unusable.
+                    db.rollback()
+                except Exception:
+                    # Connection is broken - invalidate to remove from pool
+                    # This handles cases like PgBouncer query_wait_timeout where
+                    # the connection is dead and rollback itself fails
+                    try:
+                        db.invalidate()
+                    except Exception:
+                        pass  # nosec B110 - Best effort cleanup on connection failure
+                raise
+            # Don't close - middleware owns the session lifecycle
             return
 
     # Fallback: Create our own session (observability disabled or middleware didn't create one)

mcpgateway/middleware/observability_middleware.py

@@ -207,13 +207,10 @@ async def dispatch(self, request: Request, call_next: Callable) -> Response:
                 except Exception as end_trace_error:
                     logger.warning(f"Failed to end trace {trace_id}: {end_trace_error}")
 
-            # Commit the shared session (used by both observability and route handler)
-            # Note: Some route handlers may have already committed. The is_active check
-            # ensures we only commit if the transaction is still open. Services that
-            # explicitly commit will have already closed their transaction.
-            # Only commit if the transaction is still active AND has uncommitted changes
-            if db.is_active and db.in_transaction():
-                db.commit()
+            # NOTE: Transaction control delegated to get_db()
+            # Middleware only manages session lifecycle (create/close), not transactions.
+            # get_db() will commit on success or rollback on error to maintain
+            # predictable transaction semantics for route handlers (Issue #3731).
 
             return response
 

tests/unit/mcpgateway/middleware/test_observability_middleware.py

@@ -288,8 +288,9 @@ async def mock_call_next(request):
     # Verify only one session was created
     assert len(session_instances) == 1, f"Expected 1 session, but {len(session_instances)} were created"
 
-    # Verify session was committed and closed
-    session_instances[0].commit.assert_called_once()
+    # Verify session was closed (lifecycle management)
+    # Note: Middleware no longer commits - transaction control delegated to get_db() (Issue #3731)
+    session_instances[0].commit.assert_not_called()
     session_instances[0].close.assert_called_once()
 
     # Verify response
@@ -341,6 +342,7 @@ async def test_get_db_reuses_middleware_session():
     # Create a mock request with a session in state
     mock_request = MagicMock(spec=Request)
     mock_session = MagicMock()
+    mock_session.is_active = True  # Required for commit check
     mock_request.state.db = mock_session
 
     # Call get_db with the request
@@ -350,14 +352,17 @@ async def test_get_db_reuses_middleware_session():
     # Verify it returns the same session
     assert db is mock_session, "get_db should return the middleware's session"
 
-    # Verify the session is NOT closed (middleware will handle that)
+    # Complete the generator (simulating successful request)
     try:
         next(db_generator)
     except StopIteration:
         pass
 
+    # Verify get_db() commits the middleware session (Issue #3731 fix)
+    # Transaction control is now delegated to get_db(), not middleware
+    mock_session.commit.assert_called_once()
+    # Verify the session is NOT closed (middleware will handle that)
     mock_session.close.assert_not_called()
-    mock_session.commit.assert_not_called()
 
 
 @pytest.mark.asyncio