Do not load all projects attrs for projects filtering

Author: r4victor

src/dstack/_internal/server/services/fleets.py

@@ -67,7 +67,6 @@
 from dstack._internal.server.services.projects import (
     get_member,
     get_member_permissions,
-    list_project_models,
     list_user_project_models,
 )
 from dstack._internal.server.services.resources import set_resources_defaults
@@ -88,10 +87,11 @@ async def list_fleets(
     limit: int,
     ascending: bool,
 ) -> List[Fleet]:
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+        only_names=True,
+    )
     if project_name is not None:
         projects = [p for p in projects if p.name == project_name]
     fleet_models = await list_projects_fleet_models(

src/dstack/_internal/server/services/instances.py

@@ -34,7 +34,6 @@
     TerminationPolicy,
 )
 from dstack._internal.core.models.runs import JobProvisioningData, Requirements
-from dstack._internal.core.models.users import GlobalRole
 from dstack._internal.core.models.volumes import Volume
 from dstack._internal.core.services.profiles import get_termination
 from dstack._internal.server.models import (
@@ -44,7 +43,7 @@
     UserModel,
 )
 from dstack._internal.server.services.offers import generate_shared_offer
-from dstack._internal.server.services.projects import list_project_models, list_user_project_models
+from dstack._internal.server.services.projects import list_user_project_models
 from dstack._internal.utils import common as common_utils
 from dstack._internal.utils.logging import get_logger
 
@@ -372,18 +371,15 @@ async def list_user_instances(
     limit: int,
     ascending: bool,
 ) -> List[Instance]:
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
-    if not projects:
-        return []
-
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+        only_names=True,
+    )
     if project_names is not None:
-        projects = [proj for proj in projects if proj.name in project_names]
+        projects = [p for p in projects if p.name in project_names]
         if len(projects) == 0:
             return []
