Allow sequential simulations

Co-authored-by: Zohar Malamant <git@wah.pink>
Co-authored-by: Lars Petter Hauge <lars-petter-hauge@users.noreply.github.com>
Co-authored-by: Muhammad Usama <110888193+UsamaEquinorAFK@users.noreply.github.com>

Author: 3 people

.envrc

@@ -4,4 +4,4 @@ eval "$(devenv direnvrc)"
 
 # The use_devenv function supports passing flags to the devenv command
 # For example: use devenv --impure --option services.postgres.enable:bool true
-use devenv
+use devenv -P pg

.github/workflows/python-pgsql.yaml

@@ -43,9 +43,15 @@ jobs:
         pipx install poetry
         poetry sync -E pg
 
+    - name: Create test database
+      run: psql -h localhost -U postgres -c 'CREATE DATABASE "acidwatch_test";'
+      env:
+        PGPASSWORD: password
+
     - name: Run pytest
       run: |
         poetry run alembic upgrade head
-        poetry run pytest tests
+        poetry run pytest tests --test-alembic
       env:
         ACIDWATCH_DATABASE: postgresql://postgres:password@localhost:5432/postgres
+        ACIDWATCH_TEST_DATABASE: postgresql://postgres:password@localhost:5432/acidwatch_test

backend/alembic/env.py

@@ -62,11 +62,15 @@ def run_migrations_online() -> None:
     and associate a connection with the context.
 
     """
-    connectable = engine_from_config(
-        config.get_section(config.config_ini_section, {}),
-        prefix="sqlalchemy.",
-        poolclass=pool.NullPool,
-    )
+    # Allow for overriding the connection during tests
+    connectable = context.config.attributes.get("connection", None)
+
+    if connectable is None:
+        connectable = engine_from_config(
+            config.get_section(config.config_ini_section, {}),
+            prefix="sqlalchemy.",
+            poolclass=pool.NullPool,
+        )
 
     with connectable.connect() as connection:
         context.configure(connection=connection, target_metadata=target_metadata)

backend/alembic/versions/01aaa143d690_refactor_simulation_to_use_model_inputs_.py

@@ -0,0 +1,125 @@
+"""refactor simulation to use model_inputs table
+
+Revision ID: 01aaa143d690
+Revises: c35588effcc4
+Create Date: 2026-01-21 16:50:46.128272
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = "01aaa143d690"
+down_revision: Union[str, Sequence[str], None] = "c35588effcc4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    """Upgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.create_table(
+        "model_inputs",
+        sa.Column("simulation_id", sa.Uuid(), nullable=False),
+        sa.Column("previous_model_input_id", sa.Uuid(), nullable=True),
+        sa.Column("model_id", sa.String(), nullable=False),
+        sa.Column("parameters", sa.JSON(), nullable=False),
+        sa.Column("id", sa.Uuid(), nullable=False),
+        sa.Column("created_at", sa.DateTime(), nullable=False),
+        sa.Column("updated_at", sa.DateTime(), nullable=False),
+        sa.ForeignKeyConstraint(
+            ["simulation_id"],
+            ["simulations.id"],
+        ),
+        sa.ForeignKeyConstraint(
+            ["previous_model_input_id"],
+            ["model_inputs.id"],
+        ),
+        sa.PrimaryKeyConstraint("id"),
+    )
+    op.add_column("results", sa.Column("model_input_id", sa.Uuid(), nullable=True))
+
+    op.execute("""
+        INSERT INTO model_inputs (id, simulation_id, parameters, created_at, updated_at, model_id)
+        SELECT
+            gen_random_uuid(),
+            simulations.id,
+            simulations.parameters,
+            simulations.created_at,
+            simulations.updated_at,
+            simulations.model_id
+        FROM simulations
+    """)
+
+    op.execute("""
+        UPDATE results
+        SET model_input_id = model_inputs.id
+        FROM model_inputs
+        WHERE results.simulation_id = model_inputs.simulation_id
+    """)
+
+    op.alter_column("results", "model_input_id", nullable=False)
+    op.drop_constraint(
+        op.f("results_simulation_id_fkey"), "results", type_="foreignkey"
+    )
+    op.create_foreign_key(None, "results", "model_inputs", ["model_input_id"], ["id"])
+    op.drop_column("results", "simulation_id")
+    op.drop_column("simulations", "model_id")
+    op.drop_column("simulations", "parameters")
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    """Downgrade schema."""
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "simulations",
+        sa.Column(
+            "parameters",
+            postgresql.JSON(astext_type=sa.Text()),
+            autoincrement=False,
+            nullable=True,
+        ),
+    )
+    op.add_column(
+        "simulations",
+        sa.Column("model_id", sa.VARCHAR(), autoincrement=False, nullable=True),
+    )
+    op.add_column(
+        "results",
+        sa.Column("simulation_id", sa.UUID(), autoincrement=False, nullable=True),
+    )
+
+    op.execute("""
+        UPDATE simulations
+        SET parameters = model_inputs.parameters,
+            model_id = model_inputs.model_id
+        FROM model_inputs
+        WHERE model_inputs.simulation_id = simulations.id
+    """)
+
+    op.execute("""
+        UPDATE results
+        SET simulation_id = model_inputs.simulation_id
+        FROM model_inputs
+        WHERE results.model_input_id = model_inputs.id
+    """)
+
+    op.alter_column("simulations", "parameters", nullable=False)
+    op.alter_column("simulations", "model_id", nullable=False)
+    op.alter_column("results", "simulation_id", nullable=False)
+
+    op.create_foreign_key(
+        op.f("results_simulation_id_fkey"),
+        "results",
+        "simulations",
+        ["simulation_id"],
+        ["id"],
+    )
+    op.drop_column("results", "model_input_id")
+    op.drop_table("model_inputs")
+    # ### end Alembic commands ###

