[GitHub] [airflow] uranusjr commented on a change in pull request #20962: Support generating SQL script for upgrades

[GitBox <gi...@apache.org>]

uranusjr commented on a change in pull request #20962:
URL: https://github.com/apache/airflow/pull/20962#discussion_r792276176



##########
File path: airflow/migrations/versions/8646922c8a04_change_default_pool_slots_to_1.py
##########
@@ -26,48 +26,21 @@
 
 import sqlalchemy as sa
 from alembic import op
-from sqlalchemy import Column, Integer, String
-from sqlalchemy.ext.declarative import declarative_base
-
-from airflow.utils.sqlalchemy import UtcDateTime
 
 # revision identifiers, used by Alembic.
 revision = '8646922c8a04'
 down_revision = '449b4072c2da'
 branch_labels = None
 depends_on = None
 
-Base = declarative_base()
-BATCH_SIZE = 5000
-
-
-class TaskInstance(Base):  # type: ignore
-    """Minimal model definition for migrations"""
-
-    __tablename__ = "task_instance"
-
-    task_id = Column(String(), primary_key=True)
-    dag_id = Column(String(), primary_key=True)
-    execution_date = Column(UtcDateTime, primary_key=True)
-    pool_slots = Column(Integer, default=1)
-
 
 def upgrade():
     """Change default pool_slots to 1 and make pool_slots not nullable"""
-    connection = op.get_bind()
-    sessionmaker = sa.orm.sessionmaker()
-    session = sessionmaker(bind=connection)
-
-    session.query(TaskInstance).filter(TaskInstance.pool_slots.is_(None)).update(
-        {TaskInstance.pool_slots: 1}, synchronize_session=False
-    )
-    session.commit()
-
     with op.batch_alter_table("task_instance", schema=None) as batch_op:
-        batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=False)
+        batch_op.alter_column("pool_slots", existing_type=sa.Integer, nullable=False, server_default='1')

Review comment:
       We should change TaskInstance declaration to match this. (The column on TaskInstance also has `default=1` but I’m not sure if it’s added in a later migration or was accidentally left out.)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org