diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 6969a53cd2..1cf5cb1feb 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -541,24 +541,29 @@ async def resolve_tenant_for_user(email: str, request: Request | None = None) -> )(email=email, request=request) -@asynccontextmanager -async def _tenant_session_with_context( - tenant_id: str, -) -> AsyncGenerator[AsyncSession, None]: - token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) - try: - async with get_async_session_context_manager(tenant_id) as db_session: - yield db_session - finally: - CURRENT_TENANT_ID_CONTEXTVAR.reset(token) - - class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): reset_password_token_secret = USER_AUTH_SECRET verification_token_secret = USER_AUTH_SECRET verification_token_lifetime_seconds = AUTH_COOKIE_EXPIRE_TIME_SECONDS user_db: SQLAlchemyUserDatabase[User, uuid.UUID] + @asynccontextmanager + async def _tenant_session_with_bound_user_db( + self, tenant_id: str + ) -> AsyncGenerator[AsyncSession, None]: + """Run tenant work on one AsyncSession; user_db writes use that connection.""" + token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) + previous_user_db = self.user_db + try: + async with get_async_session_context_manager(tenant_id) as db_session: + self.user_db = SQLAlchemyUserDatabase[User, uuid.UUID]( + db_session, User, OAuthAccount + ) + yield db_session + finally: + self.user_db = previous_user_db + CURRENT_TENANT_ID_CONTEXTVAR.reset(token) + async def get_by_email(self, user_email: str) -> User: tenant_id = fetch_ee_implementation_or_noop( "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None @@ -963,8 +968,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): if not tenant_id: raise HTTPException(status_code=401, detail="User not found") - # Proceed with the tenant context - async with _tenant_session_with_context(tenant_id) as db_session: + async with self._tenant_session_with_bound_user_db(tenant_id) as db_session: verify_email_in_whitelist(account_email, tenant_id, oauth_name, account_id) oauth_security_settings = get_security_settings() effective_valid_email_domains = ( @@ -991,14 +995,6 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): "single sign-on.", ) - # NOTE(rkuo): If this UserManager is instantiated per connection - # should we even be doing this here? - if MULTI_TENANT: - tenant_user_db = SQLAlchemyUserDatabase[User, uuid.UUID]( - db_session, User, OAuthAccount - ) - self.user_db = tenant_user_db - oauth_account_dict = { "oauth_name": oauth_name, "access_token": access_token, @@ -1140,14 +1136,10 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): # Handle case where user has used product outside of web and is now creating an account through web if not user.account_type.is_web_login(): - # We must use the existing user in the session if it matches - # the user we just got by email/oauth. Note that this only applies - # to multi-tenant, due to the overwriting of the user_db - if MULTI_TENANT: - if user.id: - user_by_session = await db_session.get(User, user.id) - if user_by_session: - user = user_by_session + if user.id: + user_by_session = await db_session.get(User, user.id) + if user_by_session: + user = user_by_session # Lock + check + upgrade in one transaction. was_inactive = not user.is_active diff --git a/backend/onyx/utils/timing.py b/backend/onyx/utils/timing.py index 563d7dac1d..0a41529394 100644 --- a/backend/onyx/utils/timing.py +++ b/backend/onyx/utils/timing.py @@ -1,3 +1,4 @@ +import inspect import time from collections.abc import Callable, Generator, Iterator from functools import wraps @@ -40,13 +41,8 @@ def log_function_time( """ def decorator(func: F) -> F: - @wraps(func) - def wrapped_func(*args: Any, **kwargs: Any) -> Any: - # Elapsed time should use monotonic. - start_time = time.monotonic() - result = func(*args, **kwargs) - elapsed_time = time.monotonic() - start_time - elapsed_time_str = f"{elapsed_time:.3f}" + def _log_elapsed(start_time: float, *args: Any, **kwargs: Any) -> None: + elapsed_time_str = f"{time.monotonic() - start_time:.3f}" log_name = func_name or func.__name__ # ty: ignore[unresolved-attribute] args_str = "" if include_args: @@ -63,8 +59,6 @@ def log_function_time( if debug_only: logger.debug(final_log) else: - # These are generally more important logs so the level is a bit - # higher. logger.notice(final_log) if not print_only: @@ -75,6 +69,22 @@ def log_function_time( user_id=str(user.id) if user else "Unknown", ) + if inspect.iscoroutinefunction(func): + + @wraps(func) + async def wrapped_async(*args: Any, **kwargs: Any) -> Any: + start_time = time.monotonic() + result = await func(*args, **kwargs) + _log_elapsed(start_time, *args, **kwargs) + return result + + return cast(F, wrapped_async) + + @wraps(func) + def wrapped_func(*args: Any, **kwargs: Any) -> Any: + start_time = time.monotonic() + result = func(*args, **kwargs) + _log_elapsed(start_time, *args, **kwargs) return result return cast(F, wrapped_func) diff --git a/backend/tests/external_dependency_unit/auth/test_oauth_callback_oidc_expiry_session.py b/backend/tests/external_dependency_unit/auth/test_oauth_callback_oidc_expiry_session.py new file mode 100644 index 0000000000..fc8b5a9fd6 --- /dev/null +++ b/backend/tests/external_dependency_unit/auth/test_oauth_callback_oidc_expiry_session.py @@ -0,0 +1,205 @@ +"""Single-tenant OIDC callback must persist `oidc_expiry` on the existing +OAuth user without blocking on another connection's row lock, and must not +leave a second transaction idle on `"user"`. +""" + +from __future__ import annotations + +import asyncio +import time +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest +from fastapi_users_db_sqlalchemy import SQLAlchemyUserDatabase +from sqlalchemy import delete, text +from sqlalchemy.orm import Session + +import onyx.auth.users as users_module +from onyx.auth.users import UserManager +from onyx.db.engine.async_sql_engine import get_async_session_context_manager +from onyx.db.engine.sql_engine import get_session_with_current_tenant +from onyx.db.models import OAuthAccount, User +from onyx.server.security.store import _build_env_defaults +from tests.external_dependency_unit.conftest import create_test_user, delete_test_user + +_CALLBACK_TIMEOUT_SECONDS = 5 +_POLL_INTERVAL_SECONDS = 0.2 +_OAUTH_NAME = "oidc" + + +def _idle_user_lock_pids(exclude: set[int]) -> list[int]: + """Backends sitting idle while still holding a lock on `"user"`.""" + with get_session_with_current_tenant() as session: + checker_pid = session.execute(text("SELECT pg_backend_pid()")).scalar() + exclude = exclude | {int(checker_pid)} + pids = session.execute( + text( + "SELECT DISTINCT a.pid " + "FROM pg_stat_activity a " + "JOIN pg_locks l ON l.pid = a.pid " + "JOIN pg_class c ON c.oid = l.relation " + "WHERE a.datname = current_database() " + "AND a.state = 'idle in transaction' " + "AND c.relname = 'user'" + ) + ).scalars() + return [int(pid) for pid in pids if int(pid) not in exclude] + + +def _assert_no_idle_user_lock(exclude: set[int]) -> None: + deadline = time.monotonic() + 1.0 + leftover: list[int] = [] + while time.monotonic() < deadline: + leftover = _idle_user_lock_pids(exclude) + if not leftover: + return + time.sleep(0.05) + pytest.fail(f'leftover idle-in-transaction lock on "user"; pids={leftover}') + + +def _attach_oauth_account(db_session: Session, user: User) -> str: + account_id = f"entra-{uuid4().hex}" + db_session.add( + OAuthAccount( + user_id=user.id, + oauth_name=_OAUTH_NAME, + account_id=account_id, + account_email=user.email, + access_token="test-access-token", + refresh_token="", + ) + ) + db_session.commit() + return account_id + + +def _delete_oauth_accounts(db_session: Session, user: User) -> None: + db_session.execute( + delete(OAuthAccount).where(OAuthAccount.__table__.c.user_id == user.id) + ) + + +async def _run_oauth_callback( + *, + user_email: str, + account_id: str, + expires_at: int | None, + track_external_idp_expiry: bool, + monkeypatch: pytest.MonkeyPatch, + exclude_pids: set[int], +) -> User: + monkeypatch.setattr(users_module, "MULTI_TENANT", False) + monkeypatch.setattr( + users_module, + "get_security_settings", + lambda: _build_env_defaults().model_copy( + update={"track_external_idp_expiry": track_external_idp_expiry} + ), + ) + + async with get_async_session_context_manager() as injected_session: + injected_pid = ( + await injected_session.execute(text("SELECT pg_backend_pid()")) + ).scalar() + assert injected_pid is not None + watched_pids = exclude_pids | {int(injected_pid)} + + manager = UserManager( + SQLAlchemyUserDatabase(injected_session, User, OAuthAccount) + ) + task = asyncio.create_task( + manager.oauth_callback( + oauth_name=_OAUTH_NAME, + access_token="rotated-access-token", + account_id=account_id, + account_email=user_email, + expires_at=expires_at, + is_verified_by_default=True, + ) + ) + deadline = time.monotonic() + _CALLBACK_TIMEOUT_SECONDS + idle_pids: list[int] = [] + while not task.done(): + if time.monotonic() >= deadline: + idle_pids = _idle_user_lock_pids(watched_pids) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + pytest.fail( + "oauth_callback blocked instead of writing oidc_expiry; " + f"idle-in-transaction pids={idle_pids}" + ) + await asyncio.sleep(_POLL_INTERVAL_SECONDS) + idle_pids = _idle_user_lock_pids(watched_pids) + return task.result() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("tenant_context") +async def test_existing_oauth_login_writes_oidc_expiry( + db_session: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + user = create_test_user(db_session, "oidc_expiry_write") + account_id = _attach_oauth_account(db_session, user) + expires_at = int(datetime.now(timezone.utc).timestamp()) + 3600 + + sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar() + try: + result = await _run_oauth_callback( + user_email=user.email, + account_id=account_id, + expires_at=expires_at, + track_external_idp_expiry=True, + monkeypatch=monkeypatch, + exclude_pids={int(sync_pid)}, + ) + assert result.id == user.id + + db_session.expire_all() + persisted = db_session.get(User, user.id) + assert persisted is not None + assert persisted.oidc_expiry is not None + assert persisted.oidc_expiry == datetime.fromtimestamp( + expires_at, tz=timezone.utc + ) + sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar() + _assert_no_idle_user_lock({int(sync_pid)}) + finally: + _delete_oauth_accounts(db_session, user) + delete_test_user(db_session, user) + db_session.commit() + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("tenant_context") +async def test_existing_oauth_login_clears_oidc_expiry_when_tracking_disabled( + db_session: Session, monkeypatch: pytest.MonkeyPatch +) -> None: + user = create_test_user(db_session, "oidc_expiry_clear") + user.oidc_expiry = datetime.now(timezone.utc) + db_session.commit() + account_id = _attach_oauth_account(db_session, user) + + sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar() + try: + result = await _run_oauth_callback( + user_email=user.email, + account_id=account_id, + expires_at=int(datetime.now(timezone.utc).timestamp()) + 3600, + track_external_idp_expiry=False, + monkeypatch=monkeypatch, + exclude_pids={int(sync_pid)}, + ) + assert result.id == user.id + assert result.oidc_expiry is None + + db_session.expire_all() + persisted = db_session.get(User, user.id) + assert persisted is not None + assert persisted.oidc_expiry is None + sync_pid = db_session.execute(text("SELECT pg_backend_pid()")).scalar() + _assert_no_idle_user_lock({int(sync_pid)}) + finally: + _delete_oauth_accounts(db_session, user) + delete_test_user(db_session, user) + db_session.commit() diff --git a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py index 315e2be775..ba66b1e4bf 100644 --- a/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py +++ b/backend/tests/integration/mock_services/mcp_test_server/run_mcp_server_per_user_key.py @@ -18,6 +18,7 @@ from starlette.responses import JSONResponse, Response DEFAULT_PORT = 8003 MCP_PATH_PREFIX = "/mcp" +USER_EMAIL_HEADER = "X-User-Email" # ---- pretend database -------------------------------------------------------- @@ -122,7 +123,10 @@ class RequireHeadersMiddleware(BaseHTTPMiddleware): return await call_next(request) for header in self.required_headers: - if not request.headers.get(header): + value = request.headers.get(header) + if header.lower() == USER_EMAIL_HEADER.lower(): + print(f"{USER_EMAIL_HEADER}: {value!r}", flush=True) + if not value: return JSONResponse( {"error": f"Missing required header '{header}'"}, status_code=401, diff --git a/backend/tests/unit/onyx/auth/test_user_registration.py b/backend/tests/unit/onyx/auth/test_user_registration.py index 8cd59d2f13..adc2477a7a 100644 --- a/backend/tests/unit/onyx/auth/test_user_registration.py +++ b/backend/tests/unit/onyx/auth/test_user_registration.py @@ -47,12 +47,16 @@ def mock_async_session() -> MagicMock: session.commit = AsyncMock() session.rollback = AsyncMock() session.run_sync = AsyncMock(return_value=None) + session.get = AsyncMock(return_value=None) return session class _AsyncSessionContextManager: def __init__(self, session: MagicMock) -> None: self._session = session + # oauth_callback awaits session.get on the placeholder upgrade path. + if not isinstance(session.get, AsyncMock): + session.get = AsyncMock(return_value=None) async def __aenter__(self) -> MagicMock: return self._session @@ -79,6 +83,14 @@ def _mock_user_manager_methods(user_manager: UserManager) -> None: user_manager.validate_password = AsyncMock() +def _bind_sqlalchemy_user_db_cls( + mock_user_db_cls: MagicMock, mock_user_db: MagicMock +) -> None: + # SQLAlchemyUserDatabase[User, uuid.UUID](...) hits __getitem__, then call. + mock_user_db_cls.return_value = mock_user_db + mock_user_db_cls.__getitem__.return_value = mock_user_db_cls + + class TestDisposableEmailValidation: """Test disposable email validation before tenant provisioning.""" @@ -580,8 +592,10 @@ class TestOAuthDottedGmail: @patch("onyx.auth.users.fetch_ee_implementation_or_noop") @patch("onyx.auth.users.get_async_session_context_manager") @patch("onyx.auth.users.remove_user_from_invited_users") + @patch("onyx.auth.users.SQLAlchemyUserDatabase") async def test_oauth_create_does_not_block_dotted_gmail( self, + mock_user_db_cls: MagicMock, mock_remove_invited: MagicMock, # noqa: ARG002 mock_session_manager: MagicMock, mock_fetch_ee: MagicMock, @@ -617,6 +631,7 @@ class TestOAuthDottedGmail: mock_user_db.session = MagicMock() mock_user_db.session.run_sync = AsyncMock() user_manager.user_db = mock_user_db + _bind_sqlalchemy_user_db_cls(mock_user_db_cls, mock_user_db) await user_manager.oauth_callback( oauth_name="google", @@ -661,7 +676,9 @@ class TestOAuthNoAutoLinkExemptions: ) @staticmethod - def _manager_with_existing(existing_user: MagicMock) -> UserManager: + def _manager_with_existing( + existing_user: MagicMock, mock_user_db_cls: MagicMock | None = None + ) -> UserManager: user_manager = UserManager(MagicMock()) _mock_user_manager_methods(user_manager) user_manager.on_after_register = AsyncMock() @@ -675,6 +692,8 @@ class TestOAuthNoAutoLinkExemptions: mock_user_db.update = AsyncMock(return_value=existing_user) mock_user_db.session = MagicMock() user_manager.user_db = mock_user_db + if mock_user_db_cls is not None: + _bind_sqlalchemy_user_db_cls(mock_user_db_cls, mock_user_db) return user_manager @pytest.mark.asyncio @@ -690,8 +709,10 @@ class TestOAuthNoAutoLinkExemptions: @patch("onyx.auth.users.assign_user_to_default_groups__no_commit") @patch("onyx.auth.users._upgrade_will_add_seat", return_value=False) @patch("onyx.auth.users.get_session_with_current_tenant") + @patch("onyx.auth.users.SQLAlchemyUserDatabase") async def test_placeholder_promoted_without_auto_link( self, + mock_user_db_cls: MagicMock, mock_sync_session_factory: MagicMock, mock_will_add_seat: MagicMock, # noqa: ARG002 mock_assign_groups: MagicMock, @@ -717,7 +738,8 @@ class TestOAuthNoAutoLinkExemptions: account_type=AccountType.EXT_PERM_USER, is_active=is_active, ) - user_manager = self._manager_with_existing(placeholder) + user_manager = self._manager_with_existing(placeholder, mock_user_db_cls) + mock_async_session.get = AsyncMock(return_value=placeholder) sync_user = MagicMock(is_active=is_active) mock_sync_db = MagicMock() @@ -757,8 +779,10 @@ class TestOAuthNoAutoLinkExemptions: @patch("onyx.auth.users.fetch_ee_implementation_or_noop") @patch("onyx.auth.users.get_async_session_context_manager") @patch("onyx.auth.users.remove_user_from_invited_users") + @patch("onyx.auth.users.SQLAlchemyUserDatabase") async def test_unclaimed_row_is_claimed_by_first_login( self, + mock_user_db_cls: MagicMock, mock_remove_invited: MagicMock, # noqa: ARG002 mock_session_manager: MagicMock, mock_fetch_ee: MagicMock, @@ -774,7 +798,7 @@ class TestOAuthNoAutoLinkExemptions: mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant") provisioned = self._unclaimed() - user_manager = self._manager_with_existing(provisioned) + user_manager = self._manager_with_existing(provisioned, mock_user_db_cls) result = await user_manager.oauth_callback( oauth_name="okta", @@ -804,8 +828,10 @@ class TestOAuthNoAutoLinkExemptions: @patch("onyx.auth.users.fetch_ee_implementation_or_noop") @patch("onyx.auth.users.get_async_session_context_manager") @patch("onyx.auth.users.remove_user_from_invited_users") + @patch("onyx.auth.users.SQLAlchemyUserDatabase") async def test_spoken_for_row_is_rejected( self, + mock_user_db_cls: MagicMock, mock_remove_invited: MagicMock, # noqa: ARG002 mock_session_manager: MagicMock, mock_fetch_ee: MagicMock, @@ -819,7 +845,9 @@ class TestOAuthNoAutoLinkExemptions: ) mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant") - user_manager = self._manager_with_existing(self._unclaimed(**overrides)) + user_manager = self._manager_with_existing( + self._unclaimed(**overrides), mock_user_db_cls + ) with pytest.raises(exceptions.UserAlreadyExists): await user_manager.oauth_callback( @@ -839,8 +867,10 @@ class TestOAuthNoAutoLinkExemptions: @patch("onyx.auth.users.fetch_ee_implementation_or_noop") @patch("onyx.auth.users.get_async_session_context_manager") @patch("onyx.auth.users.remove_user_from_invited_users") + @patch("onyx.auth.users.SQLAlchemyUserDatabase") async def test_deactivated_row_is_not_linked( self, + mock_user_db_cls: MagicMock, mock_remove_invited: MagicMock, # noqa: ARG002 mock_session_manager: MagicMock, mock_fetch_ee: MagicMock, @@ -856,7 +886,7 @@ class TestOAuthNoAutoLinkExemptions: mock_fetch_ee.return_value = AsyncMock(return_value="test_tenant") deactivated = self._unclaimed(is_active=False) - user_manager = self._manager_with_existing(deactivated) + user_manager = self._manager_with_existing(deactivated, mock_user_db_cls) result = await user_manager.oauth_callback( oauth_name="okta",