From d9b92711ec9f7d1708488df441de6b93f1655755 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Mon, 18 May 2026 12:29:26 +0900 Subject: [PATCH 1/3] 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) --- src/ai/backend/common/identifier/user.py | 7 +++++ .../repositories/auth/db_source/db_source.py | 12 +++++++ .../manager/repositories/auth/repository.py | 5 +++ .../actions/resolve_user_id_by_access_key.py | 31 +++++++++++++++++++ .../manager/services/auth/processors.py | 11 +++++++ .../backend/manager/services/auth/service.py | 10 ++++++ 6 files changed, 76 insertions(+) create mode 100644 src/ai/backend/common/identifier/user.py create mode 100644 src/ai/backend/manager/services/auth/actions/resolve_user_id_by_access_key.py diff --git a/src/ai/backend/common/identifier/user.py b/src/ai/backend/common/identifier/user.py new file mode 100644 index 00000000000..809300951d9 --- /dev/null +++ b/src/ai/backend/common/identifier/user.py @@ -0,0 +1,7 @@ +from typing import NewType +from uuid import UUID + +__all__ = ("UserID",) + + +UserID = NewType("UserID", UUID) diff --git a/src/ai/backend/manager/repositories/auth/db_source/db_source.py b/src/ai/backend/manager/repositories/auth/db_source/db_source.py index 56ebf53856c..d71c4d39d43 100644 --- a/src/ai/backend/manager/repositories/auth/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/auth/db_source/db_source.py @@ -12,6 +12,7 @@ from sqlalchemy.orm import joinedload, selectinload from ai.backend.common.exception import BackendAIError, UserNotFound +from ai.backend.common.identifier.user import UserID from ai.backend.common.metrics.metric import DomainType, LayerType from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy @@ -21,6 +22,7 @@ from ai.backend.manager.data.common.types import SearchResult from ai.backend.manager.data.permission.types import EntityType, ScopeType from ai.backend.manager.errors.auth import ( + AccessKeyNotFound, AuthorizationFailed, GroupMembershipNotFoundError, LoginSessionNotFoundError, @@ -292,6 +294,16 @@ async def fetch_user_info_by_access_key(self, access_key: str) -> tuple[str, Use raise ValueError("Unknown owner access key") return row.domain_name, row.role + @auth_db_source_resilience.apply() + async def fetch_user_id_by_access_key(self, access_key: str) -> UserID: + async with self._db.begin_readonly() as conn: + query = sa.select(keypairs.c.user).where(keypairs.c.access_key == access_key) + result = await conn.execute(query) + row = result.scalar() + if row is None: + raise AccessKeyNotFound("Unknown access key") + return UserID(UUID(str(row))) + @auth_db_source_resilience.apply() async def fetch_user_info_by_email(self, email: str) -> tuple[UUID, UserRole, str]: """Fetch (uuid, role, domain_name) for a user identified by *email*. diff --git a/src/ai/backend/manager/repositories/auth/repository.py b/src/ai/backend/manager/repositories/auth/repository.py index 8c3e79c42b7..b882ad8a570 100644 --- a/src/ai/backend/manager/repositories/auth/repository.py +++ b/src/ai/backend/manager/repositories/auth/repository.py @@ -4,6 +4,7 @@ import sqlalchemy as sa +from ai.backend.common.identifier.user import UserID from ai.backend.common.metrics.metric import DomainType, LayerType from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy 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 async def get_delegation_target_by_access_key(self, access_key: str) -> tuple[str, UserRole]: return await self._db_source.fetch_user_info_by_access_key(access_key) + @auth_repository_resilience.apply() + async def get_user_id_by_access_key(self, access_key: str) -> UserID: + return await self._db_source.fetch_user_id_by_access_key(access_key) + @auth_repository_resilience.apply() async def get_delegation_target_by_email(self, email: str) -> tuple[UUID, UserRole, str]: return await self._db_source.fetch_user_info_by_email(email) diff --git a/src/ai/backend/manager/services/auth/actions/resolve_user_id_by_access_key.py b/src/ai/backend/manager/services/auth/actions/resolve_user_id_by_access_key.py new file mode 100644 index 00000000000..0c8fe23038d --- /dev/null +++ b/src/ai/backend/manager/services/auth/actions/resolve_user_id_by_access_key.py @@ -0,0 +1,31 @@ +from dataclasses import dataclass +from typing import override + +from ai.backend.common.identifier.user import UserID +from ai.backend.common.types import AccessKey +from ai.backend.manager.actions.action import BaseActionResult +from ai.backend.manager.actions.types import ActionOperationType +from ai.backend.manager.services.auth.actions.base import AuthAction + + +@dataclass +class ResolveUserIDByAccessKeyAction(AuthAction): + access_key: AccessKey + + @override + def entity_id(self) -> str | None: + return str(self.access_key) + + @override + @classmethod + def operation_type(cls) -> ActionOperationType: + return ActionOperationType.GET + + +@dataclass +class ResolveUserIDByAccessKeyResult(BaseActionResult): + user_id: UserID + + @override + def entity_id(self) -> str | None: + return str(self.user_id) diff --git a/src/ai/backend/manager/services/auth/processors.py b/src/ai/backend/manager/services/auth/processors.py index 608f1a6b27c..5da695a4cb3 100644 --- a/src/ai/backend/manager/services/auth/processors.py +++ b/src/ai/backend/manager/services/auth/processors.py @@ -24,6 +24,10 @@ ResolveAccessKeyScopeAction, ResolveAccessKeyScopeResult, ) +from ai.backend.manager.services.auth.actions.resolve_user_id_by_access_key import ( + ResolveUserIDByAccessKeyAction, + ResolveUserIDByAccessKeyResult, +) from ai.backend.manager.services.auth.actions.resolve_user_scope import ( ResolveUserScopeAction, ResolveUserScopeResult, @@ -88,6 +92,9 @@ class AuthProcessors(AbstractProcessorPackage): ResolveAccessKeyScopeAction, ResolveAccessKeyScopeResult ] resolve_user_scope: ActionProcessor[ResolveUserScopeAction, ResolveUserScopeResult] + resolve_user_id_by_access_key: ActionProcessor[ + ResolveUserIDByAccessKeyAction, ResolveUserIDByAccessKeyResult + ] admin_search_login_sessions: ActionProcessor[ AdminSearchLoginSessionsAction, SearchLoginSessionsActionResult ] @@ -135,6 +142,9 @@ def __init__( service.resolve_access_key_scope, action_monitors ) self.resolve_user_scope = ActionProcessor(service.resolve_user_scope, action_monitors) + self.resolve_user_id_by_access_key = ActionProcessor( + service.resolve_user_id_by_access_key, action_monitors + ) self.admin_search_login_sessions = ActionProcessor( service.admin_search_login_sessions, action_monitors ) @@ -167,6 +177,7 @@ def supported_actions(self) -> list[ActionSpec]: UpdatePasswordNoAuthAction.spec(), ResolveAccessKeyScopeAction.spec(), ResolveUserScopeAction.spec(), + ResolveUserIDByAccessKeyAction.spec(), AdminSearchLoginSessionsAction.spec(), SearchLoginSessionsAction.spec(), AdminSearchLoginHistoryAction.spec(), diff --git a/src/ai/backend/manager/services/auth/service.py b/src/ai/backend/manager/services/auth/service.py index df31d73b94d..a5cc0c9c96e 100644 --- a/src/ai/backend/manager/services/auth/service.py +++ b/src/ai/backend/manager/services/auth/service.py @@ -77,6 +77,10 @@ ResolveAccessKeyScopeAction, ResolveAccessKeyScopeResult, ) +from ai.backend.manager.services.auth.actions.resolve_user_id_by_access_key import ( + ResolveUserIDByAccessKeyAction, + ResolveUserIDByAccessKeyResult, +) from ai.backend.manager.services.auth.actions.resolve_user_scope import ( ResolveUserScopeAction, ResolveUserScopeResult, @@ -758,6 +762,12 @@ async def resolve_access_key_scope( owner_access_key=owner_ak, ) + async def resolve_user_id_by_access_key( + self, action: ResolveUserIDByAccessKeyAction + ) -> ResolveUserIDByAccessKeyResult: + user_id = await self._auth_repository.get_user_id_by_access_key(str(action.access_key)) + return ResolveUserIDByAccessKeyResult(user_id=user_id) + async def resolve_user_scope(self, action: ResolveUserScopeAction) -> ResolveUserScopeResult: if action.owner_user_email is None: return ResolveUserScopeResult( From 6d7a739894035e160b3bc5f098499ff7a2ce0a66 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Mon, 18 May 2026 12:33:45 +0900 Subject: [PATCH 2/3] changelog: add news fragment for PR #11647 --- changes/11647.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/11647.feature.md diff --git a/changes/11647.feature.md b/changes/11647.feature.md new file mode 100644 index 00000000000..173459ba510 --- /dev/null +++ b/changes/11647.feature.md @@ -0,0 +1 @@ +Add `UserID` identifier type and `ResolveUserIDByAccessKey` auth action for resolving an access_key to its owning user UUID. From e156f12e05054052e3eb6020c4f4279dde246bf4 Mon Sep 17 00:00:00 2001 From: Sanghun Lee Date: Mon, 18 May 2026 13:09:06 +0900 Subject: [PATCH 3/3] refactor(BA-6063): use AccessKey type end-to-end and add db_source tests Tighten the resolve-user-id-by-access-key chain so the repository and db_source both accept AccessKey instead of plain str (matching the action field), drop the redundant UUID(str(row)) re-parse on the GUID-mapped keypairs.user column, and add success / not-found unit coverage on AuthRepository. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../repositories/auth/db_source/db_source.py | 5 +++-- .../manager/repositories/auth/repository.py | 3 ++- .../backend/manager/services/auth/service.py | 2 +- .../repositories/auth/test_auth_repository.py | 21 +++++++++++++++++-- 4 files changed, 25 insertions(+), 6 deletions(-) diff --git a/src/ai/backend/manager/repositories/auth/db_source/db_source.py b/src/ai/backend/manager/repositories/auth/db_source/db_source.py index d71c4d39d43..75fd3d5c245 100644 --- a/src/ai/backend/manager/repositories/auth/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/auth/db_source/db_source.py @@ -17,6 +17,7 @@ from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy from ai.backend.common.resilience.policies.retry import BackoffStrategy, RetryArgs, RetryPolicy from ai.backend.common.resilience.resilience import Resilience +from ai.backend.common.types import AccessKey from ai.backend.manager.data.auth.login_session_types import LoginHistoryData, LoginSessionData from ai.backend.manager.data.auth.types import GroupMembershipData, UserData from ai.backend.manager.data.common.types import SearchResult @@ -295,14 +296,14 @@ async def fetch_user_info_by_access_key(self, access_key: str) -> tuple[str, Use return row.domain_name, row.role @auth_db_source_resilience.apply() - async def fetch_user_id_by_access_key(self, access_key: str) -> UserID: + async def fetch_user_id_by_access_key(self, access_key: AccessKey) -> UserID: async with self._db.begin_readonly() as conn: query = sa.select(keypairs.c.user).where(keypairs.c.access_key == access_key) result = await conn.execute(query) row = result.scalar() if row is None: raise AccessKeyNotFound("Unknown access key") - return UserID(UUID(str(row))) + return UserID(cast(UUID, row)) @auth_db_source_resilience.apply() async def fetch_user_info_by_email(self, email: str) -> tuple[UUID, UserRole, str]: diff --git a/src/ai/backend/manager/repositories/auth/repository.py b/src/ai/backend/manager/repositories/auth/repository.py index b882ad8a570..20da8f1ff0b 100644 --- a/src/ai/backend/manager/repositories/auth/repository.py +++ b/src/ai/backend/manager/repositories/auth/repository.py @@ -8,6 +8,7 @@ from ai.backend.common.metrics.metric import DomainType, LayerType from ai.backend.common.resilience.policies.metrics import MetricArgs, MetricPolicy from ai.backend.common.resilience.resilience import Resilience +from ai.backend.common.types import AccessKey from ai.backend.manager.data.auth.login_session_types import LoginHistoryData, LoginSessionData from ai.backend.manager.data.auth.types import GroupMembershipData, UserData from ai.backend.manager.data.common.types import SearchResult @@ -84,7 +85,7 @@ async def get_delegation_target_by_access_key(self, access_key: str) -> tuple[st return await self._db_source.fetch_user_info_by_access_key(access_key) @auth_repository_resilience.apply() - async def get_user_id_by_access_key(self, access_key: str) -> UserID: + async def get_user_id_by_access_key(self, access_key: AccessKey) -> UserID: return await self._db_source.fetch_user_id_by_access_key(access_key) @auth_repository_resilience.apply() diff --git a/src/ai/backend/manager/services/auth/service.py b/src/ai/backend/manager/services/auth/service.py index a5cc0c9c96e..54cb241e525 100644 --- a/src/ai/backend/manager/services/auth/service.py +++ b/src/ai/backend/manager/services/auth/service.py @@ -765,7 +765,7 @@ async def resolve_access_key_scope( async def resolve_user_id_by_access_key( self, action: ResolveUserIDByAccessKeyAction ) -> ResolveUserIDByAccessKeyResult: - user_id = await self._auth_repository.get_user_id_by_access_key(str(action.access_key)) + user_id = await self._auth_repository.get_user_id_by_access_key(action.access_key) return ResolveUserIDByAccessKeyResult(user_id=user_id) async def resolve_user_scope(self, action: ResolveUserScopeAction) -> ResolveUserScopeResult: diff --git a/tests/unit/manager/repositories/auth/test_auth_repository.py b/tests/unit/manager/repositories/auth/test_auth_repository.py index 23937f65f57..1c30985af2b 100644 --- a/tests/unit/manager/repositories/auth/test_auth_repository.py +++ b/tests/unit/manager/repositories/auth/test_auth_repository.py @@ -15,12 +15,12 @@ from ai.backend.common.data.permission.types import RelationType from ai.backend.common.exception import UserNotFound -from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap +from ai.backend.common.types import AccessKey, ResourceSlot, VFolderHostPermissionMap from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.auth.types import UserData from ai.backend.manager.data.group.types import GroupData from ai.backend.manager.data.permission.types import EntityType, ScopeType -from ai.backend.manager.errors.auth import GroupMembershipNotFoundError +from ai.backend.manager.errors.auth import AccessKeyNotFound, GroupMembershipNotFoundError from ai.backend.manager.models.agent import AgentRow from ai.backend.manager.models.deployment_auto_scaling_policy import DeploymentAutoScalingPolicyRow from ai.backend.manager.models.deployment_policy import DeploymentPolicyRow @@ -533,3 +533,20 @@ async def test_get_current_time(self, auth_repository: AuthRepository) -> None: now_utc = datetime.now(UTC) time_diff = abs((now_utc - result).total_seconds()) assert time_diff < 1.0 + + async def test_get_user_id_by_access_key_success( + self, + auth_repository: AuthRepository, + sample_user_data: UserTestData, + ) -> None: + result = await auth_repository.get_user_id_by_access_key( + AccessKey(sample_user_data.access_key) + ) + + assert result == sample_user_data.uuid + + async def test_get_user_id_by_access_key_not_found( + self, auth_repository: AuthRepository + ) -> None: + with pytest.raises(AccessKeyNotFound): + await auth_repository.get_user_id_by_access_key(AccessKey("AKIANONEXISTENT"))