Remove implicit gateway.backend load

Author: r4victor

src/dstack/_internal/server/background/pipeline_tasks/gateways.py

@@ -23,6 +23,7 @@
 )
 from dstack._internal.server.db import get_db, get_session_ctx
 from dstack._internal.server.models import (
+    BackendModel,
     GatewayComputeModel,
     GatewayModel,
     ProjectModel,
@@ -210,6 +211,7 @@ async def _process_submitted_item(item: GatewayPipelineItem):
                 GatewayModel.lock_token == item.lock_token,
             )
             .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends))
+            .options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
         )
         gateway_model = res.unique().scalar_one_or_none()
         if gateway_model is None:
@@ -431,6 +433,7 @@ async def _process_to_be_deleted_item(item: GatewayPipelineItem):
             )
             .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends))
             .options(joinedload(GatewayModel.gateway_compute))
+            .options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
         )
         gateway_model = res.unique().scalar_one_or_none()
         if gateway_model is None:

src/dstack/_internal/server/background/scheduled_tasks/gateways.py

@@ -7,7 +7,12 @@
 from dstack._internal.core.errors import BackendError, BackendNotAvailable, SSHError
 from dstack._internal.core.models.gateways import GatewayStatus
 from dstack._internal.server.db import get_db, get_session_ctx
-from dstack._internal.server.models import GatewayComputeModel, GatewayModel, ProjectModel
+from dstack._internal.server.models import (
+    BackendModel,
+    GatewayComputeModel,
+    GatewayModel,
+    ProjectModel,
+)
 from dstack._internal.server.services import backends as backends_services
 from dstack._internal.server.services import gateways as gateways_services
 from dstack._internal.server.services.gateways import (
@@ -109,6 +114,7 @@ async def _process_submitted_gateway(session: AsyncSession, gateway_model: Gatew
         select(GatewayModel)
         .where(GatewayModel.id == gateway_model.id)
         .options(joinedload(GatewayModel.project).joinedload(ProjectModel.backends))
+        .options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
         .execution_options(populate_existing=True)
     )
     gateway_model = res.unique().scalar_one()

src/dstack/_internal/server/models.py

@@ -513,7 +513,7 @@ class GatewayModel(PipelineModelMixin, BaseModel):
     project_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("projects.id", ondelete="CASCADE"))
     project: Mapped["ProjectModel"] = relationship(foreign_keys=[project_id])
     backend_id: Mapped[uuid.UUID] = mapped_column(ForeignKey("backends.id", ondelete="CASCADE"))
-    backend: Mapped["BackendModel"] = relationship(lazy="selectin")
+    backend: Mapped["BackendModel"] = relationship()
 
     gateway_compute_id: Mapped[Optional[uuid.UUID]] = mapped_column(
         ForeignKey("gateway_computes.id", ondelete="CASCADE")

src/dstack/_internal/server/services/gateways/__init__.py

@@ -10,7 +10,7 @@
 import httpx
 from sqlalchemy import func, select, update
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import joinedload, selectinload
+from sqlalchemy.orm import joinedload
 
 import dstack._internal.utils.random_names as random_names
 from dstack._internal.core.backends.base.compute import (
@@ -131,6 +131,7 @@ async def list_project_gateways(session: AsyncSession, project: ProjectModel) ->
         session=session,
         project=project,
         load_gateway_compute=True,
+        load_backend_type=True,
     )
     return [gateway_model_to_gateway(g) for g in gateways]
 
@@ -143,6 +144,7 @@ async def get_gateway_by_name(
         project=project,
         name=name,
         load_gateway_compute=True,
+        load_backend_type=True,
     )
     if gateway is None:
         return None
@@ -254,6 +256,14 @@ async def create_gateway(
                 session=session, project=project, name=configuration.name, user=user
             )
         pipeline_hinter.hint_fetch(GatewayModel.__name__)
+        gateway = await get_project_gateway_model_by_name(
+            session=session,
+            project=project,
+            name=configuration.name,
+            load_gateway_compute=True,
+            load_backend_type=True,
+        )
+        assert gateway is not None
         return gateway_model_to_gateway(gateway)
 
 
@@ -392,10 +402,11 @@ async def _delete_gateways_sync(
                 GatewayModel.project_id == project.id,
                 GatewayModel.name.in_(gateways_names),
             )
-            .options(selectinload(GatewayModel.gateway_compute))
+            .options(joinedload(GatewayModel.gateway_compute))
+            .options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
             .execution_options(populate_existing=True)
             .order_by(GatewayModel.id)  # take locks in order
-            .with_for_update(key_share=True)
+            .with_for_update(key_share=True, of=GatewayModel)
         )
         gateway_models = res.scalars().all()
         for gateway_model in gateway_models:
@@ -506,10 +517,13 @@ async def list_project_gateway_models(
     session: AsyncSession,
     project: ProjectModel,
     load_gateway_compute: bool = False,
+    load_backend_type: bool = False,
 ) -> Sequence[GatewayModel]:
     stmt = select(GatewayModel).where(GatewayModel.project_id == project.id)
     if load_gateway_compute:
         stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    if load_backend_type:
+        stmt = stmt.options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
     res = await session.execute(stmt)
     return res.scalars().all()
 
@@ -519,13 +533,16 @@ async def get_project_gateway_model_by_name(
     project: ProjectModel,
     name: str,
     load_gateway_compute: bool = False,
+    load_backend_type: bool = False,
 ) -> Optional[GatewayModel]:
     stmt = select(GatewayModel).where(
         GatewayModel.project_id == project.id,
         GatewayModel.name == name,
     )
     if load_gateway_compute:
         stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    if load_backend_type:
+        stmt = stmt.options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
     res = await session.execute(stmt)
     return res.scalar()
 
@@ -558,6 +575,7 @@ async def get_project_gateway_model_by_name_for_update(
                 select(GatewayModel)
                 .where(GatewayModel.id.in_([gateway_id]), *filters)
                 .options(joinedload(GatewayModel.gateway_compute))
+                .options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
                 .with_for_update(key_share=True, of=GatewayModel)
             )
             yield res.scalar_one_or_none()
@@ -567,13 +585,16 @@ async def get_project_default_gateway_model(
     session: AsyncSession,
     project: ProjectModel,
     load_gateway_compute: bool = False,
+    load_backend_type: bool = False,
 ) -> Optional[GatewayModel]:
     stmt = select(GatewayModel).where(
         GatewayModel.id == project.default_gateway_id,
         GatewayModel.to_be_deleted == False,
     )
     if load_gateway_compute:
         stmt = stmt.options(joinedload(GatewayModel.gateway_compute))
+    if load_backend_type:
+        stmt = stmt.options(joinedload(GatewayModel.backend).load_only(BackendModel.type))
     res = await session.execute(stmt)
     return res.scalar_one_or_none()
 

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

@@ -94,6 +94,7 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec:
             project=run_model.project,
             name=run_spec.configuration.gateway,
             load_gateway_compute=True,
+            load_backend_type=True,
         )
         if gateway is None:
             raise ResourceNotExistsError(
@@ -110,6 +111,7 @@ async def register_service(session: AsyncSession, run_model: RunModel, run_spec:
             session=session,
             project=run_model.project,
             load_gateway_compute=True,
+            load_backend_type=True,
         )
         if gateway is None and run_spec.configuration.gateway == True:
             raise ResourceNotExistsError(