diff --git a/alembic/versions/c0c4aaf84861_add_delete_cascade_variant_studies.py b/alembic/versions/c0c4aaf84861_add_delete_cascade_variant_studies.py new file mode 100644 index 0000000000..cc1100eeba --- /dev/null +++ b/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 ### diff --git a/antarest/study/storage/variantstudy/model/dbmodel.py b/antarest/study/storage/variantstudy/model/dbmodel.py index 9272eb797f..f4a99d74e9 100644 --- a/antarest/study/storage/variantstudy/model/dbmodel.py +++ b/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) diff --git a/scripts/rollback.sh b/scripts/rollback.sh index 46d04a7966..974ca59422 100755 --- a/scripts/rollback.sh +++ b/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 - diff --git a/tests/storage/repository/test_study.py b/tests/storage/repository/test_study.py index bb4ea795e3..5535f4d074 100644 --- a/tests/storage/repository/test_study.py +++ b/tests/storage/repository/test_study.py @@ -1,3 +1,4 @@ +import json from datetime import datetime from sqlalchemy.orm import Session # type: ignore @@ -5,7 +6,7 @@ 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)