Fix migration

Author: PythonCoderAS

alembic/versions/90db4a933a16_add_batch_locking.py

@@ -0,0 +1,34 @@
+"""Add batch locking
+
+Revision ID: 90db4a933a16
+Revises: 9a098a265cbd
+Create Date: 2024-01-18 14:21:22.487177
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "90db4a933a16"
+down_revision: Union[str, None] = "9a098a265cbd"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.add_column(
+        "batches", sa.Column("locked", sa.DateTime(timezone=True), nullable=True)
+    )
+    op.create_index(op.f("ix_batches_locked"), "batches", ["locked"], unique=False)
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.drop_index(op.f("ix_batches_locked"), table_name="batches")
+    op.drop_column("batches", "locked")
+    # ### end Alembic commands ###

src/models.py

@@ -90,8 +90,8 @@ class Batch(Base):
         init=False,
         index=True,
     )
-    locked: Mapped[bool] = mapped_column(
-        default=False
+    locked: Mapped[datetime.datetime | None] = mapped_column(
+        sqlalchemy.DateTime(timezone=True), default=None, nullable=True, index=True
     )  # Indicates that a batch is locked (no more jobs can be added to it)
 
     jobs: Mapped[list["Job"]] = sqlalchemy.orm.relationship(
@@ -188,6 +188,12 @@ class BatchJobs(Base):
         sqlalchemy.UniqueConstraint("batch_id", "job_id", name="_batch_job_uc"),
     )
 
+    @sqlalchemy.orm.validates("batch")
+    def validate_not_locked_batch(self, key: str, batch: Batch) -> Batch:
+        if batch.locked:
+            raise ValueError("Batch is locked")
+        return batch
+
 
 class Job(Base):
     __tablename__ = "jobs"
@@ -228,3 +234,9 @@ class Job(Base):
     failed: Mapped[datetime.datetime | None] = mapped_column(
         sqlalchemy.DateTime(timezone=True), default=None, nullable=True, index=True
     )  # If a job has failed, this is the time it failed at
+
+    @sqlalchemy.orm.validates("batches")
+    def validate_not_locked_batch(self, key: str, batch: Batch) -> Batch:
+        if batch.locked:
+            raise ValueError("Batch is locked")
+        return batch