From 1e4edbf8ee1cf8e9bcfb2af9e0f5f549fc5dafbd Mon Sep 17 00:00:00 2001 From: octodog Date: Fri, 7 Feb 2025 17:49:05 +0900 Subject: [PATCH] fix(BA-584): Remove foreign key constraint from `EndpointRow.image` column (#3599) (#3602) Co-authored-by: Gyubong Lee --- changes/3599.fix.md | 1 + ...e4_remove_fk_constraint_from_endpoints_.py | 37 +++++++++++++++++++ src/ai/backend/manager/models/endpoint.py | 25 ++++++++++--- src/ai/backend/manager/models/image.py | 15 +++++++- 4 files changed, 71 insertions(+), 7 deletions(-) create mode 100644 changes/3599.fix.md create mode 100644 src/ai/backend/manager/models/alembic/versions/ecc9f6322be4_remove_fk_constraint_from_endpoints_.py diff --git a/changes/3599.fix.md b/changes/3599.fix.md new file mode 100644 index 00000000000..953b26912dd --- /dev/null +++ b/changes/3599.fix.md @@ -0,0 +1 @@ +Remove foreign key constraint from `EndpointRow.image` column. diff --git a/src/ai/backend/manager/models/alembic/versions/ecc9f6322be4_remove_fk_constraint_from_endpoints_.py b/src/ai/backend/manager/models/alembic/versions/ecc9f6322be4_remove_fk_constraint_from_endpoints_.py new file mode 100644 index 00000000000..5b42bea05c8 --- /dev/null +++ b/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" + ) diff --git a/src/ai/backend/manager/models/endpoint.py b/src/ai/backend/manager/models/endpoint.py index b42fb046826..1111197523c 100644 --- a/src/ai/backend/manager/models/endpoint.py +++ b/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" diff --git a/src/ai/backend/manager/models/image.py b/src/ai/backend/manager/models/image.py index 42ab8e6513a..f231c99093c 100644 --- a/src/ai/backend/manager/models/image.py +++ b/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,