fix: Wrong container registry migration script (#2949) (#2950)

Co-authored-by: Gyubong Lee <jopemachine@naver.com>

changes/2949.fix.md

@@ -0,0 +1 @@
+Fix wrong container registry migration script.

src/ai/backend/manager/models/alembic/versions/1d42c726d8a3_migrate_container_registries_from_etcd_to_psql.py

@@ -28,7 +28,6 @@
 from ai.backend.common.etcd import AsyncEtcd, ConfigScopes
 from ai.backend.manager.config import load
 from ai.backend.manager.models.base import GUID, Base, IDColumn, StrEnumType, convention
-from ai.backend.manager.models.image import ImageRow
 
 # revision identifiers, used by Alembic.
 revision = "1d42c726d8a3"
@@ -64,6 +63,37 @@ class ContainerRegistryType(enum.StrEnum):
     LOCAL = "local"
 
 
+metadata = sa.MetaData(naming_convention=convention)
+
+
+class ImageType(enum.Enum):
+    COMPUTE = "compute"
+    SYSTEM = "system"
+    SERVICE = "service"
+
+
+images_table = sa.Table(
+    "images",
+    metadata,
+    sa.Column("id", GUID, primary_key=True, nullable=False),
+    sa.Column("name", sa.String, nullable=False, index=True),
+    sa.Column("image", sa.String, nullable=False, index=True),
+    sa.Column("created_at", sa.DateTime(timezone=True), server_default=sa.func.now(), index=True),
+    sa.Column("tag", sa.TEXT),
+    sa.Column("registry", sa.String, nullable=False, index=True),
+    sa.Column("registry_id", GUID, nullable=True, index=True),
+    sa.Column("architecture", sa.String, nullable=False, index=True, server_default="x86_64"),
+    sa.Column("config_digest", sa.CHAR(length=72), nullable=False),
+    sa.Column("size_bytes", sa.BigInteger, nullable=False),
+    sa.Column("is_local", sa.Boolean, nullable=False, server_default=sa.sql.expression.false()),
+    sa.Column("type", sa.Enum(ImageType), nullable=False),
+    sa.Column("accelerators", sa.String),
+    sa.Column("labels", sa.JSON, nullable=False, server_default="{}"),
+    sa.Column("resources", sa.JSON, nullable=False, server_default="{}"),
+    extend_existing=True,
+)
+
+
 def get_container_registry_row_schema():
     class ContainerRegistryRow(Base):
         __tablename__ = "container_registries"
@@ -84,7 +114,6 @@ class ContainerRegistryRow(Base):
         password = sa.Column("password", sa.String, nullable=True)
         ssl_verify = sa.Column("ssl_verify", sa.Boolean, server_default=sa.text("true"), index=True)
         is_global = sa.Column("is_global", sa.Boolean, server_default=sa.text("true"), index=True)
-        extra = sa.Column("extra", sa.JSON, nullable=True, default=None)
 
     return ContainerRegistryRow
 
@@ -216,10 +245,22 @@ def revert_data_psql_to_etcd() -> None:
     ContainerRegistryRow = get_container_registry_row_schema()
 
     db_connection = op.get_bind()
-    rows = db_connection.execute(sa.select(ContainerRegistryRow)).fetchall()
+
+    # Prevent error from presence or absence of the extra column
+    rows = db_connection.execute(
+        sa.select([
+            ContainerRegistryRow.url,
+            ContainerRegistryRow.registry_name,
+            ContainerRegistryRow.type,
+            ContainerRegistryRow.project,
+            ContainerRegistryRow.username,
+            ContainerRegistryRow.password,
+            ContainerRegistryRow.ssl_verify,
+        ])
+    ).fetchall()
     items = []
 
-    for id, url, registry_name, type, project, username, password, ssl_verify, _is_global in rows:
+    for url, registry_name, type, project, username, password, ssl_verify in rows:
         item = {
             "": url,
             "type": str(type),
@@ -287,9 +328,9 @@ def insert_registry_id_to_images() -> None:
 
     if registry_infos:
         insert_registry_id_query = (
-            sa.update(ImageRow)
+            sa.update(images_table)
             .values(registry_id=sa.bindparam("registry_id"))
-            .where(ImageRow.name.startswith(sa.bindparam("registry_name_and_project")))
+            .where(images_table.c.name.startswith(sa.bindparam("registry_name_and_project")))
         )
 
         query_params = []

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

@@ -0,0 +1,34 @@
+"""Add extra column to container_registries
+
+Revision ID: e9e574a6e22d
+Revises: 7c8501cec07b
+Create Date: 2024-10-23 20:56:36.513421
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.engine.reflection import Inspector
+
+# revision identifiers, used by Alembic.
+revision = "e9e574a6e22d"
+down_revision = "7c8501cec07b"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    inspector = Inspector.from_engine(conn)
+    columns = [col["name"] for col in inspector.get_columns("container_registries")]
+
+    # Prevent error in case the extra column was created using a previously modified migration script.
+    if "extra" not in columns:
+        op.add_column(
+            "container_registries",
+            sa.Column("extra", sa.JSON, default=None, nullable=True),
+        )
+
+
+def downgrade() -> None:
+    op.drop_column("container_registries", "extra")