Skip to content

Commit 9cf4cf3

Browse files
feat(BA-2107): Add database migration code and functions for RBAC system (#5734)
Co-authored-by: octodog <mu001@lablup.com>
1 parent 52af4b3 commit 9cf4cf3

78 files changed

Lines changed: 4355 additions & 662 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

changes/5340.feature.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add data migration script from VFolder to RBAC tables

changes/5417.feature.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Migrate existing user/project records to RBAC data

changes/5465.feature.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Expand RBAC tables by adding `permission_groups` table to group permissions with the same target

changes/5699.feature.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Add RBAC repositoy functions to manage scopes and entity DB records

docs/manager/rest-reference/openapi.json

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -402,15 +402,15 @@
402402
"title": "HarborWebhookRequestModel",
403403
"type": "object"
404404
},
405-
"VFolderPermission": {
405+
"VFolderMountPermission": {
406406
"description": "Permissions for a virtual folder given to a specific access key.\nRW_DELETE includes READ_WRITE and READ_WRITE includes READ_ONLY.",
407407
"enum": [
408408
"ro",
409409
"rw",
410410
"wd",
411411
"wd"
412412
],
413-
"title": "VFolderPermission",
413+
"title": "VFolderMountPermission",
414414
"type": "string"
415415
},
416416
"VFolderUsageMode": {
@@ -447,7 +447,7 @@
447447
"default": "general"
448448
},
449449
"permission": {
450-
"$ref": "#/components/schemas/VFolderPermission",
450+
"$ref": "#/components/schemas/VFolderMountPermission",
451451
"default": "rw"
452452
},
453453
"unmanagedPath": {
@@ -613,7 +613,7 @@
613613
"permission": {
614614
"anyOf": [
615615
{
616-
"$ref": "#/components/schemas/VFolderPermission"
616+
"$ref": "#/components/schemas/VFolderMountPermission"
617617
},
618618
{
619619
"type": "null"

src/ai/backend/manager/data/domain/__init__.py

Whitespace-only changes.
Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
import uuid
2+
from collections.abc import Iterable, Mapping
3+
from dataclasses import dataclass, field
4+
from datetime import datetime
5+
from typing import Any, Optional, override
6+
7+
from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
8+
from ai.backend.manager.data.permission.id import ScopeId
9+
from ai.backend.manager.data.permission.types import (
10+
EntityType,
11+
OperationType,
12+
ScopeType,
13+
)
14+
from ai.backend.manager.data.user.types import UserRole
15+
from ai.backend.manager.types import Creator, OptionalState, PartialModifier, TriState
16+
17+
18+
@dataclass
19+
class UserInfo:
20+
id: uuid.UUID
21+
role: UserRole
22+
domain_name: str
23+
24+
25+
@dataclass
26+
class DomainData:
27+
name: str
28+
description: Optional[str]
29+
is_active: bool
30+
created_at: datetime = field(compare=False)
31+
modified_at: datetime = field(compare=False)
32+
total_resource_slots: ResourceSlot
33+
allowed_vfolder_hosts: VFolderHostPermissionMap
34+
allowed_docker_registries: list[str]
35+
dotfiles: bytes
36+
integration_id: Optional[str]
37+
38+
def scope_id(self) -> ScopeId:
39+
return ScopeId(
40+
scope_type=ScopeType.DOMAIN,
41+
scope_id=self.name,
42+
)
43+
44+
def role_name(self) -> str:
45+
return f"domain-{self.name}-admin"
46+
47+
def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
48+
return {
49+
entity: OperationType.admin_operations()
50+
for entity in EntityType.admin_accessible_entity_types_in_domain()
51+
}
52+
53+
54+
@dataclass
55+
class DomainCreator(Creator):
56+
name: str
57+
description: Optional[str] = None
58+
is_active: Optional[bool] = None
59+
total_resource_slots: Optional[ResourceSlot] = None
60+
allowed_vfolder_hosts: Optional[dict[str, list[str]]] = None
61+
allowed_docker_registries: Optional[list[str]] = None
62+
integration_id: Optional[str] = None
63+
dotfiles: Optional[bytes] = None
64+
65+
@override
66+
def fields_to_store(self) -> dict[str, Any]:
67+
to_store: dict[str, Any] = {"name": self.name}
68+
if self.description is not None:
69+
to_store["description"] = self.description
70+
if self.is_active is not None:
71+
to_store["is_active"] = self.is_active
72+
if self.total_resource_slots is not None:
73+
to_store["total_resource_slots"] = self.total_resource_slots
74+
if self.allowed_vfolder_hosts is not None:
75+
to_store["allowed_vfolder_hosts"] = self.allowed_vfolder_hosts
76+
if self.allowed_docker_registries is not None:
77+
to_store["allowed_docker_registries"] = self.allowed_docker_registries
78+
if self.integration_id is not None:
79+
to_store["integration_id"] = self.integration_id
80+
if self.dotfiles is not None:
81+
to_store["dotfiles"] = self.dotfiles
82+
return to_store
83+
84+
85+
@dataclass
86+
class DomainModifier(PartialModifier):
87+
name: OptionalState[str] = field(default_factory=OptionalState.nop)
88+
description: TriState[str] = field(default_factory=TriState.nop)
89+
is_active: OptionalState[bool] = field(default_factory=OptionalState.nop)
90+
total_resource_slots: TriState[ResourceSlot] = field(default_factory=TriState.nop)
91+
allowed_vfolder_hosts: OptionalState[dict[str, list[str]]] = field(
92+
default_factory=OptionalState.nop
93+
)
94+
allowed_docker_registries: OptionalState[list[str]] = field(default_factory=OptionalState.nop)
95+
integration_id: TriState[str] = field(default_factory=TriState.nop)
96+
97+
@override
98+
def fields_to_update(self) -> dict[str, Any]:
99+
to_update: dict[str, Any] = {}
100+
self.name.update_dict(to_update, "name")
101+
self.description.update_dict(to_update, "description")
102+
self.is_active.update_dict(to_update, "is_active")
103+
self.total_resource_slots.update_dict(to_update, "total_resource_slots")
104+
self.allowed_vfolder_hosts.update_dict(to_update, "allowed_vfolder_hosts")
105+
self.allowed_docker_registries.update_dict(to_update, "allowed_docker_registries")
106+
self.integration_id.update_dict(to_update, "integration_id")
107+
return to_update
108+
109+
110+
@dataclass
111+
class DomainNodeModifier(PartialModifier):
112+
description: TriState[str] = field(default_factory=TriState[str].nop)
113+
is_active: OptionalState[bool] = field(default_factory=OptionalState[bool].nop)
114+
total_resource_slots: TriState[ResourceSlot] = field(default_factory=TriState[ResourceSlot].nop)
115+
allowed_vfolder_hosts: OptionalState[dict[str, list[str]]] = field(
116+
default_factory=OptionalState[dict[str, list[str]]].nop
117+
)
118+
allowed_docker_registries: OptionalState[list[str]] = field(
119+
default_factory=OptionalState[list[str]].nop
120+
)
121+
integration_id: TriState[str] = field(default_factory=TriState[str].nop)
122+
dotfiles: OptionalState[bytes] = field(default_factory=OptionalState[bytes].nop)
123+
124+
@override
125+
def fields_to_update(self) -> dict[str, Any]:
126+
to_update: dict[str, Any] = {}
127+
self.description.update_dict(to_update, "description")
128+
self.is_active.update_dict(to_update, "is_active")
129+
self.total_resource_slots.update_dict(to_update, "total_resource_slots")
130+
self.allowed_vfolder_hosts.update_dict(to_update, "allowed_vfolder_hosts")
131+
self.allowed_docker_registries.update_dict(to_update, "allowed_docker_registries")
132+
self.integration_id.update_dict(to_update, "integration_id")
133+
self.dotfiles.update_dict(to_update, "dotfiles")
134+
return to_update

src/ai/backend/manager/data/group/__init__.py

Whitespace-only changes.
Lines changed: 132 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,132 @@
1+
from __future__ import annotations
2+
3+
import enum
4+
import uuid
5+
from collections.abc import Iterable, Mapping
6+
from dataclasses import dataclass, field
7+
from datetime import datetime
8+
from typing import Any, Optional, override
9+
10+
from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap
11+
from ai.backend.manager.data.permission.id import ScopeId
12+
from ai.backend.manager.data.permission.types import (
13+
EntityType,
14+
OperationType,
15+
ScopeType,
16+
)
17+
from ai.backend.manager.types import Creator, OptionalState, PartialModifier, TriState
18+
19+
20+
class ProjectType(enum.StrEnum):
21+
GENERAL = "general"
22+
MODEL_STORE = "model-store"
23+
24+
@classmethod
25+
def _missing_(cls, value: Any) -> Optional[ProjectType]:
26+
assert isinstance(value, str)
27+
match value.upper():
28+
case "GENERAL":
29+
return cls.GENERAL
30+
case "MODEL_STORE" | "MODEL-STORE":
31+
return cls.MODEL_STORE
32+
return None
33+
34+
35+
@dataclass
36+
class GroupCreator(Creator):
37+
name: str
38+
domain_name: str
39+
type: Optional[ProjectType] = None
40+
description: Optional[str] = None
41+
is_active: Optional[bool] = None
42+
total_resource_slots: Optional[ResourceSlot] = None
43+
allowed_vfolder_hosts: Optional[dict[str, str]] = None
44+
integration_id: Optional[str] = None
45+
resource_policy: Optional[str] = None
46+
container_registry: Optional[dict[str, str]] = None
47+
dotfiles: Optional[bytes] = None
48+
49+
def fields_to_store(self) -> dict[str, Any]:
50+
return {
51+
"name": self.name,
52+
"domain_name": self.domain_name,
53+
"type": self.type,
54+
"description": self.description,
55+
"is_active": self.is_active,
56+
"total_resource_slots": self.total_resource_slots,
57+
"allowed_vfolder_hosts": self.allowed_vfolder_hosts,
58+
"integration_id": self.integration_id,
59+
"resource_policy": self.resource_policy,
60+
"container_registry": self.container_registry,
61+
"dotfiles": self.dotfiles,
62+
}
63+
64+
65+
@dataclass
66+
class GroupData:
67+
id: uuid.UUID = field(compare=False)
68+
name: str
69+
description: Optional[str]
70+
is_active: bool
71+
created_at: datetime = field(compare=False)
72+
modified_at: datetime = field(compare=False)
73+
integration_id: Optional[str]
74+
domain_name: str
75+
total_resource_slots: ResourceSlot
76+
allowed_vfolder_hosts: VFolderHostPermissionMap
77+
dotfiles: bytes
78+
resource_policy: str
79+
type: ProjectType
80+
container_registry: Optional[dict[str, str]]
81+
82+
def scope_id(self) -> ScopeId:
83+
return ScopeId(
84+
scope_type=ScopeType.PROJECT,
85+
scope_id=str(self.id),
86+
)
87+
88+
def role_name(self) -> str:
89+
return f"project-{str(self.id)[:8]}-admin"
90+
91+
def entity_operations(self) -> Mapping[EntityType, Iterable[OperationType]]:
92+
return {
93+
entity: OperationType.admin_operations()
94+
for entity in EntityType.admin_accessible_entity_types_in_project()
95+
}
96+
97+
98+
@dataclass
99+
class GroupModifier(PartialModifier):
100+
name: OptionalState[str] = field(default_factory=OptionalState[str].nop)
101+
description: TriState[str] = field(
102+
default_factory=TriState[str].nop,
103+
)
104+
is_active: OptionalState[bool] = field(default_factory=OptionalState[bool].nop)
105+
domain_name: OptionalState[str] = field(
106+
default_factory=OptionalState[str].nop,
107+
)
108+
total_resource_slots: OptionalState[ResourceSlot] = field(
109+
default_factory=OptionalState[ResourceSlot].nop
110+
)
111+
allowed_vfolder_hosts: OptionalState[dict[str, str]] = field(
112+
default_factory=OptionalState[dict[str, str]].nop
113+
)
114+
integration_id: OptionalState[str] = field(default_factory=OptionalState[str].nop)
115+
resource_policy: OptionalState[str] = field(default_factory=OptionalState[str].nop)
116+
container_registry: TriState[dict[str, str]] = field(
117+
default_factory=TriState[dict[str, str]].nop
118+
)
119+
120+
@override
121+
def fields_to_update(self) -> dict[str, Any]:
122+
to_update: dict[str, Any] = {}
123+
self.name.update_dict(to_update, "name")
124+
self.description.update_dict(to_update, "description")
125+
self.is_active.update_dict(to_update, "is_active")
126+
self.domain_name.update_dict(to_update, "domain_name")
127+
self.total_resource_slots.update_dict(to_update, "total_resource_slots")
128+
self.allowed_vfolder_hosts.update_dict(to_update, "allowed_vfolder_hosts")
129+
self.integration_id.update_dict(to_update, "integration_id")
130+
self.resource_policy.update_dict(to_update, "resource_policy")
131+
self.container_registry.update_dict(to_update, "container_registry")
132+
return to_update

src/ai/backend/manager/data/permission/association_scopes_entities.py

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,27 @@
11
import uuid
22
from dataclasses import dataclass
3+
from typing import Any, override
4+
5+
from ai.backend.manager.types import Creator
36

47
from .id import ObjectId, ScopeId
58

69

10+
@dataclass
11+
class AssociationScopesEntitiesCreateInput(Creator):
12+
scope_id: ScopeId
13+
object_id: ObjectId
14+
15+
@override
16+
def fields_to_store(self) -> dict[str, Any]:
17+
return {
18+
"scope_type": self.scope_id.scope_type,
19+
"scope_id": self.scope_id.scope_id,
20+
"entity_type": self.object_id.entity_type,
21+
"entity_id": self.object_id.entity_id,
22+
}
23+
24+
725
@dataclass
826
class AssociationScopesEntitiesData:
927
id: uuid.UUID

0 commit comments

Comments
 (0)