diff --git a/backend/ee/onyx/db/license.py b/backend/ee/onyx/db/license.py index 03fe523d00..e569aea527 100644 --- a/backend/ee/onyx/db/license.py +++ b/backend/ee/onyx/db/license.py @@ -156,7 +156,7 @@ def get_used_seats(tenant_id: str | None = None) -> int: Per-user predicate ``user_counts_toward_seats`` mirrors this filter. """ if MULTI_TENANT: - from ee.onyx.server.tenants.user_mapping import get_tenant_count + from ee.onyx.db.user_tenant_mapping import get_tenant_count return get_tenant_count(tenant_id or get_current_tenant_id()) else: diff --git a/backend/ee/onyx/server/tenants/user_mapping.py b/backend/ee/onyx/db/user_tenant_mapping.py similarity index 82% rename from backend/ee/onyx/server/tenants/user_mapping.py rename to backend/ee/onyx/db/user_tenant_mapping.py index 11722bb6db..2deffdb120 100644 --- a/backend/ee/onyx/server/tenants/user_mapping.py +++ b/backend/ee/onyx/db/user_tenant_mapping.py @@ -1,5 +1,6 @@ from fastapi_users import exceptions from sqlalchemy import select +from sqlalchemy.dialects.postgresql import insert as pg_insert from onyx.auth.invited_users import ( get_invited_users, @@ -8,10 +9,10 @@ from onyx.auth.invited_users import ( write_pending_users, ) from onyx.db.engine.sql_engine import ( - get_session_with_shared_schema, + get_catalog_session, get_session_with_tenant, ) -from onyx.db.models import UserTenantMapping +from onyx.db.models import UserTenantMapping, UserTenantMappingOAuthAccount from onyx.server.manage.models import TenantSnapshot from onyx.utils.logger import setup_logger from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA @@ -25,7 +26,7 @@ def get_tenant_id_for_email(email: str) -> str: return POSTGRES_DEFAULT_SCHEMA # Implement logic to get tenant_id from the mapping table try: - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: # First try to get an active tenant result = db_session.execute( select(UserTenantMapping).where( @@ -60,7 +61,8 @@ def get_tenant_id_for_email(email: str) -> str: def user_owns_a_tenant(email: str) -> bool: - with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session: + email = email.lower() + with get_catalog_session() as db_session: result = ( db_session.query(UserTenantMapping) .filter(UserTenantMapping.email == email) @@ -69,6 +71,64 @@ def user_owns_a_tenant(email: str) -> bool: return result is not None +def record_oauth_identity( + email: str, tenant_id: str, oauth_name: str, account_id: str +) -> None: + """Link the IdP subject to this user's mapping row so resolution survives + an address change at the provider. + + A subject links to one mapping, and the first linked row wins: a later + login whose address maps elsewhere never re-links it, so address + reassignment cannot steal a linked identity. + + No-op outside multi-tenant, where the mapping tables are not provisioned. + """ + if not MULTI_TENANT: + return + + normalized_email = email.lower() + with get_catalog_session() as db_session: + if db_session.get(UserTenantMapping, (normalized_email, tenant_id)) is None: + logger.info("No mapping row to link in tenant %s", tenant_id) + return + + inserted = db_session.execute( + pg_insert(UserTenantMappingOAuthAccount) + .values( + oauth_name=oauth_name, + account_id=account_id, + email=normalized_email, + tenant_id=tenant_id, + ) + .on_conflict_do_nothing( + index_elements=[ + UserTenantMappingOAuthAccount.oauth_name, + UserTenantMappingOAuthAccount.account_id, + ] + ) + .returning( + UserTenantMappingOAuthAccount.email, + UserTenantMappingOAuthAccount.tenant_id, + ) + ).one_or_none() + if inserted is None: + owner = db_session.execute( + select( + UserTenantMappingOAuthAccount.email, + UserTenantMappingOAuthAccount.tenant_id, + ).where( + UserTenantMappingOAuthAccount.oauth_name == oauth_name, + UserTenantMappingOAuthAccount.account_id == account_id, + ) + ).one_or_none() + if owner is None or tuple(owner) != (normalized_email, tenant_id): + logger.warning( + "OAuth identity is already linked to another tenant mapping" + ) + return + db_session.commit() + + def add_users_to_tenant(emails: list[str], tenant_id: str) -> None: """ Add users to a tenant. If a user has an active mapping elsewhere, @@ -81,7 +141,7 @@ def add_users_to_tenant(emails: list[str], tenant_id: str) -> None: if not unique_emails: return - with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session: + with get_catalog_session() as db_session: try: # Start a transaction db_session.begin() @@ -122,7 +182,7 @@ def add_users_to_tenant(emails: list[str], tenant_id: str) -> None: if new_active_seat_emails: from ee.onyx.server.tenants.billing import enforce_cloud_seat_limit - # Lock + bill held across the inserts below; rolled back + # Lock + bill held across the inserts below. Rolled back # on Stripe decline by the outer ``except Exception``. enforce_cloud_seat_limit( seats_needed=len(new_active_seat_emails), @@ -156,7 +216,7 @@ def add_users_to_tenant(emails: list[str], tenant_id: str) -> None: def remove_users_from_tenant(emails: list[str], tenant_id: str) -> None: - with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session: + with get_catalog_session() as db_session: try: mappings_to_delete = ( db_session.query(UserTenantMapping) @@ -179,30 +239,19 @@ def remove_users_from_tenant(emails: list[str], tenant_id: str) -> None: def remove_all_users_from_tenant(tenant_id: str) -> None: - with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as db_session: + with get_catalog_session() as db_session: db_session.query(UserTenantMapping).filter( UserTenantMapping.tenant_id == tenant_id ).delete() db_session.commit() -def invite_self_to_tenant(email: str, tenant_id: str) -> None: - token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) - try: - pending_users = get_pending_users() - if email in pending_users: - return - write_pending_users(pending_users + [email]) - finally: - CURRENT_TENANT_ID_CONTEXTVAR.reset(token) - - def approve_user_invite(email: str, tenant_id: str) -> None: """ Approve a user invite to a tenant. This will delete all existing records for this email and create a new mapping entry for the user in this tenant. """ - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: # Delete all existing records for this email db_session.query(UserTenantMapping).filter( UserTenantMapping.email == email @@ -232,7 +281,7 @@ def accept_user_invite(email: str, tenant_id: str) -> None: Accept an invitation to join a tenant. This activates the user's mapping to the tenant. """ - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: try: # Lock the user's mappings first to prevent race conditions. # This ensures no concurrent request can modify this user's mappings. @@ -311,7 +360,7 @@ def deny_user_invite(email: str, tenant_id: str) -> None: Deny an invitation to join a tenant. This removes the user's mapping to the tenant. """ - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: # Delete the mapping for this user and tenant result = ( db_session.query(UserTenantMapping) @@ -357,7 +406,7 @@ def get_tenant_count(tenant_id: str) -> int: from onyx.db.models import User # First get all emails with active mappings to this tenant - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: active_mapping_emails = ( db_session.query(UserTenantMapping.email) .filter( @@ -391,7 +440,7 @@ def get_tenant_invitation(email: str) -> TenantSnapshot | None: """ Get the first tenant invitation for this user """ - with get_session_with_shared_schema() as db_session: + with get_catalog_session() as db_session: # Get the first tenant invitation for this user invitation = ( db_session.query(UserTenantMapping) diff --git a/backend/ee/onyx/server/tenants/admin_api.py b/backend/ee/onyx/server/tenants/admin_api.py index 1a830db6c7..2e228b5cdd 100644 --- a/backend/ee/onyx/server/tenants/admin_api.py +++ b/backend/ee/onyx/server/tenants/admin_api.py @@ -2,8 +2,8 @@ from fastapi import APIRouter, Depends, HTTPException, Response from fastapi_users import exceptions from ee.onyx.auth.users import current_cloud_superuser +from ee.onyx.db.user_tenant_mapping import get_tenant_id_for_email from ee.onyx.server.tenants.models import ImpersonateRequest -from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email from onyx.auth.users import User, auth_backend, get_redis_strategy from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME from onyx.db.engine.sql_engine import get_session_with_tenant diff --git a/backend/ee/onyx/server/tenants/billing.py b/backend/ee/onyx/server/tenants/billing.py index 7432f331a9..895d08f9a2 100644 --- a/backend/ee/onyx/server/tenants/billing.py +++ b/backend/ee/onyx/server/tenants/billing.py @@ -263,7 +263,7 @@ def enforce_cloud_seat_limit( # Local import avoids circular import with user_mapping (which calls # back into this module from add_users_to_tenant). - from ee.onyx.server.tenants.user_mapping import get_tenant_count + from ee.onyx.db.user_tenant_mapping import get_tenant_count if db_session is not None: acquire_seat_lock(db_session, tenant) diff --git a/backend/ee/onyx/server/tenants/provisioning.py b/backend/ee/onyx/server/tenants/provisioning.py index 7ae23f2279..a80b6cd532 100644 --- a/backend/ee/onyx/server/tenants/provisioning.py +++ b/backend/ee/onyx/server/tenants/provisioning.py @@ -9,6 +9,11 @@ from sqlalchemy import select from sqlalchemy.orm import Session from ee.onyx.configs.app_configs import HUBSPOT_TRACKING_URL +from ee.onyx.db.user_tenant_mapping import ( + add_users_to_tenant, + get_tenant_id_for_email, + user_owns_a_tenant, +) from ee.onyx.server.tenants.access import generate_data_plane_token from ee.onyx.server.tenants.models import ( TenantByDomainResponse, @@ -20,11 +25,6 @@ from ee.onyx.server.tenants.schema_management import ( drop_schema, run_alembic_migrations, ) -from ee.onyx.server.tenants.user_mapping import ( - add_users_to_tenant, - get_tenant_id_for_email, - user_owns_a_tenant, -) from onyx.auth.users import exceptions from onyx.configs.app_configs import ( ANTHROPIC_DEFAULT_API_KEY, diff --git a/backend/ee/onyx/server/tenants/team_membership_api.py b/backend/ee/onyx/server/tenants/team_membership_api.py index e8d89ffe3c..d1d06a7764 100644 --- a/backend/ee/onyx/server/tenants/team_membership_api.py +++ b/backend/ee/onyx/server/tenants/team_membership_api.py @@ -1,11 +1,11 @@ from fastapi import APIRouter, Depends, HTTPException from sqlalchemy.orm import Session -from ee.onyx.server.tenants.provisioning import delete_user_from_control_plane -from ee.onyx.server.tenants.user_mapping import ( +from ee.onyx.db.user_tenant_mapping import ( remove_all_users_from_tenant, remove_users_from_tenant, ) +from ee.onyx.server.tenants.provisioning import delete_user_from_control_plane from onyx.auth.permissions import require_permission from onyx.auth.users import User from onyx.db.auth import get_user_count diff --git a/backend/ee/onyx/server/tenants/user_invitations_api.py b/backend/ee/onyx/server/tenants/user_invitations_api.py index 41e24e7df4..02a53aa996 100644 --- a/backend/ee/onyx/server/tenants/user_invitations_api.py +++ b/backend/ee/onyx/server/tenants/user_invitations_api.py @@ -1,28 +1,45 @@ from fastapi import APIRouter, Depends, HTTPException +from ee.onyx.db.user_tenant_mapping import ( + accept_user_invite, + approve_user_invite, + deny_user_invite, +) from ee.onyx.server.tenants.models import ( ApproveUserRequest, PendingUserSnapshot, RequestInviteRequest, ) -from ee.onyx.server.tenants.user_mapping import ( - accept_user_invite, - approve_user_invite, - deny_user_invite, - invite_self_to_tenant, +from onyx.auth.invited_users import ( + get_pending_users, + write_pending_users, ) -from onyx.auth.invited_users import get_pending_users from onyx.auth.permissions import require_permission from onyx.auth.users import User from onyx.db.enums import Permission from onyx.utils.logger import setup_logger -from shared_configs.contextvars import get_current_tenant_id +from shared_configs.contextvars import ( + CURRENT_TENANT_ID_CONTEXTVAR, + get_current_tenant_id, +) logger = setup_logger() router = APIRouter(prefix="/tenants") +def invite_self_to_tenant(email: str, tenant_id: str) -> None: + # The pending list lives in the target tenant's KV store, not the caller's. + token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id) + try: + pending_users = get_pending_users() + if email in pending_users: + return + write_pending_users(pending_users + [email]) + finally: + CURRENT_TENANT_ID_CONTEXTVAR.reset(token) + + @router.post("/users/invite/request") async def request_invite( invite_request: RequestInviteRequest, diff --git a/backend/onyx/auth/login_claims_capture.py b/backend/onyx/auth/login_claims_capture.py index 92e2fa704e..02a7acf24c 100644 --- a/backend/onyx/auth/login_claims_capture.py +++ b/backend/onyx/auth/login_claims_capture.py @@ -92,7 +92,7 @@ async def _resolve_capture_tenant_id(email: str) -> str | None: try: tenant_id = await asyncio.to_thread( fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None + "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None ), email, ) diff --git a/backend/onyx/auth/users.py b/backend/onyx/auth/users.py index 2cc14005e6..01319af64d 100644 --- a/backend/onyx/auth/users.py +++ b/backend/onyx/auth/users.py @@ -481,7 +481,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): async def get_by_email(self, user_email: str) -> User: tenant_id = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None + "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None )(user_email) async with get_async_session_context_manager(tenant_id) as db_session: if MULTI_TENANT: @@ -524,7 +524,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): try: tenant_id = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", + "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None, )(email) @@ -1012,6 +1012,12 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]): oauth_account_dict, ) + # Keyed on the stored email rather than the one the IdP just sent: + # that is the address this user's membership row is filed under. + fetch_ee_implementation_or_noop( + "onyx.db.user_tenant_mapping", "record_oauth_identity", None + )(user.email, tenant_id, oauth_name, account_id) + # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to # re-authenticate that frequently, so by default this is disabled track_external_idp_expiry = ( @@ -2468,7 +2474,7 @@ async def complete_login_flow( referral_source = state_data.get("referral_source", None) try: tenant_id = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None + "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None )(account_email) except exceptions.UserNotExists: tenant_id = None diff --git a/backend/onyx/server/manage/users.py b/backend/onyx/server/manage/users.py index 4bad22d3fd..611e232771 100644 --- a/backend/onyx/server/manage/users.py +++ b/backend/onyx/server/manage/users.py @@ -653,7 +653,7 @@ def bulk_invite_users( try: write_invited_users(initial_invited_users) # Reset to original state fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "remove_users_from_tenant", None + "onyx.db.user_tenant_mapping", "remove_users_from_tenant", None )(new_invited_emails, tenant_id) finally: # Release the counter reservation regardless of whether the KV / @@ -710,7 +710,7 @@ def remove_invited_user( ) if MULTI_TENANT: fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "remove_users_from_tenant", None + "onyx.db.user_tenant_mapping", "remove_users_from_tenant", None )([user_email.user_email], tenant_id) number_of_invited_users = remove_user_from_invited_users(user_email.user_email) @@ -799,7 +799,7 @@ async def delete_user( try: tenant_id = get_current_tenant_id() fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "remove_users_from_tenant", None + "onyx.db.user_tenant_mapping", "remove_users_from_tenant", None )([user_email.user_email], tenant_id) delete_user_from_db(user_to_delete, db_session) logger.info("Deleted user %s", user_to_delete.email) @@ -1060,20 +1060,20 @@ def verify_user_logged_in( team_name = tenant_id else: team_name = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None + "onyx.db.user_tenant_mapping", "get_tenant_id_for_email", None )(user.email) if MULTI_TENANT: if team_name != tenant_id: user_count = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_count", None + "onyx.db.user_tenant_mapping", "get_tenant_count", None )(team_name) new_tenant = TenantSnapshot( tenant_id=team_name, number_of_users=user_count ) tenant_invitation = fetch_ee_implementation_or_noop( - "onyx.server.tenants.user_mapping", "get_tenant_invitation", None + "onyx.db.user_tenant_mapping", "get_tenant_invitation", None )(user.email) super_users_list = cast( diff --git a/backend/onyx/utils/telemetry.py b/backend/onyx/utils/telemetry.py index a13b1224e4..43404426e0 100644 --- a/backend/onyx/utils/telemetry.py +++ b/backend/onyx/utils/telemetry.py @@ -189,7 +189,7 @@ def mt_cloud_telemetry( def _get_tenant_id_for_user_identify(user_email: str) -> str | None: try: return fetch_versioned_implementation_with_fallback( - module="onyx.server.tenants.user_mapping", + module="onyx.db.user_tenant_mapping", attribute="get_tenant_id_for_email", fallback=lambda _email: POSTGRES_DEFAULT_SCHEMA, )(user_email) diff --git a/backend/scripts/debugging/onyx_redis.py b/backend/scripts/debugging/onyx_redis.py index 75a30ae9ca..e9f4461e65 100644 --- a/backend/scripts/debugging/onyx_redis.py +++ b/backend/scripts/debugging/onyx_redis.py @@ -10,7 +10,7 @@ from uuid import UUID from redis import Redis -from ee.onyx.server.tenants.user_mapping import get_tenant_id_for_email +from ee.onyx.db.user_tenant_mapping import get_tenant_id_for_email from onyx.auth.invited_users import get_invited_users, write_invited_users from onyx.configs.app_configs import ( REDIS_AUTH_KEY_PREFIX, diff --git a/backend/tests/unit/ee/onyx/db/test_license.py b/backend/tests/unit/ee/onyx/db/test_license.py index 773af59847..521cf9f8db 100644 --- a/backend/tests/unit/ee/onyx/db/test_license.py +++ b/backend/tests/unit/ee/onyx/db/test_license.py @@ -178,7 +178,7 @@ class TestCheckSeatAvailabilityMultiTenant: @patch("ee.onyx.db.license.MULTI_TENANT", True) @patch( - "ee.onyx.server.tenants.user_mapping.get_tenant_count", + "ee.onyx.db.user_tenant_mapping.get_tenant_count", return_value=5, ) @patch("ee.onyx.db.license.get_license_metadata") @@ -196,7 +196,7 @@ class TestCheckSeatAvailabilityMultiTenant: @patch("ee.onyx.db.license.MULTI_TENANT", True) @patch( - "ee.onyx.server.tenants.user_mapping.get_tenant_count", + "ee.onyx.db.user_tenant_mapping.get_tenant_count", return_value=10, ) @patch("ee.onyx.db.license.get_license_metadata") diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py b/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py index c27e0a3bb3..66084407e0 100644 --- a/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py +++ b/backend/tests/unit/ee/onyx/server/tenants/test_billing_seat_enforcement.py @@ -163,7 +163,7 @@ class TestEnforceCloudSeatLimit: enforce_cloud_seat_limit(seats_needed=1) @patch(f"{_BILLING}.acquire_seat_lock") - @patch("ee.onyx.server.tenants.user_mapping.get_tenant_count") + @patch("ee.onyx.db.user_tenant_mapping.get_tenant_count") @patch(f"{_BILLING}.attempt_seat_billing_increase") @patch(f"{_BILLING}.is_tenant_on_trial_fn") @patch(f"{_BILLING}.get_current_tenant_id") @@ -194,7 +194,7 @@ class TestEnforceCloudSeatLimit: ], ) @patch(f"{_BILLING}.acquire_seat_lock") - @patch("ee.onyx.server.tenants.user_mapping.get_tenant_count") + @patch("ee.onyx.db.user_tenant_mapping.get_tenant_count") @patch(f"{_BILLING}.attempt_seat_billing_increase") @patch(f"{_BILLING}.is_tenant_on_trial_fn") @patch(f"{_BILLING}.get_current_tenant_id") @@ -221,7 +221,7 @@ class TestEnforceCloudSeatLimit: assert expected_in_message in exc.value.detail @patch(f"{_BILLING}.acquire_seat_lock") - @patch("ee.onyx.server.tenants.user_mapping.get_tenant_count") + @patch("ee.onyx.db.user_tenant_mapping.get_tenant_count") @patch(f"{_BILLING}.attempt_seat_billing_increase") @patch(f"{_BILLING}.is_tenant_on_trial_fn") @patch(f"{_BILLING}.get_current_tenant_id", return_value="ctx_var_tenant") diff --git a/backend/tests/unit/ee/onyx/server/tenants/test_record_oauth_identity.py b/backend/tests/unit/ee/onyx/server/tenants/test_record_oauth_identity.py new file mode 100644 index 0000000000..f400f929dd --- /dev/null +++ b/backend/tests/unit/ee/onyx/server/tenants/test_record_oauth_identity.py @@ -0,0 +1,105 @@ +"""Guards that `record_oauth_identity` only touches the mapping tables on cloud +and never re-links a subject that already belongs to another mapping. + +The `public.user_tenant_mapping*` tables are created by the `alembic_tenants` +tree, which only multi-tenant deployments run. Opening a session against them +anywhere else would raise on every OAuth login, so the `MULTI_TENANT` guard is +load-bearing rather than an optimization. +""" + +from unittest.mock import MagicMock, patch + +from ee.onyx.db.user_tenant_mapping import record_oauth_identity + +_MAPPING_MODULE = "ee.onyx.db.user_tenant_mapping" + + +def _run( + *, + multi_tenant: bool, + mapping_exists: bool = True, + insert_result: tuple[str, str] | None = ("user@example.com", "tenant_abc"), + owner: tuple[str, str] | None = None, +) -> MagicMock: + with ( + patch(f"{_MAPPING_MODULE}.MULTI_TENANT", multi_tenant), + patch(f"{_MAPPING_MODULE}.get_catalog_session") as session_ctx, + ): + db_session = session_ctx.return_value.__enter__.return_value + db_session.get.return_value = MagicMock() if mapping_exists else None + insert_execute = MagicMock() + insert_execute.one_or_none.return_value = insert_result + owner_execute = MagicMock() + owner_execute.one_or_none.return_value = owner + db_session.execute.side_effect = [insert_execute, owner_execute] + record_oauth_identity( + email="User@example.com", + tenant_id="tenant_abc", + oauth_name="google", + account_id="sub-123", + ) + return session_ctx + + +def _compiled_insert_sql(session_ctx: MagicMock) -> str: + from sqlalchemy.dialects import postgresql + + db_session = session_ctx.return_value.__enter__.return_value + statement = db_session.execute.call_args_list[0].args[0] + return str( + statement.compile( + dialect=postgresql.dialect(), compile_kwargs={"literal_binds": True} + ) + ) + + +def test_single_tenant_opens_no_session() -> None: + assert _run(multi_tenant=False).call_count == 0 + + +def test_multi_tenant_links_the_subject_to_the_mapping() -> None: + session_ctx = _run(multi_tenant=True) + assert session_ctx.call_count == 1 + + db_session = session_ctx.return_value.__enter__.return_value + assert db_session.execute.call_count == 1 + db_session.commit.assert_called_once() + + sql = _compiled_insert_sql(session_ctx) + assert "user_tenant_mapping_oauth_account" in sql + assert "ON CONFLICT" in sql + assert "DO NOTHING" in sql + assert "'user@example.com'" in sql + + +def test_missing_mapping_row_links_nothing() -> None: + session_ctx = _run(multi_tenant=True, mapping_exists=False) + + db_session = session_ctx.return_value.__enter__.return_value + db_session.execute.assert_not_called() + db_session.commit.assert_not_called() + + +def test_subject_linked_elsewhere_is_not_re_linked() -> None: + """Address reassignment: the new holder of an address must not inherit a + subject link that still belongs to the renamed user's mapping.""" + session_ctx = _run( + multi_tenant=True, + insert_result=None, + owner=("other@example.com", "tenant_other"), + ) + + db_session = session_ctx.return_value.__enter__.return_value + assert db_session.execute.call_count == 2 + db_session.commit.assert_not_called() + + +def test_repeat_login_for_the_linked_mapping_is_idempotent() -> None: + session_ctx = _run( + multi_tenant=True, + insert_result=None, + owner=("user@example.com", "tenant_abc"), + ) + + db_session = session_ctx.return_value.__enter__.return_value + db_session.commit.assert_called_once() diff --git a/backend/tests/unit/onyx/utils/test_telemetry.py b/backend/tests/unit/onyx/utils/test_telemetry.py index 204c29b600..d4c76b778a 100644 --- a/backend/tests/unit/onyx/utils/test_telemetry.py +++ b/backend/tests/unit/onyx/utils/test_telemetry.py @@ -22,7 +22,7 @@ class CloudIdentifyHarness: self, module: str, attribute: str, fallback: Any ) -> Any: if ( - module == "onyx.server.tenants.user_mapping" + module == "onyx.db.user_tenant_mapping" and attribute == "get_tenant_id_for_email" ): return self.get_tenant_id_for_email