feat: Apply Querier pattern in get registry action

Author: seedspirit

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

@@ -5,7 +5,7 @@
 import uuid
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Mapping, MutableMapping, Optional, Self, cast
+from typing import TYPE_CHECKING, Any, Optional, Self, cast
 from urllib.parse import urlparse
 
 import graphene
@@ -230,33 +230,6 @@ async def get_container_registry_info(
 
         return yarl.URL(url), creds
 
-    @classmethod
-    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,
-                )
-            )
-            .order_by(ContainerRegistryRow.registry_name, ContainerRegistryRow.project)
-        )
-        registries = cast(list[ContainerRegistryRow], (await session.scalars(query_stmt)).all())
-        result: MutableMapping[str, MutableMapping[str, yarl.URL]] = {}
-        for registry_row in registries:
-            project = registry_row.project
-            registry_name = registry_row.registry_name
-            url = registry_row.url
-            if project not in result:
-                result[project] = {}
-            result[project][registry_name] = yarl.URL(url)
-        return result
-
     @classmethod
     def from_dataclass(cls, data: ContainerRegistryData) -> Self:
         instance = cls(

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

@@ -5,6 +5,7 @@
 from uuid import UUID
 
 import sqlalchemy as sa
+import yarl
 from sqlalchemy.ext.asyncio import AsyncSession as SASession
 
 from ai.backend.common.container_registry import AllowedGroupsModel
@@ -30,6 +31,7 @@
 )
 from ai.backend.manager.models.image import ImageRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
+from ai.backend.manager.repositories.base import Querier, execute_querier
 
 log = BraceStyleAdapter(logging.getLogger(__spec__.name))
 
@@ -198,24 +200,29 @@ async def mark_images_as_deleted(
 
             await session.execute(update_stmt)
 
-    async def fetch_known_registries(self) -> list[ContainerRegistryLocationInfo]:
+    async def fetch_known_registries(
+        self, querier: Optional[Querier]
+    ) -> 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
+            query_stmt = sa.select(
+                ContainerRegistryRow.project,
+                ContainerRegistryRow.registry_name,
+                ContainerRegistryRow.url,
+                sa.func.count().over().label("total_count"),
             )
 
+            result = await execute_querier(session, query_stmt, querier, ContainerRegistryRow)
+            rows = result.rows
             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.append(
-                            ContainerRegistryLocationInfo(
-                                registry_name=registry_name,
-                                project=project,
-                                url=url.human_repr(),
-                            )
-                        )
+            for row in rows:
+                known_registries.append(
+                    ContainerRegistryLocationInfo(
+                        registry_name=row.registry_name,
+                        project=row.project,
+                        url=yarl.URL(row.url).human_repr(),
+                    )
+                )
 
             return known_registries
 

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

@@ -0,0 +1,18 @@
+from ai.backend.manager.models.container_registry import ContainerRegistryRow
+from ai.backend.manager.repositories.base import QueryOrder
+
+
+class ContainerRegistryOrders:
+    @staticmethod
+    def project(ascending: bool = True) -> QueryOrder:
+        if ascending:
+            return ContainerRegistryRow.project.asc()
+        else:
+            return ContainerRegistryRow.project.desc()
+
+    @staticmethod
+    def registry_name(ascending: bool = True) -> QueryOrder:
+        if ascending:
+            return ContainerRegistryRow.registry_name.asc()
+        else:
+            return ContainerRegistryRow.registry_name.desc()

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

@@ -14,6 +14,7 @@
 )
 from ai.backend.manager.models.container_registry import ContainerRegistryRow
 from ai.backend.manager.models.utils import ExtendedAsyncSAEngine
+from ai.backend.manager.repositories.base import Querier
 from ai.backend.manager.repositories.container_registry.db_source.db_source import (
     ContainerRegistryDBSource,
 )
@@ -92,8 +93,10 @@ 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) -> list[ContainerRegistryLocationInfo]:
-        return await self._db_source.fetch_known_registries()
+    async def get_known_registries(
+        self, querier: Optional[Querier]
+    ) -> list[ContainerRegistryLocationInfo]:
+        return await self._db_source.fetch_known_registries(querier)
 
     @container_registry_repository_resilience.apply()
     async def get_registry_row_for_scanner(

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

@@ -5,11 +5,14 @@
 from ai.backend.manager.data.container_registry.types import (
     ContainerRegistryLocationInfo,
 )
+from ai.backend.manager.repositories.base import Querier
 from ai.backend.manager.services.container_registry.actions.base import ContainerRegistryAction
 
 
 @dataclass
 class GetContainerRegistriesAction(ContainerRegistryAction):
+    querier: Optional[Querier] = None
+
     @override
     def entity_id(self) -> Optional[str]:
         return None

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

@@ -133,7 +133,7 @@ async def load_all_container_registries(
         return LoadAllContainerRegistriesActionResult(registries=registries)
 
     async def get_container_registries(
-        self, _action: GetContainerRegistriesAction
+        self, action: GetContainerRegistriesAction
     ) -> GetContainerRegistriesActionResult:
-        registries = await self._container_registry_repository.get_known_registries()
+        registries = await self._container_registry_repository.get_known_registries(action.querier)
         return GetContainerRegistriesActionResult(registries=registries)