refactor(BA-5816): migrate session_template off association_groups_users (#11284)

Author: fregataa

changes/11284.enhance.md

@@ -0,0 +1 @@
+Migrate session and cluster template code paths off the legacy `association_groups_users` table; project membership is now validated against `association_scopes_entities` and the REST handlers resolve `(domain, group_name) → project_id` upstream via a new `GroupService` entry point.

src/ai/backend/manager/api/rest/cluster_template/handler.py

@@ -38,10 +38,14 @@
     ListClusterTemplatesResponse,
     UpdateClusterTemplateResponse,
 )
+from ai.backend.common.identifier.project import ProjectID
 from ai.backend.common.json import load_json
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.dto.context import RequestCtx, UserContext
 from ai.backend.manager.errors.api import InvalidAPIParameters
+from ai.backend.manager.services.group.actions.resolve_project_id_by_name import (
+    ResolveProjectIdByNameAction,
+)
 from ai.backend.manager.services.template.actions.create_cluster_template import (
     CreateClusterTemplateAction,
 )
@@ -59,6 +63,7 @@
 )
 
 if TYPE_CHECKING:
+    from ai.backend.manager.services.group.processors import GroupProcessors
     from ai.backend.manager.services.template.processors import TemplateProcessors
 
 log: Final = BraceStyleAdapter(logging.getLogger(__spec__.name))
@@ -67,8 +72,24 @@
 class ClusterTemplateHandler:
     """Cluster template API handler with constructor-injected dependencies."""
 
-    def __init__(self, *, template: TemplateProcessors) -> None:
+    def __init__(
+        self,
+        *,
+        template: TemplateProcessors,
+        group: GroupProcessors,
+    ) -> None:
         self._template = template
