feat(BA-3519): Add utils for RBAC entity data migration #7534
Conversation
github-actions
bot
added
size:L
There was a problem hiding this comment.
Pull request overview
This PR introduces reusable utility infrastructure for RBAC entity migration, enabling systematic migration of entity data and permissions from the original database to the RBAC database. The changes provide abstract interfaces and concrete adapters to transform migration data into the appropriate RBAC table formats.
Key Changes
- Introduces abstract interfaces (
AbstractEntityandAbstractEntityType) to define entity migration contracts - Adds
MigrationAdapterto transform migration data into RBAC table input formats - Implements
EntityMigrationUtilwith methods for querying permission groups and adding entity permissions/scope associations
Reviewed changes
Copilot reviewed 4 out of 5 changed files in this pull request and generated 10 comments.
| File | Description |
|---|---|
src/ai/backend/manager/models/rbac_models/migration/entity/__init__.py |
Empty package marker for new entity subdirectory |
src/ai/backend/manager/models/rbac_models/migration/entity/abc.py |
Defines abstract interfaces for entities and entity types used during migration |
src/ai/backend/manager/models/rbac_models/migration/adapter.py |
Implements data transformation adapter that parses entity data into RBAC table input formats |
src/ai/backend/manager/models/rbac_models/migration/utils.py |
Adds EntityMigrationUtil class with database operations for entity migration, including new imports for adapter and entity abstractions |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| ) | ||
| result = db_conn.execute(query) | ||
| rows = result.all() | ||
| return [(row.permission_group_id, row.role_source) for row in rows] |
There was a problem hiding this comment.
Type mismatch in return value. The method signature declares it returns list[tuple[uuid.UUID, RoleSource]], but the roles_table.c.source column is defined as VARCHAR(16) in models.py, so querying it returns a string, not a RoleSource enum. The return value at line 180 will be list[tuple[uuid.UUID, str]]. Consider either converting the string to RoleSource enum (e.g., RoleSource(row.role_source)) or updating the return type annotation to match the actual type.
| return [(row.permission_group_id, row.role_source) for row in rows] | |
| return [ | |
| (row.permission_group_id, RoleSource(row.role_source)) | |
| for row in rows | |
| ] |
| def add_entity_typed_permission_to_permission_groups( | ||
| self, | ||
| db_conn: Connection, | ||
| role_source: RoleSource, | ||
| permission_group_id: uuid.UUID, | ||
| entity_type: AbstractEntityType, | ||
| ) -> None: | ||
| match role_source: | ||
| case RoleSource.SYSTEM: | ||
| operations = entity_type.operations_in_system_role() | ||
| case RoleSource.CUSTOM: | ||
| operations = entity_type.operations_in_custom_role() | ||
|
|
||
| permission_inputs = self._adapter.parse_entity_typed_permissions( | ||
| permission_group_id=permission_group_id, | ||
| entity_type=entity_type.entity_type(), | ||
| operations=operations, | ||
| ) | ||
| insert_if_data_exists( | ||
| db_conn, | ||
| get_permissions_table(), | ||
| [input.to_dict() for input in permission_inputs], | ||
| ) |
There was a problem hiding this comment.
Missing docstring. This public method should document its purpose, parameters, and behavior. Specifically, it should explain how permissions are determined based on role_source (SYSTEM roles get admin operations, CUSTOM roles get member operations from the entity type).
| class MigrationAdapter: | ||
| """ | ||
| RBAC Migration Adapter. | ||
| Provides methods to parse various migration input types. | ||
| """ | ||
|
|
||
| def parse_entity_typed_permissions( | ||
| self, | ||
| permission_group_id: UUID, | ||
| entity_type: EntityType, | ||
| operations: Iterable[OperationType], | ||
| ) -> list[PermissionCreateInput]: | ||
| return [ | ||
| PermissionCreateInput( | ||
| permission_group_id=permission_group_id, | ||
| entity_type=entity_type, | ||
| operation=operation, | ||
| ) | ||
| for operation in operations | ||
| ] | ||
|
|
||
| def parse_association_scopes_entities( | ||
| self, entity: AbstractEntity | ||
| ) -> list[AssociationScopesEntitiesCreateInput]: | ||
| result: list[AssociationScopesEntitiesCreateInput] = [] | ||
| entity_id = entity.entity_id() | ||
| for scope in entity.scopes(): | ||
| result.append( | ||
| AssociationScopesEntitiesCreateInput( | ||
| scope_id=scope, | ||
| object_id=entity_id, | ||
| ) | ||
| ) | ||
| return result |
There was a problem hiding this comment.
Missing test coverage for the new MigrationAdapter class and its methods. The existing test files in tests/manager/rbac/ show comprehensive testing of migration utilities. Consider adding tests to verify that parse_entity_typed_permissions correctly creates PermissionCreateInput instances and that parse_association_scopes_entities properly handles entity scope mappings.
| case RoleSource.SYSTEM: | ||
| operations = entity_type.operations_in_system_role() | ||
| case RoleSource.CUSTOM: | ||
| operations = entity_type.operations_in_custom_role() |
There was a problem hiding this comment.
The match statement lacks a default case. If role_source somehow has an unexpected value, the operations variable will be undefined, causing a NameError at line 195. Consider adding a case _ clause to handle unexpected values or raise an explicit error for invalid role sources.
| operations = entity_type.operations_in_custom_role() | |
| operations = entity_type.operations_in_custom_role() | |
| case _: | |
| raise ValueError(f"Unexpected role source: {role_source!r}") |
| def get_permission_group_ids_with_role_source( | ||
| self, | ||
| db_conn: Connection, | ||
| offset: int, | ||
| limit: int, | ||
| ) -> list[tuple[uuid.UUID, RoleSource]]: | ||
| roles_table = get_roles_table() | ||
| permission_groups_table = get_permission_groups_table() | ||
| query = ( | ||
| sa.select( | ||
| roles_table.c.source.label("role_source"), | ||
| permission_groups_table.c.id.label("permission_group_id"), | ||
| ) | ||
| .join( | ||
| permission_groups_table, | ||
| roles_table.c.id == permission_groups_table.c.role_id, | ||
| ) | ||
| .order_by(permission_groups_table.c.id) | ||
| .offset(offset) | ||
| .limit(limit) | ||
| ) | ||
| result = db_conn.execute(query) | ||
| rows = result.all() | ||
| return [(row.permission_group_id, row.role_source) for row in rows] |
There was a problem hiding this comment.
Missing docstring. This public method would benefit from documentation explaining its purpose, parameters (especially the pagination parameters offset and limit), and return value format. Consider documenting that it returns tuples of (permission_group_id, role_source) and how the pagination works.
| def add_entity_to_scopes( | ||
| self, | ||
| db_conn: Connection, | ||
| entity: AbstractEntity, | ||
| ) -> None: | ||
| association_scopes_entities_table = get_association_scopes_entities_table() | ||
| entity_inputs = self._adapter.parse_association_scopes_entities(entity) | ||
| insert_skip_on_conflict( | ||
| db_conn, | ||
| association_scopes_entities_table, | ||
| [input.to_dict() for input in entity_inputs], | ||
| ) |
There was a problem hiding this comment.
Missing docstring. This public method should document its purpose and behavior. It should explain that it creates associations between an entity and its owning scopes in the association_scopes_entities table, using insert_skip_on_conflict to handle duplicates.
| from ..enums import EntityType, OperationType | ||
|
|
||
|
|
||
| class AbstractEntity(ABC): |
There was a problem hiding this comment.
Missing class docstring. Consider adding a docstring that explains the purpose of this interface and what concrete implementations should represent (e.g., "Represents an entity instance during RBAC migration, providing access to the entity's ID and associated scopes").
| pass | ||
|
|
||
|
|
||
| class AbstractEntityType(ABC): |
There was a problem hiding this comment.
Missing class docstring. Consider adding a docstring that explains the purpose of this interface and what concrete implementations should represent (e.g., "Defines the entity type configuration for RBAC migration, specifying which operations are available in system vs. custom roles").
| class AbstractEntityType(ABC): | |
| class AbstractEntityType(ABC): | |
| """ | |
| Defines the entity type configuration for RBAC migration. | |
| Concrete implementations represent a specific entity type and declare | |
| which operations are permitted when that entity is used in system roles | |
| versus custom roles. | |
| """ |
| class EntityMigrationUtil: | ||
| """ | ||
| RBAC Entity Migration Utility. | ||
| Provides methods to handle entity-related migrations. | ||
| """ | ||
|
|
||
| def __init__(self) -> None: | ||
| self._adapter = MigrationAdapter() | ||
|
|
||
| def get_permission_group_ids_with_role_source( | ||
| self, | ||
| db_conn: Connection, | ||
| offset: int, | ||
| limit: int, | ||
| ) -> list[tuple[uuid.UUID, RoleSource]]: | ||
| roles_table = get_roles_table() | ||
| permission_groups_table = get_permission_groups_table() | ||
| query = ( | ||
| sa.select( | ||
| roles_table.c.source.label("role_source"), | ||
| permission_groups_table.c.id.label("permission_group_id"), | ||
| ) | ||
| .join( | ||
| permission_groups_table, | ||
| roles_table.c.id == permission_groups_table.c.role_id, | ||
| ) | ||
| .order_by(permission_groups_table.c.id) | ||
| .offset(offset) | ||
| .limit(limit) | ||
| ) | ||
| result = db_conn.execute(query) | ||
| rows = result.all() | ||
| return [(row.permission_group_id, row.role_source) for row in rows] | ||
|
|
||
| def add_entity_typed_permission_to_permission_groups( | ||
| self, | ||
| db_conn: Connection, | ||
| role_source: RoleSource, | ||
| permission_group_id: uuid.UUID, | ||
| entity_type: AbstractEntityType, | ||
| ) -> None: | ||
| match role_source: | ||
| case RoleSource.SYSTEM: | ||
| operations = entity_type.operations_in_system_role() | ||
| case RoleSource.CUSTOM: | ||
| operations = entity_type.operations_in_custom_role() | ||
|
|
||
| permission_inputs = self._adapter.parse_entity_typed_permissions( | ||
| permission_group_id=permission_group_id, | ||
| entity_type=entity_type.entity_type(), | ||
| operations=operations, | ||
| ) | ||
| insert_if_data_exists( | ||
| db_conn, | ||
| get_permissions_table(), | ||
| [input.to_dict() for input in permission_inputs], | ||
| ) | ||
|
|
||
| def add_entity_to_scopes( | ||
| self, | ||
| db_conn: Connection, | ||
| entity: AbstractEntity, | ||
| ) -> None: | ||
| association_scopes_entities_table = get_association_scopes_entities_table() | ||
| entity_inputs = self._adapter.parse_association_scopes_entities(entity) | ||
| insert_skip_on_conflict( | ||
| db_conn, | ||
| association_scopes_entities_table, | ||
| [input.to_dict() for input in entity_inputs], | ||
| ) |
There was a problem hiding this comment.
Missing test coverage for the new EntityMigrationUtil class and its methods. The existing test files show comprehensive testing of migration utilities. Consider adding tests to verify: 1) get_permission_group_ids_with_role_source correctly queries and paginates results, 2) add_entity_typed_permission_to_permission_groups properly handles both SYSTEM and CUSTOM role sources, and 3) add_entity_to_scopes correctly creates scope-entity associations.
| from abc import ABC, abstractmethod | ||
|
|
||
| from ai.backend.manager.data.permission.id import ObjectId, ScopeId | ||
|
|
||
| from ..enums import EntityType, OperationType | ||
|
|
||
|
|
||
| class AbstractEntity(ABC): | ||
| @abstractmethod | ||
| def scopes(self) -> list[ScopeId]: | ||
| pass | ||
|
|
||
| @abstractmethod | ||
| def entity_id(self) -> ObjectId: | ||
| pass | ||
|
|
||
|
|
||
| class AbstractEntityType(ABC): | ||
| @abstractmethod | ||
| def entity_type(self) -> EntityType: | ||
| pass | ||
|
|
||
| @abstractmethod | ||
| def operations_in_system_role(self) -> set[OperationType]: | ||
| pass | ||
|
|
||
| @abstractmethod | ||
| def operations_in_custom_role(self) -> set[OperationType]: | ||
| pass |
There was a problem hiding this comment.
Missing test coverage for the new abstract interfaces. The existing test files show comprehensive testing patterns. Consider adding tests with concrete implementations of AbstractEntity and AbstractEntityType to demonstrate the expected behavior and validate the interface design.
fregataa
force-pushed
the
feat/model-utils-for-rbac-migration
branch
2 times, most recently
from
e259673 to
4d0e99f
Compare
fregataa
force-pushed
the
feat/model-utils-for-rbac-migration
branch
from
4d0e99f to
ebc8ec3
Compare
HyeockJinKim
force-pushed
the
main
branch
2 times, most recently
from
9552aac to
4af738e
Compare
resolves #7530 (BA-3519)
Summary
Add reusable utility classes and adapters for RBAC entity data migration.
Changes
Added Components:
- Parses entity-typed permissions into PermissionCreateInput
- Parses entity scope associations into AssociationScopesEntitiesCreateInput
- AbstractEntity: Interface for accessing entity scopes and IDs
- AbstractEntityType: Interface for defining entity type operations in system/custom roles
- Query permission groups with their role sources
- Add entity-typed permissions to permission groups based on role source
- Associate entities with scopes in association_scopes_entities table
Implementation Details
Entity Scope Mapping:
Entity Type Permission Creation:
Checklist: (if applicable)
ai.backend.testdocsdirectory