fix(BA-584): Remove foreign key constraint from EndpointRow.image c…

…olumn (#3599) (#3602)

Co-authored-by: Gyubong Lee <[email protected]>

changes/3599.fix.md

@@ -0,0 +1 @@
+Remove foreign key constraint from `EndpointRow.image` column.

src/ai/backend/manager/models/alembic/versions/ecc9f6322be4_remove_fk_constraint_from_endpoints_.py

@@ -0,0 +1,37 @@
+"""Remove foreign key constraint from endpoints.image column
+
+Revision ID: ecc9f6322be4
+Revises: ef9a7960d234
+Create Date: 2025-02-07 00:58:05.211395
+
+"""
+
+from alembic import op
+
+from ai.backend.manager.models.base import GUID
+
+# revision identifiers, used by Alembic.
+revision = "ecc9f6322be4"
+down_revision = "ef9a7960d234"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    op.drop_constraint("fk_endpoints_image_images", "endpoints", type_="foreignkey")
+    op.alter_column("endpoints", "image", existing_type=GUID, nullable=True)
+    op.create_check_constraint(
+        constraint_name="ck_image_required_unless_destroyed",
+        table_name="endpoints",
+        condition="lifecycle_stage = 'destroyed' OR image IS NOT NULL",
+    )
+
+
+def downgrade() -> None:
+    op.create_foreign_key(
+        "fk_endpoints_image_images", "endpoints", "images", ["image"], ["id"], ondelete="RESTRICT"
+    )
+    op.alter_column("endpoints", "image", existing_type=GUID, nullable=False)
+    op.drop_constraint(
+        constraint_name="ck_image_required_unless_destroyed", table_name="endpoints", type_="check"
+    )

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

@@ -13,9 +13,10 @@
 import yarl
 from graphene.types.datetime import DateTime as GQLDateTime
 from graphql import Undefined
+from sqlalchemy import CheckConstraint
 from sqlalchemy.dialects import postgresql as pgsql
 from sqlalchemy.ext.asyncio import AsyncConnection, AsyncSession
-from sqlalchemy.orm import relationship, selectinload
+from sqlalchemy.orm import foreign, relationship, selectinload
 from sqlalchemy.orm.exc import NoResultFound
 
 from ai.backend.common.config import model_definition_iv
@@ -102,6 +103,16 @@ class EndpointLifecycle(Enum):
 class EndpointRow(Base):
     __tablename__ = "endpoints"
 
+    __table_args__ = (
+        CheckConstraint(
+            sa.or_(
+                sa.column("lifecycle_stage") == EndpointLifecycle.DESTROYED.value,
+                sa.column("image").isnot(None),
+            ),
+            name="ck_image_required_unless_destroyed",
+        ),
+    )
+
     id = EndpointIDColumn()
     name = sa.Column("name", sa.String(length=512), nullable=False)
     created_user = sa.Column(
@@ -114,9 +125,7 @@ class EndpointRow(Base):
     desired_session_count = sa.Column(
         "desired_session_count", sa.Integer, nullable=False, default=0, server_default="0"
     )
-    image = sa.Column(
-        "image", GUID, sa.ForeignKey("images.id", ondelete="RESTRICT"), nullable=False
-    )
+    image = sa.Column("image", GUID)
     model = sa.Column(
         "model",
         GUID,
@@ -206,7 +215,13 @@ class EndpointRow(Base):
 
     routings = relationship("RoutingRow", back_populates="endpoint_row")
     tokens = relationship("EndpointTokenRow", back_populates="endpoint_row")
-    image_row = relationship("ImageRow", back_populates="endpoints")
+    image_row = relationship(
+        "ImageRow",
+        primaryjoin=lambda: foreign(EndpointRow.image) == ImageRow.id,
+        foreign_keys=[image],
+        back_populates="endpoints",
+    )
+
     model_row = relationship("VFolderRow", back_populates="endpoints")
     created_user_row = relationship(
         "UserRow", back_populates="created_endpoints", foreign_keys="EndpointRow.created_user"

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

@@ -28,7 +28,7 @@
 from redis.asyncio import Redis
 from redis.asyncio.client import Pipeline
 from sqlalchemy.ext.asyncio import AsyncSession
-from sqlalchemy.orm import load_only, relationship, selectinload
+from sqlalchemy.orm import foreign, load_only, relationship, selectinload
 
 from ai.backend.common import redis_helper
 from ai.backend.common.docker import ImageRef
@@ -179,6 +179,13 @@ class ImageType(enum.Enum):
     SERVICE = "service"
 
 
+# Defined for avoiding circular import
+def _get_image_endpoint_join_condition():
+    from ai.backend.manager.models.endpoint import EndpointRow
+
+    return ImageRow.id == foreign(EndpointRow.image)
+
+
 class ImageRow(Base):
     __tablename__ = "images"
     id = IDColumn("id")
@@ -221,7 +228,11 @@ class ImageRow(Base):
     )
     aliases: relationship
     # sessions = relationship("SessionRow", back_populates="image_row")
-    endpoints = relationship("EndpointRow", back_populates="image_row")
+    endpoints = relationship(
+        "EndpointRow",
+        primaryjoin=_get_image_endpoint_join_condition,
+        back_populates="image_row",
+    )
 
     def __init__(
         self,