+        self._group = group
+
+    async def _resolve_project_id(self, domain_name: str, project_name: str) -> ProjectID:
+        result = await self._group.resolve_project_id_by_name.wait_for_complete(
+            ResolveProjectIdByNameAction(domain_name=domain_name, project_name=project_name)
+        )
+        if result.project_id is None:
+            raise InvalidAPIParameters(
+                f"No active group named {project_name!r} exists in domain {domain_name!r}"
+            )
+        return result.project_id
 
     async def create(
         self,
@@ -93,9 +114,10 @@ async def create(
             except (yaml.YAMLError, yaml.MarkedYAMLError) as e:
                 raise InvalidAPIParameters("Malformed payload") from e
 
+        project_id = await self._resolve_project_id(domain, params.group)
         action = CreateClusterTemplateAction(
             domain_name=domain,
-            requesting_group=params.group,
+            requesting_project=project_id,
             requester_uuid=ctx.user_uuid,
             requester_access_key=ctx.access_key,
             requester_role=req.request["user"]["role"],

src/ai/backend/manager/api/rest/session_template/handler.py

@@ -38,10 +38,14 @@
     SessionTemplateListItemDTO,
     UpdateSessionTemplateResponse,
 )
+from ai.backend.common.identifier.project import ProjectID
 from ai.backend.common.json import load_json
 from ai.backend.logging import BraceStyleAdapter
 from ai.backend.manager.dto.context import RequestCtx, UserContext
 from ai.backend.manager.errors.api import InvalidAPIParameters
+from ai.backend.manager.services.group.actions.resolve_project_id_by_name import (
+    ResolveProjectIdByNameAction,
+)
 from ai.backend.manager.services.template.actions.create_task_template import (
     CreateTaskTemplateAction,
     TaskTemplateItemInput,
@@ -60,6 +64,7 @@
 )
 
 if TYPE_CHECKING:
+    from ai.backend.manager.services.group.processors import GroupProcessors
     from ai.backend.manager.services.template.processors import TemplateProcessors
 
 log: Final = BraceStyleAdapter(logging.getLogger(__spec__.name))
@@ -68,8 +73,24 @@
 class SessionTemplateHandler:
     """Session template API handler with constructor-injected dependencies."""
 
-    def __init__(self, *, template: TemplateProcessors) -> None:
+    def __init__(
+        self,
+        *,
+        template: TemplateProcessors,
+        group: GroupProcessors,
+    ) -> None:
         self._template = template
+        self._group = group
+
+    async def _resolve_project_id(self, domain_name: str, project_name: str) -> ProjectID:
+        result = await self._group.resolve_project_id_by_name.wait_for_complete(
+            ResolveProjectIdByNameAction(domain_name=domain_name, project_name=project_name)
+        )
+        if result.project_id is None:
+            raise InvalidAPIParameters(
+                f"No active group named {project_name!r} exists in domain {domain_name!r}"
+            )
+        return result.project_id
 
     async def create(
         self,
@@ -104,9 +125,10 @@ async def create(
             for st in payload
         ]
 
+        project_id = await self._resolve_project_id(domain, params.group)
         action = CreateTaskTemplateAction(
             domain_name=domain,
-            requesting_group=params.group,
+            requesting_project=project_id,
             requester_uuid=ctx.user_uuid,
             requester_access_key=ctx.access_key,
             requester_role=req.request["user"]["role"],
@@ -220,10 +242,11 @@ async def update(
             for st in payload
         ]
 
+        project_id = await self._resolve_project_id(domain, params.group)
         action = UpdateTaskTemplateAction(
             template_id=template_id,
             domain_name=domain,
-            requesting_group=params.group,
+            requesting_project=project_id,
             requester_uuid=ctx.user_uuid,
             requester_access_key=ctx.access_key,
             requester_role=req.request["user"]["role"],

src/ai/backend/manager/api/rest/tree.py

@@ -261,8 +261,12 @@ def build_api_routes(
     )
 
     # Template sub-registries
-    cluster_template_handler = ClusterTemplateHandler(template=processors.template)
-    session_template_handler = SessionTemplateHandler(template=processors.template)
+    cluster_template_handler = ClusterTemplateHandler(
+        template=processors.template, group=processors.group
+    )
+    session_template_handler = SessionTemplateHandler(
+        template=processors.template, group=processors.group
+    )
     cluster_template_reg = register_cluster_template_routes(cluster_template_handler, route_deps)
     session_template_reg = register_session_template_routes(session_template_handler, route_deps)
 

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

@@ -1,27 +1,23 @@
 from __future__ import annotations
 
 import enum
-import uuid
-from collections.abc import Iterable, Mapping, Sequence
+from collections.abc import Mapping, Sequence
 from typing import Any, cast
 
 import sqlalchemy as sa
 import trafaret as t
 from sqlalchemy.dialects import postgresql as pgsql
-from sqlalchemy.ext.asyncio import AsyncConnection as SAConnection
 
 from ai.backend.common import validators as tx
 from ai.backend.common.types import SessionTypes
 from ai.backend.manager.defs import DEFAULT_ROLE
 from ai.backend.manager.exceptions import InvalidArgument
 
 from .base import GUID, EnumType, IDColumn, metadata
-from .user import UserRole
 from .vfolder import verify_vfolder_name
 
 __all__: Sequence[str] = (
     "TemplateType",
-    "query_accessible_session_templates",
     "session_templates",
 )
 
@@ -139,112 +135,3 @@ def check_cluster_template(raw_data: Mapping[str, Any]) -> Mapping[str, Any]:
             f"One and only one {DEFAULT_ROLE} node must be created per cluster",
         )
     return cast(Mapping[str, Any], data)
-
-
-async def query_accessible_session_templates(
-    conn: SAConnection,
-    user_uuid: uuid.UUID,
-    template_type: TemplateType,
-    *,
-    user_role: UserRole | None = None,
-    domain_name: str | None = None,
-    allowed_types: Iterable[str] = ["user"],
-    extra_conds: Any = None,
-) -> list[Mapping[str, Any]]:
-    from .group import association_groups_users as agus
-    from .group import groups
-    from .user import users
-
-    entries: list[Mapping[str, Any]] = []
-    if "user" in allowed_types:
-        # Query user templates
-        j = session_templates.join(users, session_templates.c.user_uuid == users.c.uuid)
-        query = (
-            sa.select(
-                session_templates.c.name,
-                session_templates.c.id,
-                session_templates.c.created_at,
-                session_templates.c.user_uuid,
-                session_templates.c.group_id,
-                users.c.email,
-            )
-            .select_from(j)
-            .where(
-                (session_templates.c.user_uuid == user_uuid)
-                & session_templates.c.is_active
-                & (session_templates.c.type == template_type),
-            )
-        )
-        if extra_conds is not None:
-            query = query.where(extra_conds)
-        result = await conn.execute(query)
-        for row in result:
-            entries.append({
-                "name": row.name,
-                "id": row.id,
-                "created_at": row.created_at,
-                "is_owner": True,
-                "user": str(row.user_uuid) if row.user_uuid else None,
-                "group": str(row.group_id) if row.group_id else None,
-                "user_email": row.email,
-                "group_name": None,
-            })
-    if "group" in allowed_types:
-        # Query group session_templates
-        if user_role == UserRole.ADMIN:
-            query = (
-                sa.select(groups.c.id)
-                .select_from(groups)
-                .where(groups.c.domain_name == domain_name)
-            )
-            result = await conn.execute(query)
-            grps = result.fetchall()
-            group_ids = [g.id for g in grps]
-        else:
-            j = sa.join(agus, users, agus.c.user_id == users.c.uuid)
-            query = sa.select(agus.c.group_id).select_from(j).where(agus.c.user_id == user_uuid)
-            result = await conn.execute(query)
-            grps = result.fetchall()
-            group_ids = [g.group_id for g in grps]
-        j = session_templates.join(groups, session_templates.c.group_id == groups.c.id)
-        query = (
-            sa.select(
-                session_templates.c.name,
-                session_templates.c.id,
-                session_templates.c.created_at,
-                session_templates.c.user_uuid,
-                session_templates.c.group_id,
-                groups.c.name,
-            )
-            .set_label_style(sa.LABEL_STYLE_TABLENAME_PLUS_COL)
-            .select_from(j)
-            .where(
-                session_templates.c.group_id.in_(group_ids)
-                & session_templates.c.is_active
-                & (session_templates.c.type == template_type),
-            )
-        )
-        if extra_conds is not None:
-            query = query.where(extra_conds)
-        if "user" in allowed_types:
-            query = query.where(session_templates.c.user_uuid != user_uuid)
-        result = await conn.execute(query)
-        is_owner = user_role == UserRole.ADMIN
-        for row in result:
-            entries.append({
-                "name": row.session_templates_name,
-                "id": row.session_templates_id,
-                "created_at": row.session_templates_created_at,
-                "is_owner": is_owner,
-                "user": (
-                    str(row.session_templates_user_uuid)
-                    if row.session_templates_user_uuid
-                    else None
-                ),
-                "group": (
-                    str(row.session_templates_group_id) if row.session_templates_group_id else None
-                ),
-                "user_email": None,
-                "group_name": row.groups_name,
-            })
-    return entries

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

@@ -18,6 +18,7 @@
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.data.permission.types import RBACElementType
 from ai.backend.common.exception import InvalidAPIParameters
+from ai.backend.common.identifier.project import ProjectID
 from ai.backend.common.types import SlotName, VFolderID
 from ai.backend.common.utils import nmget
 from ai.backend.logging.utils import BraceStyleAdapter
@@ -921,6 +922,32 @@ async def get_project(self, project_id: UUID) -> GroupData:
                 raise ProjectNotFound(f"Project {project_id} not found")
             return row.to_data()
 
+    async def project_id_by_name_in_domain(
+        self, domain_name: str, project_name: str
+    ) -> ProjectID | None:
+        """Resolve an active project's UUID by its domain-scoped name.
+
+        LEGACY: Exists solely to support existing API handlers that only accept a
+        group name as input (e.g. the REST v1 session/cluster template endpoints).
+        New API handlers and any other new code MUST NOT use this — they should
+        accept a project UUID directly.
+
+        Returns:
+            The project UUID if found, or ``None`` if no matching active project exists.
+        """
+        async with self._db.begin_readonly_session_read_committed() as db_sess:
+            result = await db_sess.execute(
+                sa.select(GroupRow.id).where(
+                    GroupRow.domain_name == domain_name,
+                    GroupRow.name == project_name,
+                    GroupRow.is_active.is_(True),
+                )
+            )
+            project_id = result.scalar_one_or_none()
+            if project_id is None:
+                return None
+            return ProjectID(project_id)
+
     async def search_projects(
         self,
         querier: BatchQuerier,

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

@@ -9,6 +9,7 @@
 
 from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient
 from ai.backend.common.exception import BackendAIError
+from ai.backend.common.identifier.project import ProjectID
 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
@@ -174,6 +175,22 @@ async def get_project(self, project_id: UUID) -> GroupData:
         """
         return await self._db_source.get_project(project_id)
 
+    @group_repository_resilience.apply()
+    async def project_id_by_name_in_domain(
+        self, domain_name: str, project_name: str
+    ) -> ProjectID | None:
+        """Resolve an active project's UUID by its domain-scoped name.
+
+        LEGACY: Exists solely to support existing API handlers that only accept a
+        group name as input (e.g. the REST v1 session/cluster template endpoints).
+        New API handlers and any other new code MUST NOT use this — they should
+        accept a project UUID directly.
+
+        Returns:
+            The project UUID if found, or ``None`` if no matching active project exists.
+        """
+        return await self._db_source.project_id_by_name_in_domain(domain_name, project_name)
+
     @group_repository_resilience.apply()
     async def search_projects(
         self,