refactor: migrate to entity slugs (#384)

leafty · web-flow · commit 9a7f277b5535 · 2024-09-13T13:30:00.000+02:00
Moves the `project_slugs` table to `entity_slugs` and `project_slugs_old` to `entity_slugs_old`. This is done to prepare for more entity types to be namespaced.

/deploy renku=release-0.57.0

Note: the sequence migration has been tested:
1. Force migration revision to be `9058bf0a1a12` (before this).
2. Deploy with image tag `renku/renku-data-service:sha-345a1bb` (current `main`) and run `alembic check` and `alembic heads` to verify the DB is before this change.
3. Create some projects.
4. Re-deploy this PR and check that migration revision is `a11752a5afba` (after this).
5. Create more projects.
diff --git a/components/renku_data_services/migrations/versions/a11752a5afba_migrate_to_entity_slugs.py b/components/renku_data_services/migrations/versions/a11752a5afba_migrate_to_entity_slugs.py
@@ -0,0 +1,118 @@
+"""migrate to entity slugs
+
+Revision ID: a11752a5afba
+Revises: 9058bf0a1a12
+Create Date: 2024-09-03 11:18:46.025525
+
+"""
+
+import sqlalchemy as sa
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision = "a11752a5afba"
+down_revision = "9058bf0a1a12"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    connection = op.get_bind()
+
+    op.execute("ALTER TABLE projects.project_slugs SET SCHEMA common")
+    op.rename_table("project_slugs", "entity_slugs", schema="common")
+    op.execute("ALTER INDEX common.project_slugs_unique_slugs RENAME TO entity_slugs_unique_slugs")
+    op.execute(
+        "ALTER INDEX common.ix_projects_project_slugs_namespace_id RENAME TO ix_common_entity_slugs_namespace_id"
+    )
+    op.execute("ALTER INDEX common.ix_projects_project_slugs_project_id RENAME TO ix_common_entity_slugs_project_id")
+    op.execute("ALTER INDEX common.ix_projects_project_slugs_slug RENAME TO ix_common_entity_slugs_slug")
+    op.execute("ALTER SEQUENCE common.project_slugs_id_seq RENAME TO entity_slugs_id_seq")
+    op.drop_constraint("project_slugs_project_id_fk", "entity_slugs", schema="common", type_="foreignkey")
+    op.create_foreign_key(
+        "entity_slugs_project_id_fk",
+        "entity_slugs",
+        "projects",
+        ["project_id"],
+        ["id"],
+        source_schema="common",
+        referent_schema="projects",
+        ondelete="CASCADE",
+    )
+
+    op.execute("ALTER TABLE projects.project_slugs_old SET SCHEMA common")
+    op.rename_table("project_slugs_old", "entity_slugs_old", schema="common")
+    op.execute(
+        "ALTER INDEX common.ix_projects_project_slugs_old_created_at RENAME TO ix_common_entity_slugs_old_created_at"
+    )
+    op.execute(
+        "ALTER INDEX common.ix_projects_project_slugs_old_latest_slug_id RENAME TO ix_common_entity_slugs_old_latest_slug_id"
+    )
+    op.execute("ALTER INDEX common.ix_projects_project_slugs_old_slug RENAME TO ix_common_entity_slugs_old_slug")
+    op.execute("ALTER SEQUENCE common.project_slugs_old_id_seq RENAME TO entity_slugs_old_id_seq")
+
+    tables = ["entity_slugs", "entity_slugs_old"]
+    inspector = sa.inspect(op.get_bind())
+    found_sequences = inspector.get_sequence_names("common")
+    for table in tables:
+        seq = f"{table}_id_seq"
+        if seq not in found_sequences:
+            continue
+        last_id_stmt = sa.select(sa.func.max(sa.column("id", type_=sa.INT))).select_from(
+            sa.table(table, schema="common")
+        )
+        last_id = connection.scalars(last_id_stmt).one_or_none()
+        if last_id is None or last_id <= 0:
+            continue
+        op.execute(sa.text(f"ALTER SEQUENCE common.{seq} RESTART WITH {last_id + 1}"))
+
+
+def downgrade() -> None:
+    connection = op.get_bind()
+
+    op.drop_constraint("entity_slugs_project_id_fk", "entity_slugs", schema="common", type_="foreignkey")
+    op.create_foreign_key(
+        "project_slugs_project_id_fk",
+        "entity_slugs",
+        "projects",
+        ["project_id"],
+        ["id"],
+        source_schema="common",
+        referent_schema="projects",
+        ondelete="CASCADE",
+    )
+    op.execute("ALTER SEQUENCE common.entity_slugs_id_seq RENAME TO project_slugs_id_seq")
+    op.execute("ALTER INDEX common.ix_common_entity_slugs_slug RENAME TO ix_projects_project_slugs_slug")
+    op.execute("ALTER INDEX common.ix_common_entity_slugs_project_id RENAME TO ix_projects_project_slugs_project_id")
+    op.execute(
+        "ALTER INDEX common.ix_common_entity_slugs_namespace_id RENAME TO ix_projects_project_slugs_namespace_id"
+    )
+    op.execute("ALTER INDEX common.entity_slugs_unique_slugs RENAME TO project_slugs_unique_slugs")
+    op.rename_table("entity_slugs", "project_slugs", schema="common")
+    op.execute("ALTER TABLE common.project_slugs SET SCHEMA projects")
+
+    op.execute("ALTER SEQUENCE common.entity_slugs_old_id_seq RENAME TO project_slugs_old_id_seq")
+    op.execute("ALTER INDEX common.ix_common_entity_slugs_old_slug RENAME TO ix_projects_project_slugs_old_slug")
+    op.execute(
+        "ALTER INDEX common.ix_common_entity_slugs_old_latest_slug_id RENAME TO ix_projects_project_slugs_old_latest_slug_id"
+    )
+    op.execute(
+        "ALTER INDEX common.ix_common_entity_slugs_old_created_at RENAME TO ix_projects_project_slugs_old_created_at"
+    )
+    op.rename_table("entity_slugs_old", "project_slugs_old", schema="common")
+    op.execute("ALTER TABLE common.project_slugs_old SET SCHEMA projects")
+
+    tables = ["project_slugs", "project_slugs_old"]
+    inspector = sa.inspect(op.get_bind())
+    found_sequences = inspector.get_sequence_names("projects")
+    for table in tables:
+        seq = f"{table}_id_seq"
+        if seq not in found_sequences:
+            continue
+        last_id_stmt = sa.select(sa.func.max(sa.column("id", type_=sa.INT))).select_from(
+            sa.table(table, schema="projects")
+        )
+        last_id = connection.scalars(last_id_stmt).one_or_none()
+        if last_id is None or last_id <= 0:
+            continue
+        op.execute(sa.text(f"ALTER SEQUENCE projects.{seq} RESTART WITH {last_id + 1}"))
diff --git a/components/renku_data_services/namespace/orm.py b/components/renku_data_services/namespace/orm.py
@@ -3,14 +3,15 @@
 from datetime import datetime
 from typing import Optional
 
-from sqlalchemy import CheckConstraint, DateTime, MetaData, String, func
+from sqlalchemy import CheckConstraint, DateTime, Index, MetaData, String, func
 from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column, relationship
 from sqlalchemy.schema import ForeignKey
 from ulid import ULID
 
 from renku_data_services.base_orm.registry import COMMON_ORM_REGISTRY
 from renku_data_services.errors import errors
 from renku_data_services.namespace import models
+from renku_data_services.project.orm import ProjectORM
 from renku_data_services.users.models import UserInfo, UserWithNamespace
 from renku_data_services.users.orm import UserORM
 from renku_data_services.utils.sqlalchemy import ULIDType
@@ -160,3 +161,40 @@ def dump(self) -> models.Namespace:
             underlying_resource_id=self.latest_slug.user_id,
             name=name,
         )
