fix: address oauth-status review feedback (N+1 queries, error visibility)

  Batch /oauth/status now issues one gateway SELECT and reuses a single
  TokenStorageService across all ids instead of one per id. Batch access
  checks distinguish expected 403/404 omissions from 5xx/unexpected
  failures, which are now logged instead of silently dropped. Database
  token lookups re-raise on failure instead of collapsing to the same
  "missing" result as a token that was never issued, so a transient DB
  error is distinguishable in logs from "never authorized".

Signed-off-by: Marek Dano <mk.dano@gmail.com>

Author: marekdano

docs/docs/manage/oauth-troubleshooting.md

@@ -76,7 +76,8 @@ state_data = {
 | `/oauth/authorize/{gateway_id}` | GET | Initiates OAuth flow, redirects to provider |
 | `/vault/authorize/{server_id}` | GET | Per-user OAuth credential connection for team virtual servers (shared access control with `/oauth/authorize/{gateway_id}`) |
 | `/oauth/callback` | GET | Handles OAuth callback, exchanges code for tokens |
-| `/oauth/status/{gateway_id}` | GET | Returns OAuth configuration status |
+| `/oauth/status/{gateway_id}` | GET | Returns OAuth configuration status; for `authorization_code` gateways also includes the caller's own `user_token_status` (valid/near_expiry/expired/missing) |
+| `/oauth/status?gateway_ids=a&gateway_ids=b` | GET | Batch equivalent of the above, keyed by gateway id, so a grid of cards issues one request instead of N |
 | `/oauth/fetch-tools/{gateway_id}` | POST | Fetches tools from MCP server after OAuth completion |
 | `/oauth/registered-clients` | GET | Lists all DCR-registered OAuth clients. Requires `admin.oauth_clients:read` and un-narrowed platform admin access |
 | `/oauth/registered-clients/{gateway_id}` | GET | Gets registered client for specific gateway. Requires `admin.oauth_clients:read` and un-narrowed platform admin access |

mcpgateway/routers/oauth_router.py

@@ -1333,7 +1333,7 @@ def custom_redirect_after_callback(url: str, status_code: int) -> RedirectRespon
     return RedirectResponse(url=url, status_code=status_code, headers={"Referrer-Policy": "no-referrer"})
 
 
-async def _get_caller_token_status(db: Session, current_user: Any, gateway_id: str) -> Dict[str, Any]:
+async def _get_caller_token_status(db: Session, current_user: Any, gateway_id: str, *, token_storage: Optional[TokenStorageService] = None) -> Dict[str, Any]:
     """Look up the caller's own OAuth token state for a gateway.
 
     Wires the already-implemented ``TokenStorageService.get_token_info`` into
@@ -1344,6 +1344,9 @@ async def _get_caller_token_status(db: Session, current_user: Any, gateway_id: s
         db: Active database session.
         current_user: Authenticated requester context (dict or EmailUserResponse).
         gateway_id: Gateway identifier to look up.
+        token_storage: Optional pre-built ``TokenStorageService`` to reuse across
+            multiple lookups (batch endpoint) instead of constructing a new one
+            per gateway.
 
     Returns:
         Dict with ``authorized`` (bool) and ``status`` (one of "missing",
@@ -1354,9 +1357,16 @@ async def _get_caller_token_status(db: Session, current_user: Any, gateway_id: s
     if requester_email == "unknown" or not requester_email.strip():
         return {"status": "missing", "authorized": False}
 
-    user_context = _build_user_context(current_user)
-    token_storage = TokenStorageService(db, user_context)
-    info = await token_storage.get_token_info(gateway_id, requester_email)
+    if token_storage is None:
+        token_storage = TokenStorageService(db, _build_user_context(current_user))
+
+    try:
+        info = await token_storage.get_token_info(gateway_id, requester_email)
+    except Exception as e:
+        # The backend already logs its own failure; this line ties it to the caller/gateway
+        # so a transient lookup failure is distinguishable in logs from a genuinely missing token.
+        logger.error("OAuth token status lookup failed for gateway=%s user=%s: %s", gateway_id, requester_email, str(e))
+        return {"status": "missing", "authorized": False}
 
     if not info:
         return {"status": "missing", "authorized": False}
@@ -1371,6 +1381,45 @@ async def _get_caller_token_status(db: Session, current_user: Any, gateway_id: s
     }
 
 
+def _build_oauth_status_payload(gateway: Gateway) -> Dict[str, Any]:
+    """Build the OAuth config portion of a gateway's status payload (no I/O).
+
+    Shared by the single-gateway and batch status endpoints so the two never
+    drift. Caller is responsible for gateway lookup, access enforcement, and
+    attaching ``user_token_status`` for ``authorization_code`` grants.
+
+    Args:
+        gateway: Gateway record with ``oauth_config`` already loaded.
+
+    Returns:
+        Dict describing OAuth enablement and, when configured, grant details.
+    """
+    if not gateway.oauth_config:
+        return {"oauth_enabled": False, "message": "Gateway is not configured for OAuth"}
+
+    oauth_config = gateway.oauth_config
+    grant_type = oauth_config.get("grant_type")
+
+    if grant_type == "authorization_code":
+        return {
+            "oauth_enabled": True,
+            "grant_type": grant_type,
+            "client_id": oauth_config.get("client_id"),
+            "scopes": oauth_config.get("scopes", []),
+            "authorization_url": oauth_config.get("authorization_url"),
+            "redirect_uri": oauth_config.get("redirect_uri"),
+            "message": "Gateway configured for Authorization Code flow",
+        }
+
+    return {
+        "oauth_enabled": True,
+        "grant_type": grant_type,
+        "client_id": oauth_config.get("client_id"),
+        "scopes": oauth_config.get("scopes", []),
+        "message": f"Gateway configured for {grant_type} flow",
+    }
+
+
 @oauth_router.get("/status/{gateway_id}")
 async def get_oauth_status(
     gateway_id: str,
@@ -1408,32 +1457,10 @@ async def get_oauth_status(
 
         await _enforce_gateway_access(gateway_id, gateway, current_user, db, request=request)
 
-        if not gateway.oauth_config:
-            return {"oauth_enabled": False, "message": "Gateway is not configured for OAuth"}
-
-        # Get OAuth configuration info
-        oauth_config = gateway.oauth_config
-        grant_type = oauth_config.get("grant_type")
-
-        if grant_type == "authorization_code":
-            return {
-                "oauth_enabled": True,
-                "grant_type": grant_type,
-                "client_id": oauth_config.get("client_id"),
-                "scopes": oauth_config.get("scopes", []),
-                "authorization_url": oauth_config.get("authorization_url"),
-                "redirect_uri": oauth_config.get("redirect_uri"),
-                "message": "Gateway configured for Authorization Code flow",
-                "user_token_status": await _get_caller_token_status(db, current_user, gateway_id),
-            }
-        else:
-            return {
-                "oauth_enabled": True,
-                "grant_type": grant_type,
-                "client_id": oauth_config.get("client_id"),
-                "scopes": oauth_config.get("scopes", []),
-                "message": f"Gateway configured for {grant_type} flow",
-            }
+        payload = _build_oauth_status_payload(gateway)
+        if payload.get("grant_type") == "authorization_code":
+            payload["user_token_status"] = await _get_caller_token_status(db, current_user, gateway_id)
+        return payload
 
     except HTTPException:
         raise
@@ -1480,13 +1507,39 @@ async def get_oauth_status_batch(
     if len(deduped_ids) > OAUTH_STATUS_BATCH_MAX_IDS:
         raise HTTPException(status_code=400, detail=f"Too many gateway_ids requested (max {OAUTH_STATUS_BATCH_MAX_IDS})")
 
+    # Single query for all requested gateways, and one TokenStorageService reused across
+    # the loop, instead of get_oauth_status()'s per-id gateway SELECT + TokenStorageService
+    # construction - the batch route exists specifically to avoid N+1 round trips.
+    gateways_by_id = {gw.id: gw for gw in db.execute(select(Gateway).where(Gateway.id.in_(deduped_ids))).scalars().all()}
+    token_storage = TokenStorageService(db, _build_user_context(current_user))
+
     results: Dict[str, Dict[str, Any]] = {}
     for gateway_id in deduped_ids:
+        gateway = gateways_by_id.get(gateway_id)
+        if not gateway:
+            # Not found - omit rather than failing the batch.
+            continue
+
         try:
-            results[gateway_id] = await get_oauth_status(gateway_id, request, current_user, db)
-        except HTTPException:
-            # Not found / not accessible to this caller - omit rather than failing the batch.
+            await _enforce_gateway_access(gateway_id, gateway, current_user, db, request=request)
+        except HTTPException as exc:
+            if exc.status_code >= 500:
+                logger.error("OAuth status batch: access check failed for gateway=%s: %s", gateway_id, exc.detail)
+            # Not accessible to this caller (or a lookup failure, logged above) - omit rather than failing the batch.
             continue
+        except Exception as exc:
+            logger.error("OAuth status batch: access check raised for gateway=%s: %s", gateway_id, str(exc))
+            continue
+
+        try:
+            payload = _build_oauth_status_payload(gateway)
+            if payload.get("grant_type") == "authorization_code":
+                payload["user_token_status"] = await _get_caller_token_status(db, current_user, gateway_id, token_storage=token_storage)
+            results[gateway_id] = payload
+        except Exception as exc:
+            logger.error("OAuth status batch: failed to build status for gateway=%s: %s", gateway_id, str(exc))
+            continue
+
     return results
 
 

mcpgateway/services/token_backends/db_backend.py

@@ -267,13 +267,24 @@ async def get_token_info(
 
         Phase 1: team_id parameter is IGNORED.
 
+        Returns ``None`` only when no token record exists for this
+        (gateway_id, app_user_email) pair. An unexpected failure (e.g. a DB
+        error) is logged and re-raised rather than swallowed to ``None``, so
+        callers exposing this through a user-facing status field (see
+        ``mcpgateway.routers.oauth_router._get_caller_token_status``) can
+        distinguish "never authorized" from "lookup failed" - the two read
+        identically to a caller if both collapse to ``None``.
+
         Args:
             gateway_id: Gateway ID
             team_id: Team identifier (IGNORED in Phase 1)
             app_user_email: ContextForge user email
 
         Returns:
-            Token info dict or None
+            Token info dict, or None if no token is stored.
+
+        Raises:
+            Exception: Propagated from the underlying database query on failure.
         """
         try:
             # PHASE 1: Query by (gateway_id, app_user_email) - team_id IGNORED
@@ -302,7 +313,7 @@ async def get_token_info(
 
         except Exception as e:
             logger.error("Failed to get token info: %s", str(e))
-            return None
+            raise
 
     async def revoke_user_tokens(
         self,

tests/unit/mcpgateway/routers/test_oauth_router.py

@@ -1274,10 +1274,29 @@ async def test_get_oauth_status_user_token_status_not_shared_across_users(self,
 
         mock_token_storage.get_token_info.assert_awaited_once_with("gateway123", "other@example.com")
 
+    @pytest.mark.asyncio
+    async def test_get_oauth_status_user_token_status_lookup_failure_logs_and_reads_missing(self, mock_db, mock_gateway, mock_current_user, mock_request):
+        """A backend lookup failure (e.g. a DB error) is logged, distinguishing it from a genuinely missing token in logs,
+        while still surfacing the same safe "missing"/unauthorized shape to the client rather than a 500."""
+        mock_db.execute.return_value.scalar_one_or_none.return_value = mock_gateway
+
+        from mcpgateway.routers.oauth_router import get_oauth_status
+
+        with patch("mcpgateway.routers.oauth_router.TokenStorageService") as mock_token_storage_class:
+            mock_token_storage = Mock()
+            mock_token_storage.get_token_info = AsyncMock(side_effect=RuntimeError("db unavailable"))
+            mock_token_storage_class.return_value = mock_token_storage
+
+            with patch("mcpgateway.routers.oauth_router.logger") as mock_logger:
+                result = await get_oauth_status("gateway123", mock_request, mock_current_user, mock_db)
+
+        assert result["user_token_status"] == {"status": "missing", "authorized": False}
+        mock_logger.error.assert_called_once()
+
     @pytest.mark.asyncio
     async def test_get_oauth_status_batch_success(self, mock_db, mock_gateway, mock_current_user, mock_request):
         """Batch endpoint returns the same per-gateway payload as the single endpoint, keyed by gateway id."""
-        mock_db.execute.return_value.scalar_one_or_none.return_value = mock_gateway
+        mock_db.execute.return_value.scalars.return_value.all.return_value = [mock_gateway]
 
         from mcpgateway.routers.oauth_router import get_oauth_status_batch
 
@@ -1293,17 +1312,63 @@ async def test_get_oauth_status_batch_success(self, mock_db, mock_gateway, mock_
         assert result["gateway123"]["oauth_enabled"] is True
         assert result["gateway123"]["user_token_status"]["status"] == "missing"
 
+    @pytest.mark.asyncio
+    async def test_get_oauth_status_batch_single_gateway_query(self, mock_db, mock_gateway, mock_current_user, mock_request):
+        """Batch endpoint issues one Gateway SELECT and one TokenStorageService for the whole batch, not one per id."""
+        mock_db.execute.return_value.scalars.return_value.all.return_value = [mock_gateway]
+
+        from mcpgateway.routers.oauth_router import get_oauth_status_batch
+
+        with patch("mcpgateway.routers.oauth_router.TokenStorageService") as mock_token_storage_class:
+            mock_token_storage = Mock()
+            mock_token_storage.get_token_info = AsyncMock(return_value=None)
+            mock_token_storage_class.return_value = mock_token_storage
+
+            result = await get_oauth_status_batch(mock_request, ["gateway123", "gateway123"], mock_current_user, mock_db)
+
+        assert mock_db.execute.call_count == 1
+        mock_token_storage_class.assert_called_once()
+        assert result["gateway123"]["oauth_enabled"] is True
+
     @pytest.mark.asyncio
     async def test_get_oauth_status_batch_omits_inaccessible_gateways(self, mock_db, mock_current_user, mock_request):
         """A gateway id that 404s or 403s for this caller is silently dropped, not surfaced as a batch failure."""
-        mock_db.execute.return_value.scalar_one_or_none.return_value = None  # every id -> gateway not found
+        mock_db.execute.return_value.scalars.return_value.all.return_value = []  # every id -> gateway not found
 
         from mcpgateway.routers.oauth_router import get_oauth_status_batch
 
         result = await get_oauth_status_batch(mock_request, ["missing1", "missing2"], mock_current_user, mock_db)
 
         assert result == {}
 
+    @pytest.mark.asyncio
+    async def test_get_oauth_status_batch_logs_access_check_server_errors(self, mock_db, mock_gateway, mock_current_user, mock_request):
+        """A 5xx from the per-gateway access check is logged and omitted, not silently swallowed like a 404/403."""
+        mock_db.execute.return_value.scalars.return_value.all.return_value = [mock_gateway]
+
+        from mcpgateway.routers.oauth_router import get_oauth_status_batch
+
+        with patch("mcpgateway.routers.oauth_router._enforce_gateway_access", new=AsyncMock(side_effect=HTTPException(status_code=500, detail="boom"))):
+            with patch("mcpgateway.routers.oauth_router.logger") as mock_logger:
+                result = await get_oauth_status_batch(mock_request, ["gateway123"], mock_current_user, mock_db)
+
+        assert result == {}
+        mock_logger.error.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_get_oauth_status_batch_does_not_log_404_or_403(self, mock_db, mock_gateway, mock_current_user, mock_request):
+        """A 404/403 from the per-gateway access check is omitted without an error log - it's an expected outcome, not a fault."""
+        mock_db.execute.return_value.scalars.return_value.all.return_value = [mock_gateway]
+
+        from mcpgateway.routers.oauth_router import get_oauth_status_batch
+
+        with patch("mcpgateway.routers.oauth_router._enforce_gateway_access", new=AsyncMock(side_effect=HTTPException(status_code=403, detail="nope"))):
+            with patch("mcpgateway.routers.oauth_router.logger") as mock_logger:
+                result = await get_oauth_status_batch(mock_request, ["gateway123"], mock_current_user, mock_db)
+
+        assert result == {}
+        mock_logger.error.assert_not_called()
+
     @pytest.mark.asyncio
     async def test_get_oauth_status_batch_requires_gateway_ids(self, mock_db, mock_current_user, mock_request):
         from mcpgateway.routers.oauth_router import get_oauth_status_batch

tests/unit/mcpgateway/services/test_token_storage_service.py

@@ -828,9 +828,10 @@ async def test_get_token_info_not_found(service, mock_db):
 
 @pytest.mark.asyncio
 async def test_get_token_info_exception(service, mock_db):
+    """A backend failure propagates rather than collapsing to None, which is reserved for "no token stored"."""
     mock_db.execute.side_effect = Exception("DB error")
-    result = await service.get_token_info("gw-1", "user@test.com")
-    assert result is None
+    with pytest.raises(Exception, match="DB error"):
+        await service.get_token_info("gw-1", "user@test.com")
 
 
 # ---------- revoke_user_tokens ----------

tests/unit/mcpgateway/services/token_backends/test_db_backend.py

@@ -603,18 +603,23 @@ async def test_get_token_info_expired_token(backend_with_encryption, mock_db):
 
 @pytest.mark.asyncio
 async def test_get_token_info_exception(backend_with_encryption, mock_db):
-    """Test get_token_info handles exceptions gracefully."""
+    """get_token_info logs and re-raises on failure rather than returning None.
+
+    None is reserved for "no token stored" - collapsing a lookup failure into the same
+    value would make it indistinguishable from "never authorized" to callers exposing
+    this through a user-facing status field.
+    """
     mock_db.execute.side_effect = Exception("Database error")
 
     with patch("mcpgateway.services.token_backends.db_backend.logger") as mock_logger:
-        result = await backend_with_encryption.get_token_info(
-            gateway_id="gw-1",
-            team_id="team-1",
-            app_user_email="user@test.com",
-        )
+        with pytest.raises(Exception, match="Database error"):
+            await backend_with_encryption.get_token_info(
+                gateway_id="gw-1",
+                team_id="team-1",
+                app_user_email="user@test.com",
+            )
 
         mock_logger.error.assert_called_once()
-        assert result is None
 
 
 @pytest.mark.asyncio