Add check constrains on collection status table

augur/application/db/models/augur_operations.py

@@ -1,5 +1,5 @@
 # coding: utf-8
-from sqlalchemy import BigInteger, SmallInteger, Column, Index, Integer, String, Table, text, UniqueConstraint, Boolean, ForeignKey, update
+from sqlalchemy import BigInteger, SmallInteger, Column, Index, Integer, String, Table, text, UniqueConstraint, Boolean, ForeignKey, update, CheckConstraint
 from sqlalchemy.dialects.postgresql import TIMESTAMP, UUID
 from sqlalchemy.orm.exc import NoResultFound, MultipleResultsFound
 from sqlalchemy.exc import IntegrityError
@@ -996,6 +996,47 @@ def create(session, user_session_token_id):
 class CollectionStatus(Base):
     __tablename__ = "collection_status"
     __table_args__ = (
+        CheckConstraint(
+            "(core_data_last_collected IS NOT NULL AND core_status = 'Success') OR "
+            "(core_data_last_collected IS NULL AND core_status = 'Pending') OR "
+            "(core_status = 'Error') OR "
+            "(core_status = 'Collecting')",
+            name='core_data_last_collected_check'
+        ),
+        CheckConstraint(
+            "(core_task_id IS NULL AND core_status IN ('Pending', 'Success', 'Error')) OR "
+            "(core_task_id IS NOT NULL AND core_status = 'Collecting')",
+            name='core_task_id_check'
+        ),
+        CheckConstraint(
+            "(secondary_data_last_collected IS NOT NULL AND secondary_status = 'Success') OR "
+            "(secondary_data_last_collected IS NULL AND secondary_status = 'Pending') OR"
+            "(secondary_status = 'Error') OR "
+            "(secondary_status = 'Collecting')",
+            name='secondary_data_last_collected_check'
+        ),
+        CheckConstraint(
+            "(secondary_task_id IS NULL AND secondary_status IN ('Pending', 'Success', 'Error')) OR "
+            "(secondary_task_id IS NOT NULL AND secondary_status = 'Collecting')",
+            name='secondary_task_id_check'
+        ),
+        CheckConstraint(
+            "(facade_data_last_collected IS NOT NULL AND facade_status IN ('Success', 'Update')) OR "
+            "(facade_data_last_collected IS NULL AND facade_status = 'Pending') OR "
+            "(facade_status = 'Error') OR "
+            "(facade_status = 'Collecting')",
+            name='facade_data_last_collected_check'
+        ),
+        CheckConstraint(
+            "(facade_task_id IS NULL AND facade_status IN ('Pending', 'Success', 'Error', 'Failed Clone')) OR "
+            "(facade_task_id IS NOT NULL AND facade_status = 'Collecting')",
+            name='facade_task_id_check'
+        ),
+        CheckConstraint(
+            "(core_status = 'Success') OR "
+            "(core_status IN ('Pending', 'Collecting', 'Error') AND secondary_status = 'Pending')",
+            name='core_secondary_dependency_check'
+        ),
         {"schema": "augur_operations"}
     )
 

augur/application/schema/alembic/versions/17_add_collection_status_constraints.py

@@ -0,0 +1,77 @@
+"""Add collection status constraints
+
+Revision ID: 17
+Revises: 16
+Create Date: 2023-04-19 14:30:14.709349
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '17'
+down_revision = '16'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_check_constraint(
+        constraint_name="core_data_last_collected_check",
+        table_name="collection_status",
+        condition="(core_status = 'Error') OR "
+        "(core_status = 'Collecting') OR "
+        "(core_data_last_collected IS NOT NULL AND core_status = 'Success') OR "
+        "(core_data_last_collected IS NULL AND core_status = 'Pending')"
+    )
+    op.create_check_constraint(
+        constraint_name="core_task_id_check",
+        table_name="collection_status",
+        condition="(core_task_id IS NULL AND core_status IN ('Pending', 'Success', 'Error')) OR "
+        "(core_task_id IS NOT NULL AND core_status = 'Collecting')"
+    )
+    op.create_check_constraint(
+        constraint_name="secondary_data_last_collected_check",
+        table_name="collection_status",
+        condition="(secondary_status = 'Error') OR "
+        "(secondary_status = 'Collecting') OR "
+        "(secondary_data_last_collected IS NOT NULL AND secondary_status = 'Success') OR "
+        "(secondary_data_last_collected IS NULL AND secondary_status = 'Pending')"
+    )
+    op.create_check_constraint(
+        constraint_name="secondary_task_id_check",
+        table_name="collection_status",
+        condition="(secondary_task_id IS NULL AND secondary_status IN ('Pending', 'Success', 'Error')) OR "
+        "(secondary_task_id IS NOT NULL AND secondary_status = 'Collecting')"
+    )
+    op.create_check_constraint(
+        constraint_name="facade_data_last_collected_check",
+        table_name="collection_status",
+        condition="(facade_status = 'Error') OR "
+        "(facade_data_last_collected IS NOT NULL AND facade_status IN ('Success', 'Update')) OR "
+        "(facade_data_last_collected IS NULL AND facade_status = 'Pending')"
+    )
+    op.create_check_constraint(
+        constraint_name="facade_task_id_check",
+        table_name="collection_status",
+        condition="(facade_task_id IS NULL AND facade_status IN ('Pending', 'Success', 'Error', 'Failed Clone')) OR "
+        "(facade_task_id IS NOT NULL AND facade_status = 'Collecting')"
+    )
+    op.create_check_constraint(
+        constraint_name="core_secondary_dependency_check",
+        table_name="collection_status",
+        condition="(core_status = 'Success') OR "
+        "(core_status IN ('Pending', 'Collecting', 'Error') AND secondary_status = 'Pending')"
+    )
+
+
+def downgrade():
+    op.drop_constraint("core_data_last_collected_check", "collection_status")
+    op.drop_constraint("core_task_id_check", "collection_status")
+    op.drop_constraint("secondary_data_last_collected_check", "collection_status")
+    op.drop_constraint("secondary_task_id_check", "collection_status")
+    op.drop_constraint("facade_data_last_collected_check", "collection_status")
+    op.drop_constraint("facade_task_id_check", "collection_status")
+    op.drop_constraint("core_secondary_dependency_check", "collection_status")
+