-
     instance_models = await list_projects_instance_models(
         session=session,
         projects=projects,

src/dstack/_internal/server/services/projects.py

@@ -4,7 +4,7 @@
 from sqlalchemy import delete, select, update
 from sqlalchemy import func as safunc
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload
+from sqlalchemy.orm import QueryableAttribute, joinedload, load_only
 
 from dstack._internal.core.backends.configurators import get_configurator
 from dstack._internal.core.backends.dstack.models import (
@@ -53,13 +53,12 @@ async def list_user_projects(
     user: UserModel,
 ) -> List[Project]:
     """
-    Returns projects where the user is a member.
+    Returns projects where the user is a member or all projects for global admins.
     """
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
-
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+    )
     projects = sorted(projects, key=lambda p: p.created_at)
     return [
         project_model_to_project(p, include_backends=False, include_members=False)
@@ -79,7 +78,7 @@ async def list_user_accessible_projects(
     if user.global_role == GlobalRole.ADMIN:
         projects = await list_project_models(session=session)
     else:
-        member_projects = await list_user_project_models(session=session, user=user)
+        member_projects = await list_member_project_models(session=session, user=user)
         public_projects = await list_public_non_member_project_models(session=session, user=user)
         projects = member_projects + public_projects
 
@@ -166,7 +165,7 @@ async def delete_projects(
     projects_names: List[str],
 ):
     if user.global_role != GlobalRole.ADMIN:
-        user_projects = await list_user_project_models(
+        user_projects = await list_member_project_models(
             session=session, user=user, include_members=True
         )
         user_project_names = [p.name for p in user_projects]
@@ -338,16 +337,34 @@ async def clear_project_members(
 
 
 async def list_user_project_models(
+    session: AsyncSession,
+    user: UserModel,
+    only_names: bool = False,
+) -> List[ProjectModel]:
+    load_only_attrs = []
+    if only_names:
+        load_only_attrs += [ProjectModel.id, ProjectModel.name]
+    if user.global_role == GlobalRole.ADMIN:
+        return await list_project_models(session=session, load_only_attrs=load_only_attrs)
+    return await list_member_project_models(
+        session=session, user=user, load_only_attrs=load_only_attrs
+    )
+
+
+async def list_member_project_models(
     session: AsyncSession,
     user: UserModel,
     include_members: bool = False,
+    load_only_attrs: Optional[List[QueryableAttribute]] = None,
 ) -> List[ProjectModel]:
     """
     List project models for a user where they are a member.
     """
     options = []
     if include_members:
         options.append(joinedload(ProjectModel.members))
+    if load_only_attrs:
+        options.append(load_only(*load_only_attrs))
     res = await session.execute(
         select(ProjectModel)
         .where(
@@ -394,9 +411,13 @@ async def list_user_owned_project_models(
 
 async def list_project_models(
     session: AsyncSession,
+    load_only_attrs: Optional[List[QueryableAttribute]] = None,
 ) -> List[ProjectModel]:
+    options = []
+    if load_only_attrs:
+        options.append(load_only(*load_only_attrs))
     res = await session.execute(
-        select(ProjectModel).where(ProjectModel.deleted == False),
+        select(ProjectModel).where(ProjectModel.deleted == False).options(*options)
     )
     return list(res.scalars().all())
 

src/dstack/_internal/server/services/runs.py

@@ -8,7 +8,7 @@
 from apscheduler.triggers.cron import CronTrigger
 from sqlalchemy import and_, func, or_, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload
 
 import dstack._internal.utils.common as common_utils
 from dstack._internal.core.errors import (
@@ -43,7 +43,6 @@
     RunTerminationReason,
     ServiceSpec,
 )
-from dstack._internal.core.models.users import GlobalRole
 from dstack._internal.core.models.volumes import (
     InstanceMountPoint,
     Volume,
@@ -82,7 +81,7 @@
 from dstack._internal.server.services.logging import fmt
 from dstack._internal.server.services.offers import get_offers_by_requirements
 from dstack._internal.server.services.plugins import apply_plugin_policies
-from dstack._internal.server.services.projects import list_project_models, list_user_project_models
+from dstack._internal.server.services.projects import list_user_project_models
 from dstack._internal.server.services.resources import set_resources_defaults
 from dstack._internal.server.services.secrets import get_project_secrets_mapping
 from dstack._internal.server.services.users import get_user_model_by_name
@@ -116,10 +115,11 @@ async def list_user_runs(
 ) -> List[Run]:
     if project_name is None and repo_id is not None:
         return []
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+        only_names=True,
+    )
     runs_user = None
     if username is not None:
         runs_user = await get_user_model_by_name(session=session, username=username)
@@ -218,9 +218,9 @@ async def list_projects_run_models(
     res = await session.execute(
         select(RunModel)
         .where(*filters)
+        .options(joinedload(RunModel.user).load_only(UserModel.name))
         .order_by(*order_by)
         .limit(limit)
-        .options(selectinload(RunModel.user))
     )
     run_models = list(res.scalars().all())
     return run_models

src/dstack/_internal/server/services/volumes.py

@@ -13,7 +13,6 @@
     ResourceExistsError,
     ServerClientError,
 )
-from dstack._internal.core.models.users import GlobalRole
 from dstack._internal.core.models.volumes import (
     Volume,
     VolumeAttachment,
@@ -40,7 +39,7 @@
     string_to_lock_id,
 )
 from dstack._internal.server.services.plugins import apply_plugin_policies
-from dstack._internal.server.services.projects import list_project_models, list_user_project_models
+from dstack._internal.server.services.projects import list_user_project_models
 from dstack._internal.utils import common, random_names
 from dstack._internal.utils.logging import get_logger
 
@@ -57,10 +56,11 @@ async def list_volumes(
     limit: int,
     ascending: bool,
 ) -> List[Volume]:
-    if user.global_role == GlobalRole.ADMIN:
-        projects = await list_project_models(session=session)
-    else:
-        projects = await list_user_project_models(session=session, user=user)
+    projects = await list_user_project_models(
+        session=session,
+        user=user,
+        only_names=True,
+    )
     if project_name is not None:
         projects = [p for p in projects if p.name == project_name]
     volume_models = await list_projects_volume_models(