Skip to content

Commit

Permalink
Revert back to using the MicrobatchModelRunner initializer directly
Browse files Browse the repository at this point in the history
We started using self.get_runner in a pervious commit. Unfornuately
there were unintended consequences in doing so. Namely, incorrectly
incrementing the number of nodes being run. We do _eventually_ want
to move to using self.get_runner, but it is out of scope for this
segement of work (mostly due to time constraints).
  • Loading branch information
QMalcolm committed Dec 6, 2024
1 parent 74a76c4 commit 0c16d07
Showing 1 changed file with 10 additions and 2 deletions.
12 changes: 10 additions & 2 deletions core/dbt/task/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -753,6 +753,7 @@ def handle_microbatch_model(
# Run first batch not in parallel
relation_exists = self._submit_batch(
node=node,
adapter=runner.adapter,
relation_exists=relation_exists,
batches=batches,
batch_idx=batch_idx,
Expand All @@ -767,6 +768,7 @@ def handle_microbatch_model(
while batch_idx < len(runner.batches) - 1:
relation_exists = self._submit_batch(
node=node,
adapter=runner.adapter,
relation_exists=relation_exists,
batches=batches,
batch_idx=batch_idx,
Expand All @@ -782,6 +784,7 @@ def handle_microbatch_model(
# Final batch runs once all others complete to ensure post_hook runs at the end
self._submit_batch(
node=node,
adapter=runner.adapter,
relation_exists=relation_exists,
batches=batches,
batch_idx=batch_idx,
Expand All @@ -801,6 +804,7 @@ def handle_microbatch_model(
def _submit_batch(
self,
node: ModelNode,
adapter: BaseAdapter,
relation_exists: bool,
batches: Dict[int, BatchType],
batch_idx: int,
Expand All @@ -818,8 +822,12 @@ def _submit_batch(
if batch_idx != len(batches) - 1:
node_copy.config.post_hook = []

batch_runner = self.get_runner(node_copy)
assert isinstance(batch_runner, MicrobatchModelRunner)
# TODO: We should be doing self.get_runner, however doing so
# currently causes the tracking of how many nodes there are to
# increment when we don't want it to
batch_runner = MicrobatchModelRunner(
self.config, adapter, node_copy, self.run_count, self.num_nodes
)
batch_runner.set_batch_idx(batch_idx)
batch_runner.set_relation_exists(relation_exists)
batch_runner.set_batches(batches)
Expand Down

0 comments on commit 0c16d07

Please sign in to comment.