refactor(BA-2966): Explicitly raise error when operation fails in auth repository

changes/7177.enhance.md

@@ -0,0 +1 @@
+Explicitly raise error when operation fails in auth repository

docs/manager/graphql-reference/schema.graphql

@@ -128,7 +128,7 @@
     """Added in 24.03.1"""
     id: String
     reference: String
-    architecture: String = "aarch64"
+    architecture: String = "x86_64"
   ): Image
   images(
     """
@@ -2341,7 +2341,7 @@
   ): RescanImages
   preload_image(references: [String]!, target_agents: [String]!): PreloadImage
   unload_image(references: [String]!, target_agents: [String]!): UnloadImage
-  modify_image(architecture: String = "aarch64", props: ModifyImageInput!, target: String!): ModifyImage
+  modify_image(architecture: String = "x86_64", props: ModifyImageInput!, target: String!): ModifyImage
 
   """Added in 25.6.0"""
   clear_image_custom_resource_limit(key: ClearImageCustomResourceLimitKey!): ClearImageCustomResourceLimitPayload
@@ -2350,7 +2350,7 @@
   forget_image_by_id(image_id: String!): ForgetImageById
 
   """Deprecated since 25.4.0. Use `forget_image_by_id` instead."""
-  forget_image(architecture: String = "aarch64", reference: String!): ForgetImage @deprecated(reason: "Deprecated since 25.4.0. Use `forget_image_by_id` instead.")
+  forget_image(architecture: String = "x86_64", reference: String!): ForgetImage @deprecated(reason: "Deprecated since 25.4.0. Use `forget_image_by_id` instead.")
 
   """Added in 25.4.0"""
   purge_image_by_id(
@@ -2362,7 +2362,7 @@
 
   """Added in 24.03.1"""
   untag_image_from_registry(image_id: String!): UntagImageFromRegistry
-  alias_image(alias: String!, architecture: String = "aarch64", target: String!): AliasImage
+  alias_image(alias: String!, architecture: String = "x86_64", target: String!): AliasImage
   dealias_image(alias: String!): DealiasImage
   clear_images(registry: String): ClearImages
 
@@ -2937,7 +2937,7 @@
 """Added in 25.6.0."""
 input ClearImageCustomResourceLimitKey {
   image_canonical: String!
-  architecture: String! = "aarch64"
+  architecture: String! = "x86_64"
 }
 
 """Added in 24.03.0."""

docs/manager/graphql-reference/supergraph.graphql

@@ -1129,7 +1129,7 @@ input ClearImageCustomResourceLimitKey
   @join__type(graph: GRAPHENE)
 {
   image_canonical: String!
-  architecture: String! = "aarch64"
+  architecture: String! = "x86_64"
 }
 
 """Added in 25.6.0."""
@@ -4461,7 +4461,7 @@ type Mutation
   ): RescanImages @join__field(graph: GRAPHENE)
   preload_image(references: [String]!, target_agents: [String]!): PreloadImage @join__field(graph: GRAPHENE)
   unload_image(references: [String]!, target_agents: [String]!): UnloadImage @join__field(graph: GRAPHENE)
-  modify_image(architecture: String = "aarch64", props: ModifyImageInput!, target: String!): ModifyImage @join__field(graph: GRAPHENE)
+  modify_image(architecture: String = "x86_64", props: ModifyImageInput!, target: String!): ModifyImage @join__field(graph: GRAPHENE)
 
   """Added in 25.6.0"""
   clear_image_custom_resource_limit(key: ClearImageCustomResourceLimitKey!): ClearImageCustomResourceLimitPayload @join__field(graph: GRAPHENE)
@@ -4470,7 +4470,7 @@ type Mutation
   forget_image_by_id(image_id: String!): ForgetImageById @join__field(graph: GRAPHENE)
 
   """Deprecated since 25.4.0. Use `forget_image_by_id` instead."""
-  forget_image(architecture: String = "aarch64", reference: String!): ForgetImage @join__field(graph: GRAPHENE) @deprecated(reason: "Deprecated since 25.4.0. Use `forget_image_by_id` instead.")
+  forget_image(architecture: String = "x86_64", reference: String!): ForgetImage @join__field(graph: GRAPHENE) @deprecated(reason: "Deprecated since 25.4.0. Use `forget_image_by_id` instead.")
 
   """Added in 25.4.0"""
   purge_image_by_id(
@@ -4482,7 +4482,7 @@ type Mutation
 
   """Added in 24.03.1"""
   untag_image_from_registry(image_id: String!): UntagImageFromRegistry @join__field(graph: GRAPHENE)
-  alias_image(alias: String!, architecture: String = "aarch64", target: String!): AliasImage @join__field(graph: GRAPHENE)
+  alias_image(alias: String!, architecture: String = "x86_64", target: String!): AliasImage @join__field(graph: GRAPHENE)
   dealias_image(alias: String!): DealiasImage @join__field(graph: GRAPHENE)
   clear_images(registry: String): ClearImages @join__field(graph: GRAPHENE)
 
@@ -5514,7 +5514,7 @@ type Query
     """Added in 24.03.1"""
     id: String
     reference: String
-    architecture: String = "aarch64"
+    architecture: String = "x86_64"
   ): Image @join__field(graph: GRAPHENE)
   images(
     """

src/ai/backend/manager/api/auth.py

@@ -942,19 +942,14 @@ async def update_full_name(request: web.Request, params: Any) -> web.Response:
     domain_name = request["user"]["domain_name"]
     email = request["user"]["email"]
     log.info("AUTH.UPDATE_FULL_NAME(d:{}, email:{})", domain_name, email)
-    result = await root_ctx.processors.auth.update_full_name.wait_for_complete(
+    await root_ctx.processors.auth.update_full_name.wait_for_complete(
         UpdateFullNameAction(
             user_id=request["user"]["uuid"],
             full_name=params["full_name"],
             domain_name=domain_name,
             email=email,
         )
     )
-
-    if not result.success:
-        log.info("AUTH.UPDATE_FULL_NAME(d:{}, email:{}): Unknown user", domain_name, email)
-        return web.json_response({"error_msg": "Unknown user"}, status=HTTPStatus.BAD_REQUEST)
-
     return web.json_response({}, status=HTTPStatus.OK)
 
 

src/ai/backend/manager/models/user.py

@@ -23,6 +23,7 @@
 from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm
 from ai.backend.manager.data.model_serving.types import UserData as ModelServingUserData
 from ai.backend.manager.data.user.types import UserCreator, UserData, UserRole, UserStatus
+from ai.backend.manager.errors.auth import AuthorizationFailed
 from ai.backend.manager.models.hasher.types import HashInfo, PasswordInfo
 
 from .base import (
@@ -252,23 +253,6 @@ def load_main_keypair(cls) -> Callable:
     def load_resource_policy(cls) -> Callable:
         return joinedload(UserRow.resource_policy_row)
 
-    @classmethod
-    async def query_user_by_uuid(
-        cls,
-        user_uuid: UUID,
-        db_session: SASession,
-    ) -> Optional[Self]:
-        user_query = (
-            sa.select(UserRow)
-            .where(UserRow.uuid == user_uuid)
-            .options(
-                joinedload(UserRow.main_keypair),
-                selectinload(UserRow.keypairs),
-            )
-        )
-        user_row = await db_session.scalar(user_query)
-        return user_row
-
     @classmethod
     async def query_by_condition(
         cls,
@@ -440,7 +424,7 @@ async def check_credential_with_migration(
     domain: str,
     email: str,
     target_password_info: PasswordInfo,
-) -> Any:
+) -> dict:
     """
     Check user credentials and optionally migrate password hash if needed.
 
@@ -451,8 +435,12 @@ async def check_credential_with_migration(
         target_password_info: Password configuration containing password and target hash settings
 
     Returns:
-        User row if credentials are valid, None otherwise
+        User row if credentials are valid
+
+    Raises:
+        AuthorizationFailed: If user not found, password not set, or password mismatch
     """
+
     async with db.begin_readonly() as conn:
         result = await conn.execute(
             sa.select([users])
@@ -463,16 +451,15 @@ async def check_credential_with_migration(
         )
     row = result.first()
     if row is None:
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
     if row["password"] is None:
-        # user password is not set.
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
 
     try:
         if not _verify_password(target_password_info.password, row["password"]):
-            return None
+            raise AuthorizationFailed("User credential mismatch.")
     except ValueError:
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
 
     # Password is valid, check if we need to migrate the hash
     current_hash_info = HashInfo.from_hash_string(row["password"])
@@ -498,7 +485,7 @@ async def check_credential(
     domain: str,
     email: str,
     password: str,
-) -> Any:
+) -> dict[str, Any]:
     """
     Check user credentials without migration (for signout, update password, etc.)
 
@@ -509,8 +496,12 @@ async def check_credential(
         password: Plain text password to verify
 
     Returns:
-        User row if credentials are valid, None otherwise
+        User row if credentials are valid
+
+    Raises:
+        AuthorizationFailed: If user not found, password not set, or password mismatch
     """
+
     async with db.begin_readonly() as conn:
         result = await conn.execute(
             sa.select([users])
@@ -521,15 +512,14 @@ async def check_credential(
         )
     row = result.first()
     if row is None:
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
     if row["password"] is None:
-        # user password is not set.
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
 
     try:
         if not _verify_password(password, row["password"]):
-            return None
+            raise AuthorizationFailed("User credential mismatch.")
     except ValueError:
-        return None
+        raise AuthorizationFailed("User credential mismatch.")
 
     return row

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

@@ -3,12 +3,13 @@
 from __future__ import annotations
 
 from datetime import datetime
-from typing import Optional
+from typing import Any, Optional
 from uuid import UUID
 
 import sqlalchemy as sa
+from sqlalchemy.orm import joinedload, selectinload
 
-from ai.backend.common.exception import BackendAIError
+from ai.backend.common.exception import BackendAIError, UserNotFound
 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
@@ -124,8 +125,8 @@ async def insert_user_with_keypair(
             return self._user_row_to_data(user_row)
 
     @auth_db_source_resilience.apply()
-    async def modify_user_full_name(self, email: str, domain_name: str, full_name: str) -> bool:
-        """Modify user's full name in database. Returns True if updated, False if user not found."""
+    async def modify_user_full_name(self, email: str, domain_name: str, full_name: str) -> None:
+        """Modify user's full name in database."""
         async with self._db.begin() as conn:
             query = (
                 sa.select(users)
@@ -135,12 +136,11 @@ async def modify_user_full_name(self, email: str, domain_name: str, full_name: s
             result = await conn.execute(query)
             user_row = result.first()
             if not user_row:
-                return False
+                raise UserNotFound(extra_data={"email": email, "domain": domain_name})
 
             data = {"full_name": full_name}
             update_query = users.update().values(data).where(users.c.email == email)
             await conn.execute(update_query)
-            return True
 
     @auth_db_source_resilience.apply()
     async def modify_user_password(self, email: str, password_info: PasswordInfo) -> None:
@@ -239,7 +239,7 @@ async def verify_credential_with_migration(
         domain_name: str,
         email: str,
         target_password_info: PasswordInfo,
-    ) -> Optional[dict]:
+    ) -> dict[str, Any]:
         """Verify credentials with password migration support."""
         return await check_credential_with_migration(
             db=self._db,
@@ -254,7 +254,7 @@ async def verify_credential_without_migration(
         domain_name: str,
         email: str,
         password: str,
-    ) -> Optional[dict]:
+    ) -> dict[str, Any]:
         """Verify credentials without password migration (for signout, etc.)"""
         return await check_credential(
             db=self._db,
@@ -264,10 +264,21 @@ async def verify_credential_without_migration(
         )
 
     @auth_db_source_resilience.apply()
-    async def fetch_user_row_by_uuid(self, user_uuid: UUID) -> Optional[UserRow]:
+    async def fetch_user_row_by_uuid(self, user_uuid: UUID) -> UserRow:
         """Fetch user row by UUID from database."""
         async with self._db.begin_session() as db_session:
-            return await UserRow.query_user_by_uuid(user_uuid, db_session)
+            user_query = (
+                sa.select(UserRow)
+                .where(UserRow.uuid == user_uuid)
+                .options(
+                    joinedload(UserRow.main_keypair),
+                    selectinload(UserRow.keypairs),
+                )
+            )
+            user_row = await db_session.scalar(user_query)
+            if user_row is None:
+                raise UserNotFound(extra_data=user_uuid)
+            return user_row
 
     @auth_db_source_resilience.apply()
     async def fetch_current_time(self) -> datetime:

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

@@ -45,8 +45,8 @@ async def create_user_with_keypair(
         )
 
     @auth_repository_resilience.apply()
-    async def update_user_full_name(self, email: str, domain_name: str, full_name: str) -> bool:
-        return await self._db_source.modify_user_full_name(email, domain_name, full_name)
+    async def update_user_full_name(self, email: str, domain_name: str, full_name: str) -> None:
+        await self._db_source.modify_user_full_name(email, domain_name, full_name)
 
     @auth_repository_resilience.apply()
     async def update_user_password(self, email: str, password_info: PasswordInfo) -> None:
@@ -76,7 +76,7 @@ async def check_credential_with_migration(
         domain_name: str,
         email: str,
         target_password_info: PasswordInfo,
-    ) -> Optional[dict]:
+    ) -> dict:
         return await self._db_source.verify_credential_with_migration(
             domain_name, email, target_password_info
         )
@@ -87,14 +87,14 @@ async def check_credential_without_migration(
         domain_name: str,
         email: str,
         password: str,
-    ) -> Optional[dict]:
+    ) -> dict:
         """Check credentials without password migration (for signout, etc.)"""
         return await self._db_source.verify_credential_without_migration(
             domain_name, email, password
         )
 
     @auth_repository_resilience.apply()
-    async def get_user_row_by_uuid(self, user_uuid: UUID) -> Optional[UserRow]:
+    async def get_user_row_by_uuid(self, user_uuid: UUID) -> UserRow:
         return await self._db_source.fetch_user_row_by_uuid(user_uuid)
 
     @auth_repository_resilience.apply()