backend/poetry.lock

@@ -39,6 +39,7 @@ mypy = "*"
 pytest = "*"
 ruff = "*"
 pytest-asyncio = "*"
+pytest-alembic = "^0.12.1"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

backend/pyproject.toml

@@ -44,23 +44,36 @@ class Simulation(Base):
     __tablename__ = "simulations"
 
     owner_id: Mapped[UUID | None] = mapped_column(Uuid)
-    model_id: Mapped[str] = mapped_column()
     concentrations: Mapped[dict[str, float]] = mapped_column(JSON)
+
+    model_inputs: Mapped[list[ModelInput]] = relationship(back_populates="simulation")
+
+
+class ModelInput(Base):
+    __tablename__ = "model_inputs"
+
+    simulation_id: Mapped[UUID] = mapped_column(ForeignKey("simulations.id"))
+    previous_model_input_id: Mapped[UUID | None] = mapped_column(
+        ForeignKey("model_inputs.id")
+    )
+    model_id: Mapped[str] = mapped_column()
     parameters: Mapped[dict[str, Any]] = mapped_column(JSON)
 
-    result: Mapped[Result | None] = relationship(back_populates="simulation")
+    simulation: Mapped[Simulation] = relationship("Simulation")
+    previous_model_input: Mapped[ModelInput | None] = relationship()
+    result: Mapped[ModelResult | None] = relationship(back_populates="model_input")
 
 
-class Result(Base):
+class ModelResult(Base):
     __tablename__ = "results"
 
-    simulation_id: Mapped[UUID] = mapped_column(ForeignKey("simulations.id"))
+    model_input_id: Mapped[UUID] = mapped_column(ForeignKey("model_inputs.id"))
     concentrations: Mapped[dict[str, float]] = mapped_column(JSON)
     panels: Mapped[list[Any]] = mapped_column(JSON)
     python_exception: Mapped[BaseException | None] = mapped_column(PickleType)
     error: Mapped[str | None] = mapped_column()
 
-    simulation: Mapped[Simulation] = relationship("Simulation")
+    model_input: Mapped[ModelInput] = relationship("ModelInput")
 
 
 class AppState(TypedDict):