Skip to content

Commit d9b9271

Browse files
fregataaclaude
andcommitted
feat(BA-6063): add UserID type and ResolveUserIDByAccessKey action
Add a shared UserID identifier type under common/identifier/ and introduce ResolveUserIDByAccessKey on the auth service so any handler/service can resolve an access_key to its owning user_uuid without depending on the keypairs table directly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 761cc8f commit d9b9271

6 files changed

Lines changed: 76 additions & 0 deletions

File tree

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
from typing import NewType
2+
from uuid import UUID
3+
4+
__all__ = ("UserID",)
5+
6+
7+
UserID = NewType("UserID", UUID)

src/ai/backend/manager/repositories/auth/db_source/db_source.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
from sqlalchemy.orm import joinedload, selectinload
1313

1414
from ai.backend.common.exception import BackendAIError, UserNotFound
15+
from ai.backend.common.identifier.user import UserID
1516
from ai.backend.common.metrics.metric import DomainType, LayerType
1617
from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy
1718
from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy
@@ -21,6 +22,7 @@
2122
from ai.backend.manager.data.common.types import SearchResult
2223
from ai.backend.manager.data.permission.types import EntityType, ScopeType
2324
from ai.backend.manager.errors.auth import (
25+
AccessKeyNotFound,
2426
AuthorizationFailed,
2527
GroupMembershipNotFoundError,
2628
LoginSessionNotFoundError,
@@ -292,6 +294,16 @@ async def fetch_user_info_by_access_key(self, access_key: str) -> tuple[str, Use
292294
raise ValueError("Unknown owner access key")
293295
return row.domain_name, row.role
294296

297+
@auth_db_source_resilience.apply()
298+
async def fetch_user_id_by_access_key(self, access_key: str) -> UserID:
299+
async with self._db.begin_readonly() as conn:
300+
query = sa.select(keypairs.c.user).where(keypairs.c.access_key == access_key)
301+
result = await conn.execute(query)
302+
row = result.scalar()
303+
if row is None:
304+
raise AccessKeyNotFound("Unknown access key")
305+
return UserID(UUID(str(row)))
306+
295307
@auth_db_source_resilience.apply()
296308
async def fetch_user_info_by_email(self, email: str) -> tuple[UUID, UserRole, str]:
297309
"""Fetch (uuid, role, domain_name) for a user identified by *email*.

src/ai/backend/manager/repositories/auth/repository.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import sqlalchemy as sa
66

7+
from ai.backend.common.identifier.user import UserID
78
from ai.backend.common.metrics.metric import DomainType, LayerType
89
from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy
910
from ai.backend.common.resilience.resilience import Resilience
@@ -82,6 +83,10 @@ async def update_ssh_keypair(self, access_key: str, public_key: str, private_key
8283
async def get_delegation_target_by_access_key(self, access_key: str) -> tuple[str, UserRole]:
8384
return await self._db_source.fetch_user_info_by_access_key(access_key)
8485

86+
@auth_repository_resilience.apply()
87+
async def get_user_id_by_access_key(self, access_key: str) -> UserID:
88+
return await self._db_source.fetch_user_id_by_access_key(access_key)
89+
8590
@auth_repository_resilience.apply()
8691
async def get_delegation_target_by_email(self, email: str) -> tuple[UUID, UserRole, str]:
8792
return await self._db_source.fetch_user_info_by_email(email)
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
from dataclasses import dataclass
2+
from typing import override
3+
4+
from ai.backend.common.identifier.user import UserID
5+
from ai.backend.common.types import AccessKey
6+
from ai.backend.manager.actions.action import BaseActionResult
7+
from ai.backend.manager.actions.types import ActionOperationType
8+
from ai.backend.manager.services.auth.actions.base import AuthAction
9+
10+
11+
@dataclass
12+
class ResolveUserIDByAccessKeyAction(AuthAction):
13+
access_key: AccessKey
14+
15+
@override
16+
def entity_id(self) -> str | None:
17+
return str(self.access_key)
18+
19+
@override
20+
@classmethod
21+
def operation_type(cls) -> ActionOperationType:
22+
return ActionOperationType.GET
23+
24+
25+
@dataclass
26+
class ResolveUserIDByAccessKeyResult(BaseActionResult):
27+
user_id: UserID
28+
29+
@override
30+
def entity_id(self) -> str | None:
31+
return str(self.user_id)

src/ai/backend/manager/services/auth/processors.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@
2424
ResolveAccessKeyScopeAction,
2525
ResolveAccessKeyScopeResult,
2626
)
27+
from ai.backend.manager.services.auth.actions.resolve_user_id_by_access_key import (
28+
ResolveUserIDByAccessKeyAction,
29+
ResolveUserIDByAccessKeyResult,
30+
)
2731
from ai.backend.manager.services.auth.actions.resolve_user_scope import (
2832
ResolveUserScopeAction,
2933
ResolveUserScopeResult,
@@ -88,6 +92,9 @@ class AuthProcessors(AbstractProcessorPackage):
8892
ResolveAccessKeyScopeAction, ResolveAccessKeyScopeResult
8993
]
9094
resolve_user_scope: ActionProcessor[ResolveUserScopeAction, ResolveUserScopeResult]
95+
resolve_user_id_by_access_key: ActionProcessor[
96+
ResolveUserIDByAccessKeyAction, ResolveUserIDByAccessKeyResult
97+
]
9198
admin_search_login_sessions: ActionProcessor[
9299
AdminSearchLoginSessionsAction, SearchLoginSessionsActionResult
93100
]
@@ -135,6 +142,9 @@ def __init__(
135142
service.resolve_access_key_scope, action_monitors
136143
)
137144
self.resolve_user_scope = ActionProcessor(service.resolve_user_scope, action_monitors)
145+
self.resolve_user_id_by_access_key = ActionProcessor(
146+
service.resolve_user_id_by_access_key, action_monitors
147+
)
138148
self.admin_search_login_sessions = ActionProcessor(
139149
service.admin_search_login_sessions, action_monitors
140150
)
@@ -167,6 +177,7 @@ def supported_actions(self) -> list[ActionSpec]:
167177
UpdatePasswordNoAuthAction.spec(),
168178
ResolveAccessKeyScopeAction.spec(),
169179
ResolveUserScopeAction.spec(),
180+
ResolveUserIDByAccessKeyAction.spec(),
170181
AdminSearchLoginSessionsAction.spec(),
171182
SearchLoginSessionsAction.spec(),
172183
AdminSearchLoginHistoryAction.spec(),

src/ai/backend/manager/services/auth/service.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,10 @@
7777
ResolveAccessKeyScopeAction,
7878
ResolveAccessKeyScopeResult,
7979
)
80+
from ai.backend.manager.services.auth.actions.resolve_user_id_by_access_key import (
81+
ResolveUserIDByAccessKeyAction,
82+
ResolveUserIDByAccessKeyResult,
83+
)
8084
from ai.backend.manager.services.auth.actions.resolve_user_scope import (
8185
ResolveUserScopeAction,
8286
ResolveUserScopeResult,
@@ -758,6 +762,12 @@ async def resolve_access_key_scope(
758762
owner_access_key=owner_ak,
759763
)
760764

765+
async def resolve_user_id_by_access_key(
766+
self, action: ResolveUserIDByAccessKeyAction
767+
) -> ResolveUserIDByAccessKeyResult:
768+
user_id = await self._auth_repository.get_user_id_by_access_key(str(action.access_key))
769+
return ResolveUserIDByAccessKeyResult(user_id=user_id)
770+
761771
async def resolve_user_scope(self, action: ResolveUserScopeAction) -> ResolveUserScopeResult:
762772
if action.owner_user_email is None:
763773
return ResolveUserScopeResult(

0 commit comments

Comments
 (0)