fix: fix mutation types in delete DAG

dags/deletes.py

@@ -292,17 +292,28 @@ def count_pending_deletes(client: Client) -> int:
 
 @op
 def wait_for_delete_mutations(
+    context: OpExecutionContext,
     cluster: ResourceParam[ClickhouseCluster],
-    delete_person_events: tuple[PendingPersonEventDeletesTable, Mutation],
+    delete_person_events: tuple[PendingPersonEventDeletesTable, ShardMutations],
 ) -> PendingPersonEventDeletesTable:
     pending_person_deletions, shard_mutations = delete_person_events
 
-    if not shard_mutations:
-        return pending_person_deletions
+    results = cluster.map_all_hosts_in_shards(
+        {shard: mutation.wait for shard, mutation in shard_mutations.items()}
+    ).as_completed()
+    try:
+        for shard, result in results:
+            result.result()
+            context.log.info(
+                f"Shard {shard.shard_num if shard.shard_num else 'unknown'} replica {shard.replica_num if shard.replica_num else 'unknown'} completed "
+            )
+    except Exception as e:
+        context.log.exception(
+            f"Shard {shard.shard_num if shard.shard_num else 'unknown'} replica {shard.replica_num if shard.replica_num else 'unknown'} failed: {e}"
+        )
+        raise
 
-    [table, mutations] = delete_person_events
-    cluster.map_all_hosts(mutations.wait).result()
-    return table
+    return pending_person_deletions
 
 
 @op

posthog/clickhouse/cluster.py

@@ -159,6 +159,29 @@ def map_all_hosts_in_shard(
                 }
             )
 
+    def map_all_hosts_in_shards(
+        self, shard_fns: dict[int, Callable[[Client], T]], concurrency: int | None = None
+    ) -> FuturesMap[HostInfo, T]:
+        """
+        Execute the callable once for each host in the specified shards.
+
+        The number of concurrent queries can limited with the ``concurrency`` parameter, or set to ``None`` to use the
+        default limit of the executor.
+
+        Wait for all to return before returning upon ``.values()``
+        """
+
+        shard_host_fn = {}
+        for shard, fn in shard_fns.items():
+            for host in self.__hosts:
+                if host.shard_num == shard:
+                    shard_host_fn[host] = fn
+
+        with ThreadPoolExecutor(max_workers=concurrency) as executor:
+            return FuturesMap(
+                {host: executor.submit(self.__get_task_function(host, fn)) for host, fn in shard_host_fn.items()}
+            )
+
     def map_one_host_per_shard(
         self, fn: Callable[[Client], T], concurrency: int | None = None
     ) -> FuturesMap[HostInfo, T]: