feat: added non-negative check to sql migration

yeagerai · Sep 9, 2024 · c9d1cb0 · c9d1cb0
1 parent 06b41d6
commit c9d1cb0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/...d/database_handler/migration/versions/f9636f013003_added_balance_attribute_to_accounts.py b/...d/database_handler/migration/versions/f9636f013003_added_balance_attribute_to_accounts.py
@@ -24,6 +24,9 @@ def upgrade() -> None:
     op.add_column("current_state", sa.Column("balance", sa.Integer(), nullable=True))
     op.execute("UPDATE current_state SET balance = 0 WHERE balance IS NULL")
     op.alter_column("current_state", "balance", nullable=False)
+    op.create_check_constraint(
+        "check_balance_non_negative", "current_state", "balance >= 0"
+    )
     op.alter_column(
         "transactions",
         "status",

diff --git a/backend/database_handler/models.py b/backend/database_handler/models.py
@@ -37,11 +37,14 @@ class Base(MappedAsDataclass, DeclarativeBase):
 
 class CurrentState(Base):
     __tablename__ = "current_state"
-    __table_args__ = (PrimaryKeyConstraint("id", name="current_state_pkey"),)
+    __table_args__ = (
+        PrimaryKeyConstraint("id", name="current_state_pkey"),
+        CheckConstraint("balance >= 0", name="check_balance_non_negative"),
+    )
 
     id: Mapped[str] = mapped_column(String(255), primary_key=True)
     data: Mapped[dict] = mapped_column(JSONB)
-    balance: Mapped[int] = mapped_column(Integer, default=0)
+    balance: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
     updated_at: Mapped[Optional[datetime.datetime]] = mapped_column(
         DateTime(True),
         init=False,