feat: dbt microbatch ref filter support

Author: eakmanrq

sqlmesh/core/renderer.py

@@ -214,6 +214,13 @@ def _resolve_table(table: str | exp.Table) -> str:
                         dialect=self._dialect, identify=True, comments=False
                     )
 
+                all_refs = list(self._jinja_macro_registry.global_objs.get("sources", {}).values()) + list(  # type: ignore
+                    self._jinja_macro_registry.global_objs.get("refs", {}).values()  # type: ignore
+                )
+                for ref in all_refs:
+                    if ref.event_time_filter:
+                        ref.event_time_filter["start"] = render_kwargs["start_tstz"]
+                        ref.event_time_filter["end"] = render_kwargs["end_tstz"]
                 jinja_env = self._jinja_macro_registry.build_environment(**jinja_env_kwargs)
 
                 expressions = []

sqlmesh/dbt/__init__.py

@@ -2,3 +2,25 @@
     create_builtin_filters as create_builtin_filters,
     create_builtin_globals as create_builtin_globals,
 )
+from sqlmesh.dbt.util import DBT_VERSION
+
+
+if DBT_VERSION >= (1, 9, 0):
+    from dbt.adapters.base.relation import BaseRelation, EventTimeFilter
+
+    def _render_event_time_filtered_inclusive(
+        self: BaseRelation, event_time_filter: EventTimeFilter
+    ) -> str:
+        """
+        Returns "" if start and end are both None
+        """
+        filter = ""
+        if event_time_filter.start and event_time_filter.end:
+            filter = f"{event_time_filter.field_name} BETWEEN '{event_time_filter.start}' and '{event_time_filter.end}'"
+        elif event_time_filter.start:
+            filter = f"{event_time_filter.field_name} >= '{event_time_filter.start}'"
+        elif event_time_filter.end:
+            filter = f"{event_time_filter.field_name} <= '{event_time_filter.end}'"
+        return filter
+
+    BaseRelation._render_event_time_filtered = _render_event_time_filtered_inclusive  # type: ignore

sqlmesh/dbt/basemodel.py

@@ -28,6 +28,7 @@
 )
 from sqlmesh.dbt.relation import Policy, RelationType
 from sqlmesh.dbt.test import TestConfig
+from sqlmesh.dbt.util import DBT_VERSION
 from sqlmesh.utils import AttributeDict
 from sqlmesh.utils.errors import ConfigError
 from sqlmesh.utils.pydantic import field_validator
@@ -130,6 +131,7 @@ class BaseModelConfig(GeneralConfig):
     grants: t.Dict[str, t.List[str]] = {}
     columns: t.Dict[str, ColumnConfig] = {}
     quoting: t.Dict[str, t.Optional[bool]] = {}
+    event_time: t.Optional[str] = None
 
     version: t.Optional[int] = None
     latest_version: t.Optional[int] = None
@@ -222,13 +224,20 @@ def relation_info(self) -> AttributeDict[str, t.Any]:
         else:
             relation_type = RelationType.Table
 
+        extras = {}
+        if DBT_VERSION >= (1, 9, 0) and self.event_time:
+            extras["event_time_filter"] = {
+                "field_name": self.event_time,
+            }
+
         return AttributeDict(
             {
                 "database": self.database,
                 "schema": self.table_schema,
                 "identifier": self.table_name,
                 "type": relation_type.value,
                 "quote_policy": AttributeDict(self.quoting),
+                **extras,
             }
         )
 

sqlmesh/dbt/source.py

@@ -8,6 +8,7 @@
 from sqlmesh.dbt.column import ColumnConfig
 from sqlmesh.dbt.common import GeneralConfig
 from sqlmesh.dbt.relation import RelationType
+from sqlmesh.dbt.util import DBT_VERSION
 from sqlmesh.utils import AttributeDict
 from sqlmesh.utils.errors import ConfigError
 
@@ -46,6 +47,7 @@ class SourceConfig(GeneralConfig):
     external: t.Optional[t.Dict[str, t.Any]] = {}
     source_meta: t.Optional[t.Dict[str, t.Any]] = {}
     columns: t.Dict[str, ColumnConfig] = {}
+    event_time: t.Optional[str] = None
 
     _canonical_name: t.Optional[str] = None
 
@@ -94,6 +96,11 @@ def relation_info(self) -> AttributeDict:
         if external_location:
             extras["external"] = external_location.replace("{name}", self.table_name)
 
