refactor: Make known registry return dataclass, not dict

Author: seedspirit

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

@@ -392,8 +392,11 @@ async def get_container_registries(request: web.Request) -> web.Response:
             GetContainerRegistriesAction()
         )
     )
+    response: dict[str, str] = {}
+    for registry in result.registries:
+        response[f"{registry.project}/{registry.registry_name}"] = registry.url
 
-    return web.json_response(result.registries, status=HTTPStatus.OK)
+    return web.json_response(response, status=HTTPStatus.OK)
 
 
 def create_app(

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

@@ -73,3 +73,10 @@ def fields_to_update(self) -> dict[str, Any]:
         self.ssl_verify.update_dict(to_update, "ssl_verify")
         self.extra.update_dict(to_update, "extra")
         return to_update
+
+
+@dataclass
+class ContainerRegistryLocationInfo:
+    project: Optional[str]
+    registry_name: str
+    url: str

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

@@ -235,12 +235,16 @@ async def get_known_container_registries(
         cls,
         session: AsyncSession,
     ) -> Mapping[str, Mapping[str, yarl.URL]]:
-        query_stmt = sa.select(ContainerRegistryRow).options(
-            load_only(
-                ContainerRegistryRow.project,
-                ContainerRegistryRow.registry_name,
-                ContainerRegistryRow.url,
+        query_stmt = (
+            sa.select(ContainerRegistryRow)
+            .options(
+                load_only(
+                    ContainerRegistryRow.project,
+                    ContainerRegistryRow.registry_name,
+                    ContainerRegistryRow.url,
+                )
             )
+            .order_by(ContainerRegistryRow.registry_name, ContainerRegistryRow.project)
         )
         registries = cast(list[ContainerRegistryRow], (await session.scalars(query_stmt)).all())
         result: MutableMapping[str, MutableMapping[str, yarl.URL]] = {}

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

@@ -12,6 +12,7 @@
 from ai.backend.manager.data.container_registry.types import (
     ContainerRegistryCreator,
     ContainerRegistryData,
+    ContainerRegistryLocationInfo,
     ContainerRegistryModifier,
 )
 from ai.backend.manager.data.image.types import ImageStatus
@@ -197,18 +198,24 @@ async def mark_images_as_deleted(
 
             await session.execute(update_stmt)
 
-    async def fetch_known_registries(self) -> dict[str, str]:
+    async def fetch_known_registries(self) -> list[ContainerRegistryLocationInfo]:
         """Fetch all known container registries from the database."""
         async with self._db.begin_readonly_session() as session:
             known_registries_map = await ContainerRegistryRow.get_known_container_registries(
                 session
             )
 
-            known_registries = {}
+            known_registries: list[ContainerRegistryLocationInfo] = []
             for project, registries in known_registries_map.items():
                 for registry_name, url in registries.items():
                     if project not in known_registries:
-                        known_registries[f"{project}/{registry_name}"] = url.human_repr()
+                        known_registries.append(
+                            ContainerRegistryLocationInfo(
+                                registry_name=registry_name,
+                                project=project,
+                                url=url.human_repr(),
+                            )
+                        )
 
             return known_registries
 

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

@@ -9,6 +9,7 @@
 from ai.backend.manager.data.container_registry.types import (
     ContainerRegistryCreator,
     ContainerRegistryData,
+    ContainerRegistryLocationInfo,
     ContainerRegistryModifier,
 )
 from ai.backend.manager.models.container_registry import ContainerRegistryRow
@@ -91,7 +92,7 @@ async def clear_images(
         return await self._db_source.fetch_by_registry_and_project(registry_name, project)
 
     @container_registry_repository_resilience.apply()
-    async def get_known_registries(self) -> dict[str, str]:
+    async def get_known_registries(self) -> list[ContainerRegistryLocationInfo]:
         return await self._db_source.fetch_known_registries()
 
     @container_registry_repository_resilience.apply()

src/ai/backend/manager/services/container_registry/actions/get_container_registries.py

@@ -1,7 +1,10 @@
 from dataclasses import dataclass
-from typing import Any, Optional, override
+from typing import Optional, override
 
 from ai.backend.manager.actions.action import BaseActionResult
+from ai.backend.manager.data.container_registry.types import (
+    ContainerRegistryLocationInfo,
+)
 from ai.backend.manager.services.container_registry.actions.base import ContainerRegistryAction
 
 
@@ -19,7 +22,7 @@ def operation_type(cls) -> str:
 
 @dataclass
 class GetContainerRegistriesActionResult(BaseActionResult):
-    registries: Any
+    registries: list[ContainerRegistryLocationInfo]
 
     @override
     def entity_id(self) -> Optional[str]:

tests/manager/integration/services/container_registry/test_container_registry_integration.py

@@ -323,7 +323,7 @@ async def test_get_container_registries_integration(
         # For global registry, actually set project to None in the database
         registry_id = None
         async with database_engine.begin_session() as session:
-            registry = ContainerRegistryRow(
+            registry_row = ContainerRegistryRow(
                 url="https://registry3.example.com",
                 registry_name="registry3",
                 type=ContainerRegistryType.DOCKER,
@@ -334,9 +334,9 @@ async def test_get_container_registries_integration(
                 is_global=True,
                 extra=None,
             )
-            session.add(registry)
+            session.add(registry_row)
             await session.commit()
-            registry_id = registry.id
+            registry_id = registry_row.id
 
         # Action: Get known registries
         action = GetContainerRegistriesAction()
@@ -345,12 +345,15 @@ async def test_get_container_registries_integration(
         )
 
         # Verify: Should return mapping of project/registry to URL
-        assert "projectA/registry1" in result.registries
-        assert result.registries["projectA/registry1"] == "https://registry1.example.com/"
-        assert "projectB/registry2" in result.registries
-        assert result.registries["projectB/registry2"] == "https://registry2.example.com/"
+        response = {}
+        for registry in result.registries:
+            response[f"{registry.project}/{registry.registry_name}"] = registry.url
+        assert "projectA/registry1" in response
+        assert response["projectA/registry1"] == "https://registry1.example.com/"
+        assert "projectB/registry2" in response
+        assert response["projectB/registry2"] == "https://registry2.example.com/"
         # Global registry should be included with None prefix
-        assert "None/registry3" in result.registries
+        assert "None/registry3" in response
 
         # Cleanup the global registry manually
         if registry_id:

tests/manager/services/container_registry/actions/test_get_container_registries.py

@@ -1,5 +1,8 @@
 import pytest
 
+from ai.backend.manager.data.container_registry.types import (
+    ContainerRegistryLocationInfo,
+)
 from ai.backend.manager.services.container_registry.actions.get_container_registries import (
     GetContainerRegistriesAction,
     GetContainerRegistriesActionResult,
@@ -19,14 +22,32 @@
             "Success Case",
             GetContainerRegistriesAction(),
             GetContainerRegistriesActionResult(
-                registries={
-                    # fixtures from fixtures.py
-                    "test_project/registry.example.com": "https://registry.example.com/",
+                registries=[
                     # fixtures from example-container-registries-harbor.json
-                    "community/cr.backend.ai": "https://cr.backend.ai/",
-                    "multiarch/cr.backend.ai": "https://cr.backend.ai/",
-                    "stable/cr.backend.ai": "https://cr.backend.ai/",
-                }
+                    # Sorted by registry_name first (cr.backend.ai < registry.example.com),
+                    # then by project (community < multiarch < stable)
+                    ContainerRegistryLocationInfo(
+                        project="community",
+                        registry_name="cr.backend.ai",
+                        url="https://cr.backend.ai/",
+                    ),
+                    ContainerRegistryLocationInfo(
+                        project="multiarch",
+                        registry_name="cr.backend.ai",
+                        url="https://cr.backend.ai/",
+                    ),
+                    ContainerRegistryLocationInfo(
+                        project="stable",
+                        registry_name="cr.backend.ai",
+                        url="https://cr.backend.ai/",
+                    ),
+                    # fixtures from fixtures.py
+                    ContainerRegistryLocationInfo(
+                        project="test_project",
+                        registry_name="registry.example.com",
+                        url="https://registry.example.com/",
+                    ),
+                ],
             ),
         ),
     ],