refactor/prework to refactor direct database access to follow dir pattern

Author: kwonkwonn

src/ai/backend/manager/data/auth/types.py

@@ -1,8 +1,16 @@
+from __future__ import annotations
+
 import uuid
 from dataclasses import dataclass
 from datetime import datetime
 from typing import Optional
 
+from ai.backend.manager.data.keypair.types import KeyPairData
+from ai.backend.manager.data.resource.types import (
+    KeyPairResourcePolicyData,
+    UserResourcePolicyData,
+)
+from ai.backend.manager.data.user.types import UserData as FullUserData
 from ai.backend.manager.models.user import UserRole, UserStatus
 
 
@@ -47,3 +55,28 @@ class UserData:
 class GroupMembershipData:
     group_id: uuid.UUID
     user_id: uuid.UUID
+
+
+@dataclass
+class CredentialData:
+    """
+    Aggregated credential data for authentication.
+
+    Combines user, keypair, and their resource policies into a single data structure.
+    Used by authentication middleware to populate request context.
+    """
+
+    user: FullUserData
+    user_resource_policy: UserResourcePolicyData
+    keypair: KeyPairData
+    keypair_resource_policy: KeyPairResourcePolicyData
+
+    @property
+    def is_admin(self) -> bool:
+        """Check if the keypair has admin privileges."""
+        return self.keypair.is_admin
+
+    @property
+    def is_superadmin(self) -> bool:
+        """Check if the user has superadmin role."""
+        return self.user.role == UserRole.SUPERADMIN

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

@@ -14,11 +14,11 @@
 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.manager.data.auth.types import GroupMembershipData, UserData
+from ai.backend.manager.data.auth.types import CredentialData, GroupMembershipData, UserData
 from ai.backend.manager.errors.auth import GroupMembershipNotFoundError, UserCreationError
 from ai.backend.manager.models.group import association_groups_users, groups
 from ai.backend.manager.models.hasher.types import PasswordInfo
-from ai.backend.manager.models.keypair import keypairs
+from ai.backend.manager.models.keypair import KeyPairRow, keypairs
 from ai.backend.manager.models.user import (
     UserRow,
     UserStatus,
@@ -285,3 +285,43 @@ async def fetch_current_time(self) -> datetime:
         """Fetch current time from database."""
         async with self._db.begin_readonly() as db_conn:
             return await db_conn.scalar(sa.select(sa.func.now()))
+
+    @auth_db_source_resilience.apply()
+    async def fetch_credential_by_access_key(self, access_key: str) -> Optional[CredentialData]:
+        """
+        Fetch user credential data by access key.
+
+        Queries keypair with user and resource policies using ORM relationships.
+        Returns None if access key not found or inactive.
+
+        Args:
+            access_key: The access key to look up.
+
+        Returns:
+            CredentialData containing user, keypair, and resource policies,
+            or None if not found.
+        """
+        async with self._db.begin_session() as db_session:
+            query = (
+                sa.select(KeyPairRow)
+                .where((KeyPairRow.access_key == access_key) & (KeyPairRow.is_active.is_(True)))
+                .options(
+                    joinedload(KeyPairRow.resource_policy_row),
+                    joinedload(KeyPairRow.user_row).joinedload(UserRow.resource_policy_row),
+                )
+            )
+            keypair_row = await db_session.scalar(query)
+
+            if keypair_row is None:
+                return None
+
+            user_row = keypair_row.user_row
+            if user_row is None:
+                return None
+
+            return CredentialData(
+                user=user_row.to_data(),
+                user_resource_policy=user_row.resource_policy_row.to_dataclass(),
+                keypair=keypair_row.to_data(),
+                keypair_resource_policy=keypair_row.resource_policy_row.to_dataclass(),
+            )

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

@@ -5,7 +5,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.manager.data.auth.types import GroupMembershipData, UserData
+from ai.backend.manager.data.auth.types import CredentialData, GroupMembershipData, UserData
 from ai.backend.manager.models.hasher.types import PasswordInfo
 from ai.backend.manager.models.user import UserRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
@@ -100,3 +100,19 @@ async def get_user_row_by_uuid(self, user_uuid: UUID) -> UserRow:
     @auth_repository_resilience.apply()
     async def get_current_time(self) -> datetime:
         return await self._db_source.fetch_current_time()
+
+    @auth_repository_resilience.apply()
+    async def get_credential_by_access_key(self, access_key: str) -> Optional[CredentialData]:
+        """
+        Get user credential data by access key.
+
+        Used by authentication middleware to populate request context.
+
+        Args:
+            access_key: The access key to look up.
+
+        Returns:
+            CredentialData containing user, keypair, and resource policies,
+            or None if not found or inactive.
+        """
+        return await self._db_source.fetch_credential_by_access_key(access_key)