batch processing for updating TI UUIDs (#49015)

closes #43438

Author: rawwar

airflow-core/docs/img/airflow_erd.sha256

@@ -1 +1 @@
-7d6f2aa31fb10d8006b6b7f572bedcc4be78eb828edfd42386e0b872b6999afc
+d79700d79f51a3f70709131183df2e80e6be0f0e73ffdbcc21731890a0a030fd

airflow-core/src/airflow/config_templates/config.yml

@@ -703,6 +703,14 @@ database:
       type: string
       example: "airflow.providers.fab.auth_manager.models.db.FABDBManager"
       default: ~
+    migration_batch_size:
+      description: |
+        The number of rows to process in each batch when performing a migration.
+        This is useful for large tables to avoid locking and failure due to query timeouts.
+      version_added: 3.0.0
+      type: integer
+      example: ~
+      default: "10000"
 logging:
   description: ~
   options:

airflow-core/src/airflow/migrations/versions/0042_3_0_0_add_uuid_primary_key_to_task_instance_.py

@@ -31,6 +31,8 @@
 from sqlalchemy import text
 from sqlalchemy.dialects import postgresql
 
+from airflow.configuration import conf
+
 # revision identifiers, used by Alembic.
 revision = "d59cbbef95eb"
 down_revision = "05234396c6fc"
@@ -194,6 +196,7 @@ def create_foreign_keys():
 
 def upgrade():
     """Add UUID primary key to task instance table."""
+    batch_size = conf.getint("database", "migration_batch_size", fallback=1000)
     conn = op.get_bind()
     dialect_name = conn.dialect.name
 
@@ -202,12 +205,28 @@ def upgrade():
     if dialect_name == "postgresql":
         op.execute(pg_uuid7_fn)
 
-        # TODO: Add batching to handle updates in smaller chunks for large tables to avoid locking
         # Migrate existing rows with UUID v7 using a timestamp-based generation
-        op.execute(
-            "UPDATE task_instance SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, clock_timestamp()))"
-        )
-
+        while True:
+            result = conn.execute(
+                text(
+                    """
+                    WITH cte AS (
+                        SELECT ctid
+                        FROM task_instance
+                        WHERE id IS NULL
+                        LIMIT :batch_size
+                    )
+                    UPDATE task_instance
+                    SET id = uuid_generate_v7(coalesce(queued_dttm, start_date, clock_timestamp()))
+                    FROM cte
+                    WHERE task_instance.ctid = cte.ctid
+                    """
+                ).bindparams(batch_size=batch_size)
+            )
+            row_count = result.rowcount
+            if row_count == 0:
+                break
+            print(f"Migrated {row_count} task_instance rows in this batch...")
         op.execute(pg_uuid7_fn_drop)
 
         # Drop existing primary key constraint to task_instance table