[GitHub] [airflow] uranusjr commented on a diff in pull request #24585: Fix migration 0080_2_0_2 - Replace null values before setting column not null

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

uranusjr commented on code in PR #24585:
URL: https://github.com/apache/airflow/pull/24585#discussion_r903409776


##########
airflow/migrations/versions/0080_2_0_2_change_default_pool_slots_to_1.py:
##########
@@ -35,8 +35,24 @@
 airflow_version = '2.0.2'
 
 
+# Minimal required schema
+task_instance = sa.table(
+    'task_instance',
+    sa.column('pool_slots', sa.Integer),
+)
+
+
 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(task_instance) \
+        .filter(task_instance.c.pool_slots.is_(None)) \
+        .update({task_instance.c.pool_slots: 1}, synchronize_session=False)
+    session.commit()

Review Comment:
   It’s much easier to just write SQL here
   
   ```suggestion
       op.get_bind("UPDATE task_instance SET pool_slots = 1 WHERE pool_slots IS NULL")
   ```



-- 
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