fix(variant-db): correct foreign key constraints on variant study table

alembic/versions/c0c4aaf84861_add_delete_cascade_variant_studies.py

@@ -0,0 +1,51 @@
+"""
+Add delete cascade constraint to variant study foreign keys
+
+Revision ID: c0c4aaf84861
+Revises: fd73601a9075
+Create Date: 2024-02-21 17:29:48.736664
+"""
+from alembic import op  # type: ignore
+
+# revision identifiers, used by Alembic.
+revision = "c0c4aaf84861"
+down_revision = "fd73601a9075"
+branch_labels = None
+depends_on = None
+
+COMMAND_BLOCK_FK = "commandblock_study_id_fkey"
+SNAPSHOT_FK = "variant_study_snapshot_id_fkey"
+
+
+def upgrade() -> None:
+    dialect_name: str = op.get_context().dialect.name
+
+    # SQLite doesn't support dropping foreign keys, so we need to ignore it here
+    if dialect_name == "postgresql":
+        # ### commands auto generated by Alembic - please adjust! ###
+        with op.batch_alter_table("commandblock", schema=None) as batch_op:
+            batch_op.drop_constraint(COMMAND_BLOCK_FK, type_="foreignkey")
+            batch_op.create_foreign_key(COMMAND_BLOCK_FK, "variantstudy", ["study_id"], ["id"], ondelete="CASCADE")
+
+        with op.batch_alter_table("variant_study_snapshot", schema=None) as batch_op:
+            batch_op.drop_constraint(SNAPSHOT_FK, type_="foreignkey")
+            batch_op.create_foreign_key(SNAPSHOT_FK, "variantstudy", ["id"], ["id"], ondelete="CASCADE")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    dialect_name: str = op.get_context().dialect.name
+
+    # SQLite doesn't support dropping foreign keys, so we need to ignore it here
+    if dialect_name == "postgresql":
+        # ### commands auto generated by Alembic - please adjust! ###
+        with op.batch_alter_table("variant_study_snapshot", schema=None) as batch_op:
+            batch_op.drop_constraint(SNAPSHOT_FK, type_="foreignkey")
+            batch_op.create_foreign_key(SNAPSHOT_FK, "variantstudy", ["id"], ["id"])
+
+        with op.batch_alter_table("commandblock", schema=None) as batch_op:
+            batch_op.drop_constraint(COMMAND_BLOCK_FK, type_="foreignkey")
+            batch_op.create_foreign_key(COMMAND_BLOCK_FK, "variantstudy", ["study_id"], ["id"])
+
+    # ### end Alembic commands ###

antarest/study/storage/variantstudy/model/dbmodel.py

@@ -21,7 +21,7 @@ class VariantStudySnapshot(Base):  # type: ignore
 
     id: str = Column(
         String(36),
-        ForeignKey("variantstudy.id"),
+        ForeignKey("variantstudy.id", ondelete="CASCADE"),
         primary_key=True,
     )
     created_at: datetime.date = Column(DateTime)
@@ -48,7 +48,7 @@ class CommandBlock(Base):  # type: ignore
         default=lambda: str(uuid.uuid4()),
         unique=True,
     )
-    study_id: str = Column(String(36), ForeignKey("variantstudy.id"))
+    study_id: str = Column(String(36), ForeignKey("variantstudy.id", ondelete="CASCADE"))
     index: int = Column(Integer)
     command: str = Column(String(255))
     version: int = Column(Integer)

scripts/rollback.sh

@@ -12,5 +12,5 @@ CUR_DIR=$(cd "$(dirname "$0")" && pwd)
 BASE_DIR=$(dirname "$CUR_DIR")
 
 cd "$BASE_DIR"
-alembic downgrade 3c70366b10ea
+alembic downgrade fd73601a9075
 cd -

tests/storage/repository/test_study.py

@@ -1,11 +1,12 @@
+import json
 from datetime import datetime
 
 from sqlalchemy.orm import Session  # type: ignore
 
 from antarest.core.cache.business.local_chache import LocalCache
 from antarest.core.model import PublicMode
 from antarest.login.model import Group, User
-from antarest.study.model import DEFAULT_WORKSPACE_NAME, RawStudy, Study, StudyContentStatus
+from antarest.study.model import DEFAULT_WORKSPACE_NAME, RawStudy, Study, StudyAdditionalData, StudyContentStatus
 from antarest.study.repository import StudyMetadataRepository
 from antarest.study.storage.variantstudy.model.dbmodel import VariantStudy
 
@@ -77,6 +78,49 @@ def test_lifecycle(db_session: Session) -> None:
     assert repo.get(a_id) is None
 
 
+def test_study__additional_data(db_session: Session) -> None:
+    repo = StudyMetadataRepository(LocalCache(), session=db_session)
+
+    user = User(id=0, name="admin")
+    group = Group(id="my-group", name="group")
+
+    patch = {"foo": "bar"}
+    a = RawStudy(
+        name="a",
+        version="820",
+        author="John Smith",
+        created_at=datetime.utcnow(),
+        updated_at=datetime.utcnow(),
+        public_mode=PublicMode.FULL,
+        owner=user,
+        groups=[group],
+        workspace=DEFAULT_WORKSPACE_NAME,
+        path="study",
+        content_status=StudyContentStatus.WARNING,
+        additional_data=StudyAdditionalData(author="John Smith", horizon="2024-2050", patch=json.dumps(patch)),
+    )
+
+    repo.save(a)
+    a_id = a.id
+
+    # Check that the additional data is correctly saved
+    additional_data = repo.get_additional_data(a_id)
+    assert additional_data.author == "John Smith"
+    assert additional_data.horizon == "2024-2050"
+    assert json.loads(additional_data.patch) == patch
+
+    # Check that the additional data is correctly updated
+    new_patch = {"foo": "baz"}
+    a.additional_data.patch = json.dumps(new_patch)
+    repo.save(a)
+    additional_data = repo.get_additional_data(a_id)
+    assert json.loads(additional_data.patch) == new_patch
+
+    # Check that the additional data is correctly deleted when the study is deleted
+    repo.delete(a_id)
+    assert repo.get_additional_data(a_id) is None
+
+
 def test_study_inheritance(db_session: Session) -> None:
     repo = StudyMetadataRepository(LocalCache(), session=db_session)