From ff218d2adec667593ec829bc4393411c54a2d1a5 Mon Sep 17 00:00:00 2001 From: Tzu-ping Chung Date: Mon, 17 Jan 2022 18:45:13 +0800 Subject: [PATCH] Switch map_index back to use server_default This emits 'DEFAULT -1' in the column definition and helps migration. --- .../e655c0453f75_add_taskmap_and_map_id_on_taskinstance.py | 6 +++--- airflow/models/taskinstance.py | 3 ++- airflow/models/taskreschedule.py | 4 ++-- 3 files changed, 7 insertions(+), 6 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 index 91afcb99e6f39..106dffc035307 100644 --- 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 +from sqlalchemy import Column, ForeignKeyConstraint, Integer, text from airflow.models.base import StringID from airflow.utils.sqlalchemy import ExtendedJSON @@ -47,12 +47,12 @@ def upgrade(): with op.batch_alter_table("task_instance") as batch_op: # I think we always use this name for TaskInstance after 7b2661a43ba3? batch_op.drop_constraint("task_instance_pkey", type_="primary") - batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1)) + batch_op.add_column(Column("map_index", Integer, nullable=False, server_default=text("-1"))) batch_op.create_primary_key("task_instance_pkey", ["dag_id", "task_id", "run_id", "map_index"]) # Re-create task_reschedule's constraints. with op.batch_alter_table("task_reschedule") as batch_op: - batch_op.add_column(Column("map_index", Integer, nullable=False, default=-1)) + batch_op.add_column(Column("map_index", Integer, nullable=False, server_default=text("-1"))) batch_op.create_foreign_key( "task_reschedule_ti_fkey", "task_instance", diff --git a/airflow/models/taskinstance.py b/airflow/models/taskinstance.py index 6dd026b600ded..d8b7094d952a7 100644 --- a/airflow/models/taskinstance.py +++ b/airflow/models/taskinstance.py @@ -60,6 +60,7 @@ func, inspect, or_, + text, tuple_, ) from sqlalchemy.ext.associationproxy import association_proxy @@ -343,7 +344,7 @@ class TaskInstance(Base, LoggingMixin): task_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) dag_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) run_id = Column(String(ID_LEN, **COLLATION_ARGS), primary_key=True, nullable=False) - map_index = Column(Integer, primary_key=True, nullable=False, default=-1) + map_index = Column(Integer, primary_key=True, nullable=False, server_default=text("-1")) start_date = Column(UtcDateTime) end_date = Column(UtcDateTime) diff --git a/airflow/models/taskreschedule.py b/airflow/models/taskreschedule.py index 4e39af5beea4f..01f3aaeeab655 100644 --- a/airflow/models/taskreschedule.py +++ b/airflow/models/taskreschedule.py @@ -20,7 +20,7 @@ import datetime from typing import TYPE_CHECKING -from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc +from sqlalchemy import Column, ForeignKeyConstraint, Index, Integer, String, asc, desc, text from sqlalchemy.ext.associationproxy import association_proxy from sqlalchemy.orm import relationship @@ -41,7 +41,7 @@ class TaskReschedule(Base): task_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) dag_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) run_id = Column(String(ID_LEN, **COLLATION_ARGS), nullable=False) - map_index = Column(Integer, nullable=False, default=-1) + map_index = Column(Integer, nullable=False, server_default=text("-1")) try_number = Column(Integer, nullable=False) start_date = Column(UtcDateTime, nullable=False) end_date = Column(UtcDateTime, nullable=False)