From 071e4221721c9fbf7e635ced6f767d21af3f527b Mon Sep 17 00:00:00 2001
From: Staci Mullins <63313398+stacimc@users.noreply.github.com>
Date: Thu, 31 Oct 2024 13:28:39 -0700
Subject: [PATCH] Limit to 2 `alter_data_batches` concurrently (#5125)

* Limit to 10 alter_data_batches concurrently

* Reduce to 2 concurrent batches
---
 catalog/dags/data_refresh/alter_data.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/catalog/dags/data_refresh/alter_data.py b/catalog/dags/data_refresh/alter_data.py
index 0c490a2c149..1a9481da4ab 100644
--- a/catalog/dags/data_refresh/alter_data.py
+++ b/catalog/dags/data_refresh/alter_data.py
@@ -124,7 +124,7 @@ def get_alter_batches(
     return [(x, x + batch_size - 1) for x in range(start, stop, batch_size)]
 
 
-@task
+@task(max_active_tis_per_dagrun=2)
 def alter_data_batch(
     batch: tuple[int, int],
     temp_table: str,
@@ -203,7 +203,8 @@ def alter_table_data(
     )
 
     alter_data = alter_data_batch.partial(
-        temp_table=temp_table, postgres_conn_id=postgres_conn_id
+        temp_table=temp_table,
+        postgres_conn_id=postgres_conn_id,
     ).expand(batch=batches)
 
     report(alter_data)