Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: Implement Image RBAC project scope #3035

Open
wants to merge 9 commits into
base: topic/11-11-feat_implement_associatecontainerregistrywithgroup_disassociatecontainerregistrywithgroup_
Choose a base branch
from
1 change: 1 addition & 0 deletions changes/3035.feature.md
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
Add project scope implementation to Image RBAC.
48 changes: 46 additions & 2 deletions src/ai/backend/manager/models/image.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import trafaret as t
from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
from sqlalchemy.orm import joinedload, load_only, relationship, selectinload
from sqlalchemy.sql.expression import false, true

from ai.backend.common.docker import ImageRef
from ai.backend.common.etcd import AsyncEtcd
Expand All @@ -37,10 +38,10 @@
ResourceSlot,
)
from ai.backend.logging import BraceStyleAdapter
from ai.backend.manager.models.container_registry import ContainerRegistryRow

from ..api.exceptions import ImageNotFound
from ..container_registry import get_container_registry_cls
from ..models.container_registry import ContainerRegistryRow
from .base import (
GUID,
Base,
Expand Down Expand Up @@ -762,13 +763,56 @@ async def _in_domain_scope(
object_id_to_additional_permission_map=image_id_permission_map
)

async def _in_project_scope(
self,
ctx: ClientContext,
scope: ProjectScope,
) -> ImagePermissionContext:
from .group import GroupRow

permissions = await self.calculate_permission(ctx, scope)
image_id_permission_map: dict[UUID, frozenset[ImagePermission]] = {}

group_query_stmt = sa.select(GroupRow).where(GroupRow.id == scope.project_id)
group_row = cast(Optional[GroupRow], await self.db_session.scalar(group_query_stmt))
if group_row is None:
raise InvalidScope(f"Project not found (project_id:{scope.project_id})")

image_select_stmt = (
sa.select(ImageRow)
.select_from(
sa.join(
ImageRow, ContainerRegistryRow, ImageRow.registry_id == ContainerRegistryRow.id
)
)
.where(
sa.or_(
ContainerRegistryRow.is_global == true(),
sa.and_(
ContainerRegistryRow.is_global == false(),
ContainerRegistryRow.association_container_registries_groups_rows.any(
group_id=scope.project_id
),
),
)
)
)

for row in await self.db_session.scalars(image_select_stmt):
_row = cast(ImageRow, row)
image_id_permission_map[_row.id] = permissions

return ImagePermissionContext(
object_id_to_additional_permission_map=image_id_permission_map
)

@override
async def build_ctx_in_project_scope(
self,
ctx: ClientContext,
scope: ProjectScope,
) -> ImagePermissionContext:
return ImagePermissionContext()
return await self._in_project_scope(ctx, scope)

@override
async def build_ctx_in_user_scope(
Expand Down
Loading