Skip to content

Commit

Permalink
Add batch_id to jinja context of microbatch batches
Browse files Browse the repository at this point in the history
  • Loading branch information
QMalcolm committed Nov 21, 2024
1 parent a42303c commit 42329f1
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 2 deletions.
7 changes: 6 additions & 1 deletion core/dbt/materializations/incremental/microbatch.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def build_batches(self, start: datetime, end: datetime) -> List[BatchType]:

return batches

def build_batch_context(self, incremental_batch: bool) -> Dict[str, Any]:
def build_batch_context(self, incremental_batch: bool, start_time: datetime) -> Dict[str, Any]:
"""
Create context with entries that reflect microbatch model + incremental execution state
Expand All @@ -112,6 +112,7 @@ def build_batch_context(self, incremental_batch: bool) -> Dict[str, Any]:
batch_context["model"] = self.model.to_dict()
batch_context["sql"] = self.model.compiled_code
batch_context["compiled_code"] = self.model.compiled_code
batch_context["batch_id"] = self.batch_id(start_time=start_time)

# Add incremental context variables for batches running incrementally
if incremental_batch:
Expand Down Expand Up @@ -192,6 +193,10 @@ def truncate_timestamp(timestamp: datetime, batch_size: BatchSize) -> datetime:

return truncated

@staticmethod
def batch_id(start_time: datetime) -> str:
return start_time.strftime("%Y%M%d%H")

@staticmethod
def format_batch_start(
batch_start: Optional[datetime], batch_size: BatchSize
Expand Down
3 changes: 2 additions & 1 deletion core/dbt/task/run.py
Original file line number Diff line number Diff line change
Expand Up @@ -545,7 +545,8 @@ def _execute_microbatch_materialization(
)
# Update jinja context with batch context members
batch_context = microbatch_builder.build_batch_context(
incremental_batch=self.relation_exists
incremental_batch=self.relation_exists,
start_time=batch[0],
)
context.update(batch_context)

Expand Down

0 comments on commit 42329f1

Please sign in to comment.