refactor(tenancy): answer which organization owns a scope, for other domains

src/gateway/repositories/tenancy/organization_member_repository.py

@@ -147,6 +147,20 @@ async def get_first_active_for_user(self, user_id: uuid.UUID) -> OrganizationMem
         )
         return result.scalars().first()
 
+    async def get_organization_id(self, organization_member_id: uuid.UUID) -> uuid.UUID | None:
+        """Return the ID of the organization a membership belongs to, or None."""
+        result = await self.db.execute(
+            select(col(OrganizationMember.organization_id)).where(col(OrganizationMember.id) == organization_member_id)
+        )
+        return result.scalar_one_or_none()
+
+    async def get_ids_by_organization(self, organization_id: uuid.UUID) -> list[uuid.UUID]:
+        """Return the ID of every membership in an organization, whatever its status."""
+        result = await self.db.execute(
+            select(col(OrganizationMember.id)).where(col(OrganizationMember.organization_id) == organization_id)
+        )
+        return list(result.scalars().all())
+
     async def create_membership(
         self,
         *,

src/gateway/repositories/tenancy/workspace_repository.py

@@ -101,6 +101,13 @@ async def get_ids_by_organization(self, organization_id: uuid.UUID) -> list[uuid
         )
         return list(result.scalars().all())
 
+    async def get_organization_id(self, workspace_id: uuid.UUID) -> uuid.UUID | None:
+        """Return the ID of the organization that owns a workspace, or None."""
+        result = await self.db.execute(
+            select(col(Workspace.organization_id)).where(col(Workspace.id) == workspace_id)
+        )
+        return result.scalar_one_or_none()
+
     async def get_by_organization_and_name(self, organization_id: uuid.UUID, name: str) -> Workspace | None:
         """Return an organization's workspace with this name, or None."""
         result = await self.db.execute(
@@ -146,11 +153,7 @@ async def delete_workspace(self, workspace: Workspace) -> None:
 
 
 class WorkspaceMemberRepository:
-    """Repository for workspace membership rows.
-
-    Not a ``BaseRepository``: every access is keyed by the (workspace, user)
-    pair rather than by the row's own id, so none of the generic helpers apply.
-    """
+    """Repository for workspace membership rows."""
 
     def __init__(self, db: AsyncSession):
         self.db = db
@@ -307,6 +310,22 @@ async def ids_for_workspace(self, workspace_id: uuid.UUID) -> list[uuid.UUID]:
         )
         return list(result.scalars().all())
 
+    async def get_workspace_id(self, workspace_member_id: uuid.UUID) -> uuid.UUID | None:
+        """Return the ID of the workspace a membership belongs to, or None."""
+        result = await self.db.execute(
+            select(col(WorkspaceMember.workspace_id)).where(col(WorkspaceMember.id) == workspace_member_id)
+        )
+        return result.scalar_one_or_none()
+
+    async def get_ids_by_organization(self, organization_id: uuid.UUID) -> list[uuid.UUID]:
+        """Return the ID of every membership in an organization's workspaces, whatever its status."""
+        result = await self.db.execute(
+            select(col(WorkspaceMember.id))
+            .join(Workspace, col(Workspace.id) == col(WorkspaceMember.workspace_id))
+            .where(col(Workspace.organization_id) == organization_id)
+        )
+        return list(result.scalars().all())
+
     async def create(
         self,
         *,

src/gateway/services/tenancy/organization_service.py

@@ -1,35 +1,12 @@
-"""Organizations: active-organization resolution, CRUD, and membership.
-
-Rehomed from the platform's ``OrganizationService`` plus the membership half of
-``OrganizationMembershipService``, converted to async. The authorization rules,
-the membership constraints, and the response shapes are the platform's; what is
-gone is the depth that has no home in the OSS base yet: mixpanel tracking,
-managed provider-key and default-gateway provisioning, email-domain auto-join,
-teams, and the org's budget and pricing surfaces. Those arrive with their own
-slices, tracked under mozilla-ai/otari-ai#1452, and this service is where they
-attach. Emailed invitations shipped here in mozilla-ai/otari#641 (see
-``invite_active_organization_member_for_user``/``accept_invitation`` below);
-``create_active_organization_member_for_user`` is the older, still-supported
-immediate path this replaced no part of.
-
-One rule runs through every method: a caller only ever acts inside the
-organization their identity is currently pointed at, and every method but one
-resolves that organization from the caller alone rather than from the request.
-The exception is ``switch_active_organization_for_user``, which is the method
-that *moves* the pointer and so has to be told where to; it answers not-found
-for an id the caller holds no active membership in, so naming another tenant's
-organization tells the caller nothing about it.
-
-A standalone deployment still *boots* one organization, provisioned at first
-boot, and that is the shape almost every deployment keeps. But a second one is
-reachable (accept an invitation into an organization elsewhere on the same
-deployment and you hold two memberships), so creating one, listing the ones you
-belong to, and switching between them are part of this surface rather than an
-overlay's: the tables are here, the invariants that decide who becomes owner
-and what happens to ``active_organization_id`` are here, and an overlay that
-contributes no tables could only fork them (mozilla-ai/otari#715). Deleting an
-organization is still absent, which is a separate question: the rows every
-historical attribution resolves through hang off it.
+"""This module resolves a caller's active organization and manages organizations and their members.
+
+A method that acts for a caller never trusts an organization that the request names.
+It acts in the caller's active organization, or in one it reaches through the caller's own membership.
+An ID that names another tenant's row answers not-found, so a caller cannot probe which IDs exist.
+
+A deployment can hold more than one organization.
+Creating, listing and switching organizations therefore belong here rather than in an overlay.
+The service offers no way to delete an organization, because historical attribution resolves through its rows.
 """
 
 import hashlib
@@ -268,6 +245,37 @@ async def user_has_active_membership(self, *, organization_id: uuid.UUID, user_i
         """Whether an identity is an active member of an organization."""
         return await self.members.get_active_by_organization_and_user(organization_id, user_id) is not None
 
+    async def has_organization(self, organization_id: uuid.UUID) -> bool:
+        """Return whether an organization with this ID exists."""
+        return await self.organizations.get(organization_id) is not None
+
+    async def get_organization_id_for_workspace(self, workspace_id: uuid.UUID) -> uuid.UUID | None:
+        """Return the ID of the organization that owns a workspace, or None when the workspace does not exist."""
+        return await self.workspace_rows.get_organization_id(workspace_id)
+
+    async def get_organization_id_for_organization_member(
+        self,
+        organization_member_id: uuid.UUID,
+    ) -> uuid.UUID | None:
+        """Return the ID of the organization a membership belongs to, or None when the membership does not exist."""
+        return await self.members.get_organization_id(organization_member_id)
+
+    async def get_workspace_id_for_workspace_member(self, workspace_member_id: uuid.UUID) -> uuid.UUID | None:
+        """Return the ID of the workspace a membership belongs to, or None when the membership does not exist."""
+        return await self.workspaces.get_workspace_id(workspace_member_id)
+
+    async def get_workspace_ids_in_organization(self, organization_id: uuid.UUID) -> list[uuid.UUID]:
+        """Return the ID of every workspace in an organization."""
+        return await self.workspace_rows.get_ids_by_organization(organization_id)
+
+    async def get_organization_member_ids(self, organization_id: uuid.UUID) -> list[uuid.UUID]:
+        """Return the ID of every membership in an organization, whatever its status."""
+        return await self.members.get_ids_by_organization(organization_id)
+
+    async def get_workspace_member_ids_in_organization(self, organization_id: uuid.UUID) -> list[uuid.UUID]:
+        """Return the ID of every membership in an organization's workspaces, whatever its status."""
+        return await self.workspaces.get_ids_by_organization(organization_id)
+
     async def _to_context(
         self,
         *,

tests/integration/test_organization_scope_lookups.py

@@ -0,0 +1,159 @@
+"""Organizations answers which organization owns a workspace or a membership, and lists the IDs it holds."""
+
+import uuid
+from dataclasses import dataclass
+
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from gateway.models.tenancy import Organization, OrganizationMember, User, Workspace, WorkspaceMember
+from gateway.repositories.tenancy import (
+    OrganizationMemberRepository,
+    UserRepository,
+    WorkspaceMemberRepository,
+    WorkspaceRepository,
+)
+from gateway.services.tenancy.organization_service import OrganizationService
+
+from .tenancy_helpers import create_member, create_organization, create_workspace
+
+pytestmark = pytest.mark.asyncio
+
+
+@dataclass(frozen=True)
+class _Tenant:
+    organization: Organization
+    owner: User
+    workspace: Workspace
+    organization_member: OrganizationMember
+    workspace_member: WorkspaceMember
+
+
+async def _tenant(db: AsyncSession, slug: str) -> _Tenant:
+    organization = await create_organization(db, slug=slug)
+    owner = await create_member(db, organization, role="owner", full_name=f"{slug} owner")
+    workspace = await create_workspace(db, organization, name=f"{slug} workspace", owner=owner)
+    organization_member = await OrganizationMemberRepository(db).get_by_organization_and_user(organization.id, owner.id)
+    workspace_member = await WorkspaceMemberRepository(db).get_by_workspace_and_user(workspace.id, owner.id)
+    assert organization_member is not None
+    assert workspace_member is not None
+    return _Tenant(organization, owner, workspace, organization_member, workspace_member)
+
+
+async def _identity(db: AsyncSession, organization: Organization) -> User:
+    return await UserRepository(db).create_local_identity(full_name="Leaver", active_organization_id=organization.id)
+
+
+def _service(db: AsyncSession) -> OrganizationService:
+    return OrganizationService(db, membership_listener=None)
+
+
+async def test_a_workspace_resolves_to_its_organization(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    service = _service(async_db)
+
+    assert await service.get_organization_id_for_workspace(acme.workspace.id) == acme.organization.id
+    assert await service.get_organization_id_for_workspace(globex.workspace.id) == globex.organization.id
+    assert await service.get_organization_id_for_workspace(uuid.uuid4()) is None
+
+
+async def test_a_deleted_workspace_resolves_to_no_organization(async_db: AsyncSession) -> None:
+    """A deleted workspace has no organization, even after an earlier lookup found one."""
+    acme = await _tenant(async_db, "acme")
+    service = _service(async_db)
+    workspace_id = acme.workspace.id
+    assert await service.get_organization_id_for_workspace(workspace_id) == acme.organization.id
+
+    await WorkspaceRepository(async_db).delete_workspace(acme.workspace)
+
+    assert await service.get_organization_id_for_workspace(workspace_id) is None
+
+
+async def test_an_organization_membership_resolves_to_its_organization(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    service = _service(async_db)
+
+    assert (
+        await service.get_organization_id_for_organization_member(acme.organization_member.id)
+        == acme.organization.id
+    )
+    assert (
+        await service.get_organization_id_for_organization_member(globex.organization_member.id)
+        == globex.organization.id
+    )
+    assert await service.get_organization_id_for_organization_member(uuid.uuid4()) is None
+
+
+async def test_a_workspace_membership_resolves_to_its_workspace(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    service = _service(async_db)
+
+    assert await service.get_workspace_id_for_workspace_member(acme.workspace_member.id) == acme.workspace.id
+    assert await service.get_workspace_id_for_workspace_member(globex.workspace_member.id) == globex.workspace.id
+    assert await service.get_workspace_id_for_workspace_member(uuid.uuid4()) is None
+
+
+async def test_only_an_existing_organization_is_found(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    service = _service(async_db)
+
+    assert await service.has_organization(acme.organization.id) is True
+    assert await service.has_organization(uuid.uuid4()) is False
+
+
+async def test_an_organization_lists_only_its_own_workspaces(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    second = await create_workspace(async_db, acme.organization, name="acme second", owner=acme.owner)
+    service = _service(async_db)
+
+    assert sorted(await service.get_workspace_ids_in_organization(acme.organization.id)) == sorted(
+        [acme.workspace.id, second.id]
+    )
+    assert await service.get_workspace_ids_in_organization(globex.organization.id) == [globex.workspace.id]
+    assert await service.get_workspace_ids_in_organization(uuid.uuid4()) == []
+
+
+async def test_an_organization_lists_every_membership_whatever_its_status(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    leaver = await _identity(async_db, acme.organization)
+    suspended = await OrganizationMemberRepository(async_db).create_membership(
+        organization_id=acme.organization.id,
+        user_id=leaver.id,
+        role="member",
+        status="suspended",
+    )
+    service = _service(async_db)
+
+    assert sorted(await service.get_organization_member_ids(acme.organization.id)) == sorted(
+        [acme.organization_member.id, suspended.id]
+    )
+    assert await service.get_organization_member_ids(globex.organization.id) == [globex.organization_member.id]
+    assert await service.get_organization_member_ids(uuid.uuid4()) == []
+
+
+async def test_an_organization_lists_every_workspace_membership_whatever_its_status(async_db: AsyncSession) -> None:
+    acme = await _tenant(async_db, "acme")
+    globex = await _tenant(async_db, "globex")
+    second = await create_workspace(async_db, acme.organization, name="acme second", owner=acme.owner)
+    second_owner = await WorkspaceMemberRepository(async_db).get_by_workspace_and_user(second.id, acme.owner.id)
+    assert second_owner is not None
+    leaver = await _identity(async_db, acme.organization)
+    suspended = await WorkspaceMemberRepository(async_db).create(
+        workspace_id=acme.workspace.id,
+        user_id=leaver.id,
+        status="suspended",
+    )
+    service = _service(async_db)
+
+    assert sorted(await service.get_workspace_member_ids_in_organization(acme.organization.id)) == sorted(
+        [acme.workspace_member.id, second_owner.id, suspended.id]
+    )
+    assert await service.get_workspace_member_ids_in_organization(globex.organization.id) == [
+        globex.workspace_member.id
+    ]
+    assert await service.get_workspace_member_ids_in_organization(uuid.uuid4()) == []