+        if DBT_VERSION >= (1, 9, 0) and self.event_time:
+            extras["event_time_filter"] = {
+                "field_name": self.event_time,
+            }
+
         return AttributeDict(
             {
                 "database": self.database,

tests/core/test_snapshot.py

@@ -1078,7 +1078,9 @@ def test_fingerprint_jinja_macros_global_objs(model: Model, global_obj_key: str)
     )
     fingerprint = fingerprint_from_node(model, nodes={})
     model = model.copy()
-    model.jinja_macros.global_objs[global_obj_key] = AttributeDict({"test": "test"})
+    model.jinja_macros.global_objs[global_obj_key] = AttributeDict(
+        {"test": AttributeDict({"test": "test"})}
+    )
     updated_fingerprint = fingerprint_from_node(model, nodes={})
     assert updated_fingerprint.data_hash != fingerprint.data_hash
     assert updated_fingerprint.metadata_hash == fingerprint.metadata_hash

tests/dbt/test_model.py

@@ -298,3 +298,135 @@ def test_load_microbatch_required_only(
         column=exp.to_column("ds", quoted=True), format="%Y-%m-%d"
     )
     assert model.kind.batch_size is None
+
+
+@pytest.mark.slow
+def test_load_microbatch_with_ref(
+    tmp_path: Path, caplog, dbt_dummy_postgres_config: PostgresConfig, create_empty_project
+) -> None:
+    yaml = YAML()
+    project_dir, model_dir = create_empty_project()
+    source_schema = {
+        "version": 2,
+        "sources": [
+            {
+                "name": "my_source",
+                "tables": [{"name": "my_table", "config": {"event_time": "ds"}}],
+            }
+        ],
+    }
+    source_schema_file = model_dir / "source_schema.yml"
+    with open(source_schema_file, "w", encoding="utf-8") as f:
+        yaml.dump(source_schema, f)
+    # add `tests` to model config since this is loaded by dbt and ignored and we shouldn't error when loading it
+    microbatch_contents = """
+    {{
+        config(
+            materialized='incremental',
+            incremental_strategy='microbatch',
+            event_time='ds',
+            begin='2020-01-01',
+            batch_size='day'
+        )
+    }}
+
+    SELECT cola, ds FROM {{ source('my_source', 'my_table') }}
+    """
+    microbatch_model_file = model_dir / "microbatch.sql"
+    with open(microbatch_model_file, "w", encoding="utf-8") as f:
+        f.write(microbatch_contents)
+
+    microbatch_two_contents = """
+    {{
+        config(
+            materialized='incremental',
+            incremental_strategy='microbatch',
+            event_time='ds',
+            begin='2020-01-05',
+            batch_size='day'
+        )
+    }}
+
+    SELECT cola, ds FROM {{ ref('microbatch') }}
+    """
+    microbatch_two_model_file = model_dir / "microbatch_two.sql"
+    with open(microbatch_two_model_file, "w", encoding="utf-8") as f:
+        f.write(microbatch_two_contents)
+
+    microbatch_snapshot_fqn = '"local"."main"."microbatch"'
+    microbatch_two_snapshot_fqn = '"local"."main"."microbatch_two"'
+    context = Context(paths=project_dir)
+    assert (
+        context.render(microbatch_snapshot_fqn, start="2025-01-01", end="2025-01-10").sql()
+        == 'SELECT "cola" AS "cola", "ds" AS "ds" FROM (SELECT * FROM "local"."my_source"."my_table" AS "my_table" WHERE "ds" BETWEEN \'2025-01-01 00:00:00+00:00\' AND \'2025-01-10 23:59:59.999999+00:00\') AS "_q_0"'
+    )
+    assert (
+        context.render(microbatch_two_snapshot_fqn, start="2025-01-01", end="2025-01-10").sql()
+        == 'SELECT "_q_0"."cola" AS "cola", "_q_0"."ds" AS "ds" FROM (SELECT "microbatch"."cola" AS "cola", "microbatch"."ds" AS "ds" FROM "local"."main"."microbatch" AS "microbatch" WHERE "microbatch"."ds" <= \'2025-01-10 23:59:59.999999+00:00\' AND "microbatch"."ds" >= \'2025-01-01 00:00:00+00:00\') AS "_q_0"'
+    )
+
+
+@pytest.mark.slow
+def test_load_microbatch_with_ref_no_filter(
+    tmp_path: Path, caplog, dbt_dummy_postgres_config: PostgresConfig, create_empty_project
+) -> None:
+    yaml = YAML()
+    project_dir, model_dir = create_empty_project()
+    source_schema = {
+        "version": 2,
+        "sources": [
+            {
+                "name": "my_source",
+                "tables": [{"name": "my_table", "config": {"event_time": "ds"}}],
+            }
+        ],
+    }
+    source_schema_file = model_dir / "source_schema.yml"
+    with open(source_schema_file, "w", encoding="utf-8") as f:
+        yaml.dump(source_schema, f)
+    # add `tests` to model config since this is loaded by dbt and ignored and we shouldn't error when loading it
+    microbatch_contents = """
+    {{
+        config(
+            materialized='incremental',
+            incremental_strategy='microbatch',
+            event_time='ds',
+            begin='2020-01-01',
+            batch_size='day'
+        )
+    }}
+
+    SELECT cola, ds FROM {{ source('my_source', 'my_table').render() }}
+    """
+    microbatch_model_file = model_dir / "microbatch.sql"
+    with open(microbatch_model_file, "w", encoding="utf-8") as f:
+        f.write(microbatch_contents)
+
+    microbatch_two_contents = """
+    {{
+        config(
+            materialized='incremental',
+            incremental_strategy='microbatch',
+            event_time='ds',
+            begin='2020-01-01',
+            batch_size='day'
+        )
+    }}
+
+    SELECT cola, ds FROM {{ ref('microbatch').render() }}
+    """
+    microbatch_two_model_file = model_dir / "microbatch_two.sql"
+    with open(microbatch_two_model_file, "w", encoding="utf-8") as f:
+        f.write(microbatch_two_contents)
+
+    microbatch_snapshot_fqn = '"local"."main"."microbatch"'
+    microbatch_two_snapshot_fqn = '"local"."main"."microbatch_two"'
+    context = Context(paths=project_dir)
+    assert (
+        context.render(microbatch_snapshot_fqn, start="2025-01-01", end="2025-01-10").sql()
+        == 'SELECT "cola" AS "cola", "ds" AS "ds" FROM "local"."my_source"."my_table" AS "my_table"'
+    )
+    assert (
+        context.render(microbatch_two_snapshot_fqn, start="2025-01-01", end="2025-01-10").sql()
+        == 'SELECT "microbatch"."cola" AS "cola", "microbatch"."ds" AS "ds" FROM "local"."main"."microbatch" AS "microbatch"'
+    )