diff --git a/changes/12934.feature.md b/changes/12934.feature.md new file mode 100644 index 00000000000..b42098e9f32 --- /dev/null +++ b/changes/12934.feature.md @@ -0,0 +1 @@ +Add ensure-virtual-scope RBAC ops that bind a scope to its own (and optionally its parent's) virtual scope, so ownership resolves through the single virtual-scope path. diff --git a/src/ai/backend/manager/repositories/ops/rbac/provider.py b/src/ai/backend/manager/repositories/ops/rbac/provider.py index eb5ef6c4c56..13ba94df8d6 100644 --- a/src/ai/backend/manager/repositories/ops/rbac/provider.py +++ b/src/ai/backend/manager/repositories/ops/rbac/provider.py @@ -116,40 +116,52 @@ async def _resolve_virtual_scope_id(self, scope: ScopeRef) -> VirtualScopeID: return virtual_scope_id async def _insert_virtual_scopes(self, scopes: Sequence[ScopeRef]) -> None: - """Create the virtual scope nodes for ``scopes`` and register each scope - as an entity member of its own virtual scope (idempotent). - - Access to the scope object itself then resolves as an ordinary entity through the - single-path permission resolution. ``permission_cap`` is NULL (ownership, no ceiling). - """ + """Create each scope's virtual scope node with its self entity-membership and self + scope_binding (permission_cap NULL). Idempotent: an existing scope is a no-op.""" if not scopes: return values = [{"scope_type": s.scope_type, "scope_id": s.scope_id} for s in scopes] - stmt = ( + insert_stmt = ( pg_insert(VirtualScopeRow) .values(values) .on_conflict_do_nothing(index_elements=["scope_type", "scope_id"]) + .returning( + VirtualScopeRow.id, + VirtualScopeRow.scope_type, + VirtualScopeRow.scope_id, + ) ) - await self._sess.execute(stmt) - membership_source = sa.select( - VirtualScopeRow.id, - VirtualScopeRow.scope_type, - VirtualScopeRow.scope_id, - sa.null(), - ).where( - sa.tuple_(VirtualScopeRow.scope_type, VirtualScopeRow.scope_id).in_([ - (s.scope_type, s.scope_id) for s in scopes - ]) - ) + inserted = (await self._sess.execute(insert_stmt)).all() + if not inserted: + return membership_stmt = ( pg_insert(EntityMembershipRow) - .from_select( - ["virtual_scope_id", "entity_type", "entity_id", "permission_cap"], - membership_source, - ) + .values([ + { + "virtual_scope_id": row.id, + "entity_type": row.scope_type, + "entity_id": row.scope_id, + "permission_cap": None, + } + for row in inserted + ]) .on_conflict_do_nothing() ) await self._sess.execute(membership_stmt) + binding_stmt = ( + pg_insert(ScopeBindingRow) + .values([ + { + "virtual_scope_id": row.id, + "scope_type": row.scope_type, + "scope_id": row.scope_id, + "permission_cap": None, + } + for row in inserted + ]) + .on_conflict_do_nothing() + ) + await self._sess.execute(binding_stmt) async def _delete_virtual_scopes(self, scopes: Sequence[ScopeRef]) -> None: """Delete the virtual scope nodes for ``scopes`` (FK CASCADE removes their edges).""" @@ -273,14 +285,18 @@ async def add_users_to_scope( async def create_scope[TRow: Base]( self, creation: ScopeCreation[TRow], + bound_scope: ScopeRef | None = None, ) -> CreatorResult[TRow]: """Create the real scope entity via ``creation.creator`` and its virtual scope node. - The virtual scope insert is idempotent (get-or-create). ``creation.scope.scope_id`` - must match the id the created row carries. + ``creation.scope.scope_id`` must match the id the created row carries. When + ``bound_scope`` is given, it is bound to this scope's virtual scope so it can reach + this scope's entities. """ result = await self.create(creation.creator) await self._insert_virtual_scopes([creation.scope]) + if bound_scope is not None: + await self.bind_scope(bound_scope, creation.scope, permission_cap=None) return result async def bulk_create_scopes[TRow: Base]( @@ -409,6 +425,20 @@ async def remove_entity_members( ) await self._sess.execute(stmt) + # -- Virtual scope: ensure compatibility for externally-created rows ---------- + + async def ensure_scope( + self, + scope: ScopeRef, + bound_scope: ScopeRef | None = None, + ) -> None: + """Ensure the virtual scope node for an already-created ``scope``. When ``bound_scope`` + is given, it is bound to this scope's virtual scope. Idempotent. + """ + await self._insert_virtual_scopes([scope]) + if bound_scope is not None: + await self.bind_scope(bound_scope, scope, permission_cap=None) + class RBACOpsProvider(DBOpsProvider): """Hands out :class:`RBACWriteOps` for the read-write surface.""" diff --git a/tests/unit/manager/repositories/ops/rbac/test_provider.py b/tests/unit/manager/repositories/ops/rbac/test_provider.py index a080b723503..a13d02f4cf5 100644 --- a/tests/unit/manager/repositories/ops/rbac/test_provider.py +++ b/tests/unit/manager/repositories/ops/rbac/test_provider.py @@ -1,14 +1,4 @@ -"""Integration tests for the RBAC ops provider (RBACWriteOps) with a real database. - -These verify observable outcomes — which rows and which scope associations survive an -operation — rather than internal call wiring. Two surfaces are covered: - -- ``create_scope`` / ``bulk_create_scopes`` materialize both the virtual scope node and - the scope-as-entity membership in the scope's own virtual scope. -- ``bulk_create_scoped_partial`` / ``bulk_purge_scoped_partial`` isolate each item so a - rejected one leaves the rest of the batch intact. The unscoped counterparts - (``bulk_create_partial`` / ``bulk_purge_partial``) are covered by the base ops tests. -""" +"""Integration tests for the RBAC ops provider (RBACWriteOps) with a real database.""" from __future__ import annotations @@ -50,6 +40,7 @@ from ai.backend.manager.models.user import UserRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine from ai.backend.manager.models.virtual_scope.entity_membership import EntityMembershipRow +from ai.backend.manager.models.virtual_scope.scope_binding import ScopeBindingRow from ai.backend.manager.models.virtual_scope.virtual_scope import VirtualScopeRow from ai.backend.manager.repositories.base import Creator, CreatorSpec from ai.backend.manager.repositories.base.rbac.entity_creator import RBACEntityCreator @@ -80,6 +71,7 @@ _TEST_SCOPE_TYPE = ScopeType("test-scope") _TEST_ENTITY_TYPE = VirtualScopeEntityType("test-scope") +_TEST_BOUND_SCOPE_TYPE = ScopeType("test-bound-scope") _USER_SCOPE_ID = str(uuid.uuid4()) _USER_SCOPE_REF = RBACElementRef(RBACElementType.USER, _USER_SCOPE_ID) @@ -175,6 +167,7 @@ class _ScopedRow: OpsRBACScopeRow, VirtualScopeRow, EntityMembershipRow, + ScopeBindingRow, ] @@ -240,16 +233,18 @@ def bulk_scopes() -> BulkScopeContext: class TestScopeCreationVirtualScope: - """create_scope / bulk_create_scopes materialize the VS node and self-membership.""" + """create_scope / bulk_create_scopes materialize the VS node, self-membership, and + self scope_binding (plus a bound-scope binding when a bound_scope is given).""" - async def test_create_scope_adds_virtual_scope_and_self_membership( + async def test_create_scope_adds_virtual_scope_membership_and_self_binding( self, database_connection: ExtendedAsyncSAEngine, provider: RBACOpsProvider, scope_tables: None, single_scope: SingleScopeContext, ) -> None: - """create_scope creates the VS node and registers the scope in its own VS.""" + """create_scope creates the VS node, registers the scope in its own VS, and binds + the scope to its own VS.""" scope_id = single_scope.scope_id async with provider.write_ops() as w: @@ -268,6 +263,7 @@ async def test_create_scope_adds_virtual_scope_and_self_membership( .all() ) membership_rows = (await sess.execute(sa.select(EntityMembershipRow))).scalars().all() + binding_rows = (await sess.execute(sa.select(ScopeBindingRow))).scalars().all() assert len(vs_rows) == 1 vs = vs_rows[0] @@ -281,14 +277,58 @@ async def test_create_scope_adds_virtual_scope_and_self_membership( assert membership.entity_id == scope_id assert membership.permission_cap is None - async def test_bulk_create_scopes_adds_vs_and_self_membership_per_scope( + assert len(binding_rows) == 1 + binding = binding_rows[0] + assert binding.virtual_scope_id == vs.id + assert binding.scope_type == _TEST_SCOPE_TYPE + assert binding.scope_id == scope_id + assert binding.permission_cap is None + + async def test_create_scope_with_bound_scope_binds_it_to_its_virtual_scope( + self, + database_connection: ExtendedAsyncSAEngine, + provider: RBACOpsProvider, + scope_tables: None, + single_scope: SingleScopeContext, + ) -> None: + """create_scope(bound_scope=...) binds that scope to the new scope's VS on top of + the self binding, so it reaches this scope's entities in one hop.""" + scope_id = single_scope.scope_id + bound_scope = ScopeRef(scope_type=_TEST_BOUND_SCOPE_TYPE, scope_id=uuid.uuid4()) + + async with provider.write_ops() as w: + await w.create_scope(single_scope.creation, bound_scope=bound_scope) + + async with database_connection.begin_session_read_committed() as sess: + vs = ( + await sess.execute( + sa.select(VirtualScopeRow).where(VirtualScopeRow.scope_id == scope_id) + ) + ).scalar_one() + binding_rows = ( + ( + await sess.execute( + sa.select(ScopeBindingRow).where(ScopeBindingRow.virtual_scope_id == vs.id) + ) + ) + .scalars() + .all() + ) + + assert {(b.scope_type, b.scope_id) for b in binding_rows} == { + (_TEST_SCOPE_TYPE, scope_id), # self binding + (_TEST_BOUND_SCOPE_TYPE, bound_scope.scope_id), # bound-scope binding + } + + async def test_bulk_create_scopes_adds_vs_membership_and_self_binding_per_scope( self, database_connection: ExtendedAsyncSAEngine, provider: RBACOpsProvider, scope_tables: None, bulk_scopes: BulkScopeContext, ) -> None: - """bulk_create_scopes creates one VS node and one self-membership per scope.""" + """bulk_create_scopes creates one VS node, self-membership, and self binding per + scope.""" scope_ids = bulk_scopes.scope_ids async with provider.write_ops() as w: @@ -297,6 +337,7 @@ async def test_bulk_create_scopes_adds_vs_and_self_membership_per_scope( async with database_connection.begin_session_read_committed() as sess: vs_rows = (await sess.execute(sa.select(VirtualScopeRow))).scalars().all() membership_rows = (await sess.execute(sa.select(EntityMembershipRow))).scalars().all() + binding_rows = (await sess.execute(sa.select(ScopeBindingRow))).scalars().all() assert {vs.scope_id for vs in vs_rows} == set(scope_ids) vs_by_scope = {vs.scope_id: vs for vs in vs_rows} @@ -308,6 +349,117 @@ async def test_bulk_create_scopes_adds_vs_and_self_membership_per_scope( assert membership.virtual_scope_id == vs_by_scope[membership.entity_id].id assert membership.permission_cap is None + assert len(binding_rows) == 3 + for binding in binding_rows: + assert binding.scope_type == _TEST_SCOPE_TYPE + assert binding.scope_id in scope_ids + assert binding.virtual_scope_id == vs_by_scope[binding.scope_id].id + assert binding.permission_cap is None + + +class TestEnsureScope: + """ensure_scope backfills the VS node, self-membership, and self binding for an + already-created scope, without creating the real scope row.""" + + async def test_ensure_scope_adds_vs_membership_and_self_binding_without_real_row( + self, + database_connection: ExtendedAsyncSAEngine, + provider: RBACOpsProvider, + scope_tables: None, + ) -> None: + """ensure_scope creates the VS node, self-membership, and self binding, and leaves + no real scope row behind.""" + scope_id = uuid.uuid4() + scope = ScopeRef(scope_type=_TEST_SCOPE_TYPE, scope_id=scope_id) + + async with provider.write_ops() as w: + await w.ensure_scope(scope) + + async with database_connection.begin_session_read_committed() as sess: + vs_rows = (await sess.execute(sa.select(VirtualScopeRow))).scalars().all() + membership_rows = (await sess.execute(sa.select(EntityMembershipRow))).scalars().all() + binding_rows = (await sess.execute(sa.select(ScopeBindingRow))).scalars().all() + real_row_count = await sess.scalar( + sa.select(sa.func.count()).select_from(OpsRBACScopeRow) + ) + + assert len(vs_rows) == 1 + vs = vs_rows[0] + assert vs.scope_id == scope_id + + assert len(membership_rows) == 1 + assert membership_rows[0].virtual_scope_id == vs.id + assert membership_rows[0].entity_id == scope_id + + assert len(binding_rows) == 1 + assert binding_rows[0].virtual_scope_id == vs.id + assert binding_rows[0].scope_id == scope_id + assert binding_rows[0].permission_cap is None + + assert real_row_count == 0 + + async def test_ensure_scope_is_idempotent( + self, + database_connection: ExtendedAsyncSAEngine, + provider: RBACOpsProvider, + scope_tables: None, + ) -> None: + """Calling ensure_scope twice leaves exactly one VS node, membership, and binding.""" + scope = ScopeRef(scope_type=_TEST_SCOPE_TYPE, scope_id=uuid.uuid4()) + + async with provider.write_ops() as w: + await w.ensure_scope(scope) + await w.ensure_scope(scope) + + async with database_connection.begin_session_read_committed() as sess: + vs_count = await sess.scalar(sa.select(sa.func.count()).select_from(VirtualScopeRow)) + membership_count = await sess.scalar( + sa.select(sa.func.count()).select_from(EntityMembershipRow) + ) + binding_count = await sess.scalar( + sa.select(sa.func.count()).select_from(ScopeBindingRow) + ) + + assert vs_count == 1 + assert membership_count == 1 + assert binding_count == 1 + + async def test_ensure_scope_with_bound_scope_adds_binding( + self, + database_connection: ExtendedAsyncSAEngine, + provider: RBACOpsProvider, + scope_tables: None, + ) -> None: + """ensure_scope(scope, bound_scope) binds both the scope itself and the bound scope + to the scope's VS.""" + scope_id = uuid.uuid4() + scope = ScopeRef(scope_type=_TEST_SCOPE_TYPE, scope_id=scope_id) + bound_scope = ScopeRef(scope_type=_TEST_BOUND_SCOPE_TYPE, scope_id=uuid.uuid4()) + + async with provider.write_ops() as w: + await w.ensure_scope(scope, bound_scope=bound_scope) + + async with database_connection.begin_session_read_committed() as sess: + vs = ( + await sess.execute( + sa.select(VirtualScopeRow).where(VirtualScopeRow.scope_id == scope_id) + ) + ).scalar_one() + binding_rows = ( + ( + await sess.execute( + sa.select(ScopeBindingRow).where(ScopeBindingRow.virtual_scope_id == vs.id) + ) + ) + .scalars() + .all() + ) + + assert {(b.scope_type, b.scope_id) for b in binding_rows} == { + (_TEST_SCOPE_TYPE, scope_id), # self binding + (_TEST_BOUND_SCOPE_TYPE, bound_scope.scope_id), # bound-scope binding + } + @pytest.fixture async def scoped_rows(