Add constraint to ensure task map length >= 0 (#21115)

apache · Jan 26, 2022 · 80f30ee · 80f30ee
1 parent dff536e
commit 80f30ee
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py b/airflow/migrations/versions/e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py
@@ -24,7 +24,7 @@
 """
 
 from alembic import op
-from sqlalchemy import Column, ForeignKeyConstraint, Integer, text
+from sqlalchemy import CheckConstraint, Column, ForeignKeyConstraint, Integer, text
 
 from airflow.models.base import StringID
 from airflow.utils.sqlalchemy import ExtendedJSON
@@ -75,6 +75,7 @@ def upgrade():
         Column("map_index", Integer, primary_key=True),
         Column("length", Integer, nullable=False),
         Column("keys", ExtendedJSON, nullable=True),
+        CheckConstraint("length >= 0", name="task_map_length_not_negative"),
         ForeignKeyConstraint(
             ["dag_id", "task_id", "run_id", "map_index"],
             [

diff --git a/airflow/models/taskmap.py b/airflow/models/taskmap.py
@@ -22,7 +22,7 @@
 import enum
 from typing import TYPE_CHECKING, Any, Collection, List, Optional
 
-from sqlalchemy import Column, ForeignKeyConstraint, Integer, String
+from sqlalchemy import CheckConstraint, Column, ForeignKeyConstraint, Integer, String
 
 from airflow.models.base import COLLATION_ARGS, ID_LEN, Base
 from airflow.utils.sqlalchemy import ExtendedJSON
@@ -61,6 +61,7 @@ class TaskMap(Base):
     keys = Column(ExtendedJSON, nullable=True)
 
     __table_args__ = (
+        CheckConstraint(length >= 0, name="task_map_length_not_negative"),
         ForeignKeyConstraint(
             [dag_id, task_id, run_id, map_index],
             [