Rewrite execution of microbatch models to avoid blocking the main thread

.changes/unreleased/Fixes-20250303-131440.yaml

@@ -0,0 +1,6 @@
+kind: Fixes
+body: Fix microbatch execution to not block main thread nor hang
+time: 2025-03-03T13:14:40.432874-06:00
+custom:
+  Author: QMalcolm
+  Issue: 11243 11306

core/dbt/graph/thread_pool.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from multiprocessing.pool import ThreadPool
+
+
+class DbtThreadPool(ThreadPool):
+    """A ThreadPool that tracks whether or not it's been closed"""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.closed = False
+
+    def close(self):
+        self.closed = True
+        super().close()
+
+    def is_closed(self):
+        return self.closed

core/dbt/materializations/incremental/microbatch.py

@@ -100,7 +100,8 @@ def build_batches(self, start: datetime, end: datetime) -> List[BatchType]:
 
         return batches
 
-    def build_jinja_context_for_batch(self, incremental_batch: bool) -> Dict[str, Any]:
+    @staticmethod
+    def build_jinja_context_for_batch(model: ModelNode, incremental_batch: bool) -> Dict[str, Any]:
         """
         Create context with entries that reflect microbatch model + incremental execution state
 
@@ -109,9 +110,9 @@ def build_jinja_context_for_batch(self, incremental_batch: bool) -> Dict[str, An
         jinja_context: Dict[str, Any] = {}
 
         # Microbatch model properties
-        jinja_context["model"] = self.model.to_dict()
-        jinja_context["sql"] = self.model.compiled_code
-        jinja_context["compiled_code"] = self.model.compiled_code
+        jinja_context["model"] = model.to_dict()
+        jinja_context["sql"] = model.compiled_code
+        jinja_context["compiled_code"] = model.compiled_code
 
         # Add incremental context variables for batches running incrementally
         if incremental_batch:

core/dbt/task/build.py

@@ -131,7 +131,8 @@
             runner.do_skip(cause=cause)
 
         if isinstance(runner, MicrobatchModelRunner):
-            return self.handle_microbatch_model(runner, pool)
+            runner.set_parent_task(self)
+            runner.set_pool(pool)
 
         return self.call_runner(runner)
 
@@ -146,10 +147,11 @@
             runner.do_skip(cause=cause)
 
         if isinstance(runner, MicrobatchModelRunner):
-            callback(self.handle_microbatch_model(runner, pool))
-        else:
-            args = [runner]
-            self._submit(pool, args, callback)
+            runner.set_parent_task(self)
+            runner.set_pool(pool)
+
+        args = [runner]
+        self._submit(pool, args, callback)
 
     # Make a map of model unique_ids to selected unit test unique_ids,
     # for processing before the model.