+
+
+class EntitySlugORM(BaseORM):
+    """Project and namespace slugs."""
+
+    __tablename__ = "entity_slugs"
+    __table_args__ = (Index("entity_slugs_unique_slugs", "namespace_id", "slug", unique=True),)
+
+    id: Mapped[int] = mapped_column(primary_key=True, init=False)
+    slug: Mapped[str] = mapped_column(String(99), index=True, nullable=False)
+    project_id: Mapped[ULID] = mapped_column(
+        ForeignKey(ProjectORM.id, ondelete="CASCADE", name="entity_slugs_project_id_fk"), index=True
+    )
+    project: Mapped[ProjectORM] = relationship(lazy="joined", init=False, repr=False, viewonly=True)
+    namespace_id: Mapped[ULID] = mapped_column(
+        ForeignKey(NamespaceORM.id, ondelete="CASCADE", name="entity_slugs_namespace_id_fk"), index=True
+    )
+    namespace: Mapped[NamespaceORM] = relationship(lazy="joined", init=False, repr=False, viewonly=True)
+
+
+class EntitySlugOldORM(BaseORM):
+    """Project slugs history."""
+
+    __tablename__ = "entity_slugs_old"
+
+    id: Mapped[int] = mapped_column(primary_key=True, init=False)
+    slug: Mapped[str] = mapped_column(String(99), index=True, nullable=False)
+    created_at: Mapped[datetime] = mapped_column(
+        DateTime(timezone=True), nullable=False, index=True, init=False, server_default=func.now()
+    )
+    latest_slug_id: Mapped[int] = mapped_column(
+        ForeignKey(EntitySlugORM.id, ondelete="CASCADE"),
+        nullable=False,
+        init=False,
+        index=True,
+    )
+    latest_slug: Mapped[EntitySlugORM] = relationship(lazy="joined", repr=False, viewonly=True)
diff --git a/components/renku_data_services/project/db.py b/components/renku_data_services/project/db.py
@@ -22,6 +22,7 @@
 from renku_data_services.message_queue.db import EventRepository
 from renku_data_services.message_queue.interface import IMessageQueue
 from renku_data_services.message_queue.redis_queue import dispatch_message
+from renku_data_services.namespace import orm as ns_schemas
 from renku_data_services.namespace.db import GroupRepository
 from renku_data_services.project import apispec as project_apispec
 from renku_data_services.project import models
@@ -99,7 +100,7 @@ async def get_project_by_namespace_slug(
         async with self.session_maker() as session:
             stmt = select(schemas.ProjectORM)
             stmt = _filter_by_namespace_slug(stmt, namespace)
-            stmt = stmt.where(schemas.ProjectSlug.slug == slug.lower())
+            stmt = stmt.where(ns_schemas.EntitySlugORM.slug == slug.lower())
             result = await session.execute(stmt)
             project_orm = result.scalars().first()
 
@@ -135,7 +136,7 @@ async def insert_project(
         if not session:
             raise errors.ProgrammingError(message="A database session is required")
         ns = await session.scalar(
-            select(schemas.NamespaceORM).where(schemas.NamespaceORM.slug == project.namespace.lower())
+            select(ns_schemas.NamespaceORM).where(ns_schemas.NamespaceORM.slug == project.namespace.lower())
         )
         if not ns:
             raise errors.MissingResourceError(
@@ -171,11 +172,12 @@ async def insert_project(
             creation_date=datetime.now(UTC).replace(microsecond=0),
             keywords=project.keywords,
         )
-        project_slug = schemas.ProjectSlug(slug, project_id=project_orm.id, namespace_id=ns.id)
+        project_slug = ns_schemas.EntitySlugORM(slug, project_id=project_orm.id, namespace_id=ns.id)
 
-        session.add(project_slug)
         session.add(project_orm)
         await session.flush()
+        session.add(project_slug)
+        await session.flush()
         await session.refresh(project_orm)
 
         return project_orm.dump()
@@ -241,7 +243,9 @@ async def update_project(
 
         if "namespace" in payload:
             ns_slug = payload["namespace"]
-            ns = await session.scalar(select(schemas.NamespaceORM).where(schemas.NamespaceORM.slug == ns_slug.lower()))
+            ns = await session.scalar(
+                select(ns_schemas.NamespaceORM).where(ns_schemas.NamespaceORM.slug == ns_slug.lower())
+            )
             if not ns:
                 raise errors.MissingResourceError(message=f"The namespace with slug {ns_slug} does not exist")
             if not ns.group_id and not ns.user_id:
@@ -301,9 +305,9 @@ async def delete_project(
 def _filter_by_namespace_slug(statement: Select[tuple[_T]], namespace: str) -> Select[tuple[_T]]:
     """Filters a select query on projects to a given namespace."""
     return (
-        statement.where(schemas.NamespaceORM.slug == namespace.lower())
-        .where(schemas.ProjectSlug.namespace_id == schemas.NamespaceORM.id)
-        .where(schemas.ProjectORM.id == schemas.ProjectSlug.project_id)
+        statement.where(ns_schemas.NamespaceORM.slug == namespace.lower())
+        .where(ns_schemas.EntitySlugORM.namespace_id == ns_schemas.NamespaceORM.id)
+        .where(schemas.ProjectORM.id == ns_schemas.EntitySlugORM.project_id)
     )
 
 
diff --git a/components/renku_data_services/project/orm.py b/components/renku_data_services/project/orm.py
@@ -1,27 +1,31 @@
 """SQLAlchemy's schemas for the projects database."""
 
 from datetime import datetime
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 
-from sqlalchemy import DateTime, Index, Integer, MetaData, String, func
+from sqlalchemy import DateTime, Integer, MetaData, String, func
 from sqlalchemy.dialects.postgresql import ARRAY
 from sqlalchemy.orm import DeclarativeBase, Mapped, MappedAsDataclass, mapped_column, relationship
 from sqlalchemy.schema import ForeignKey
 from ulid import ULID
 
 from renku_data_services.authz import models as authz_models
-from renku_data_services.namespace.orm import NamespaceORM
+from renku_data_services.base_orm.registry import COMMON_ORM_REGISTRY
 from renku_data_services.project import models
 from renku_data_services.project.apispec import Visibility
 from renku_data_services.utils.sqlalchemy import ULIDType
 
+if TYPE_CHECKING:
+    from renku_data_services.namespace.orm import EntitySlugORM
+
 metadata_obj = MetaData(schema="projects")  # Has to match alembic ini section name
 
 
 class BaseORM(MappedAsDataclass, DeclarativeBase):
     """Base class for all ORM classes."""
 
     metadata = metadata_obj
+    registry = COMMON_ORM_REGISTRY
 
 
 class ProjectORM(BaseORM):
@@ -36,7 +40,9 @@ class ProjectORM(BaseORM):
     keywords: Mapped[Optional[list[str]]] = mapped_column("keywords", ARRAY(String(99)), nullable=True)
     # NOTE: The project slugs table has a foreign key from the projects table, but there is a stored procedure
     # triggered by the deletion of slugs to remove the project used by the slug. See migration 89aa4573cfa9.
-    slug: Mapped["ProjectSlug"] = relationship(lazy="joined", init=False, repr=False, viewonly=True)
+    slug: Mapped["EntitySlugORM"] = relationship(
+        lazy="joined", init=False, repr=False, viewonly=True, back_populates="project"
+    )
     repositories: Mapped[list["ProjectRepositoryORM"]] = relationship(
         back_populates="project",
         default_factory=list,
@@ -81,39 +87,3 @@ class ProjectRepositoryORM(BaseORM):
         ForeignKey("projects.id", ondelete="CASCADE"), default=None, index=True
     )
     project: Mapped[Optional[ProjectORM]] = relationship(back_populates="repositories", default=None, repr=False)
-
-
-class ProjectSlug(BaseORM):
-    """Project and namespace slugs."""
-
-    __tablename__ = "project_slugs"
-    __table_args__ = (Index("project_slugs_unique_slugs", "namespace_id", "slug", unique=True),)
-
-    id: Mapped[int] = mapped_column(primary_key=True, init=False)
-    slug: Mapped[str] = mapped_column(String(99), index=True, nullable=False)
-    project_id: Mapped[ULID] = mapped_column(
-        ForeignKey(ProjectORM.id, ondelete="CASCADE", name="project_slugs_project_id_fk"), index=True
-    )
-    namespace_id: Mapped[ULID] = mapped_column(
-        ForeignKey(NamespaceORM.id, ondelete="CASCADE", name="project_slugs_namespace_id_fk"), index=True
-    )
-    namespace: Mapped[NamespaceORM] = relationship(lazy="joined", init=False, repr=False, viewonly=True)
-
-
-class ProjectSlugOld(BaseORM):
-    """Project slugs history."""
-
-    __tablename__ = "project_slugs_old"
-
-    id: Mapped[int] = mapped_column(primary_key=True, init=False)
-    slug: Mapped[str] = mapped_column(String(99), index=True, nullable=False)
-    created_at: Mapped[datetime] = mapped_column(
-        DateTime(timezone=True), nullable=False, index=True, init=False, server_default=func.now()
-    )
-    latest_slug_id: Mapped[int] = mapped_column(
-        ForeignKey(ProjectSlug.id, ondelete="CASCADE"),
-        nullable=False,
-        init=False,
-        index=True,
-    )
-    latest_slug: Mapped[ProjectSlug] = relationship(lazy="joined", repr=False, viewonly=True)
