Feat: Support the '_dbt_max_partition' variable for BigQuery dbt projects (#1195)

Author: izeigerman

sqlmesh/dbt/adapter.py

@@ -79,6 +79,14 @@ def execute(
     ) -> t.Tuple[AdapterResponse, agate.Table]:
         """Executes the given SQL statement and returns the results as an agate table."""
 
+    @abc.abstractmethod
+    def resolve_schema(self, relation: BaseRelation) -> t.Optional[str]:
+        """Resolves the relation's schema to its physical schema."""
+
+    @abc.abstractmethod
+    def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
+        """Resolves the relation's schema to its physical identifier."""
+
     def quote(self, identifier: str) -> str:
         """Returns a quoted identifier."""
         return exp.to_column(identifier).sql(dialect=self.dialect, identify=True)
@@ -138,6 +146,12 @@ def execute(
         self._raise_parsetime_adapter_call_error("execute SQL")
         raise
 
+    def resolve_schema(self, relation: BaseRelation) -> t.Optional[str]:
+        return relation.schema
+
+    def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
+        return relation.identifier
+
     @staticmethod
     def _raise_parsetime_adapter_call_error(action: str) -> None:
         raise ParsetimeAdapterCallError(f"Can't {action} at parse time.")
@@ -276,6 +290,20 @@ def execute(
             return AdapterResponse("Success"), pandas_to_agate(resp)
         return AdapterResponse("Success"), empty_table()
 
+    def resolve_schema(self, relation: BaseRelation) -> t.Optional[str]:
+        schema = self._map_table_name(relation.database, relation.schema, relation.identifier).db
+        if not schema:
+            return None
+        return schema
+
+    def resolve_identifier(self, relation: BaseRelation) -> t.Optional[str]:
+        identifier = self._map_table_name(
+            relation.database, relation.schema, relation.identifier
+        ).name
+        if not identifier:
+            return None
+        return identifier
+
     def _map_table_name(
         self, database: t.Optional[str], schema: t.Optional[str], identifier: t.Optional[str]
     ) -> exp.Table:

sqlmesh/dbt/basemodel.py

@@ -71,7 +71,6 @@ class BaseModelConfig(GeneralConfig):
         database: Database the model is stored in
         schema: Custom schema name added to the model schema name
         alias: Relation identifier for this model instead of the filename
-        sql_header: SQL statement to run before table/view creation. Currently implemented as a pre-hook.
         pre-hook: List of SQL statements to run before the model is built.
         post-hook: List of SQL statements to run after the model is built.
         full_refresh: Forces the model to always do a full refresh or never do a full refresh
@@ -94,7 +93,6 @@ class BaseModelConfig(GeneralConfig):
     schema_: str = Field("", alias="schema")
     database: t.Optional[str] = None
     alias: t.Optional[str] = None
-    sql_header: t.Optional[str] = None
     pre_hook: t.List[Hook] = Field([], alias="pre-hook")
     post_hook: t.List[Hook] = Field([], alias="post-hook")
     full_refresh: t.Optional[bool] = None
@@ -229,10 +227,6 @@ def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
             if field_val:
                 optional_kwargs[field] = field_val
 
-        pre_hooks = self.pre_hook
-        if self.sql_header:
-            pre_hooks.insert(0, Hook(sql=self.sql_header))
-
         return {
             "audits": [(test.name, {}) for test in self.tests],
             "columns": column_types_to_sqlmesh(self.columns) or None,
@@ -243,7 +237,7 @@ def sqlmesh_model_kwargs(self, context: DbtContext) -> t.Dict[str, t.Any]:
             "jinja_macros": jinja_macros,
             "path": self.path,
             "hash_raw_query": True,
-            "pre_statements": [d.jinja_statement(hook.sql) for hook in pre_hooks],
+            "pre_statements": [d.jinja_statement(hook.sql) for hook in self.pre_hook],
             "post_statements": [d.jinja_statement(hook.sql) for hook in self.post_hook],
             "tags": self.tags,
             **optional_kwargs,

sqlmesh/dbt/model.py

@@ -58,7 +58,7 @@ class ModelConfig(BaseModelConfig):
         cluster_by: Field(s) to use for clustering in data warehouses that support clustering
         incremental_strategy: Strategy used to build the incremental model
         materialized: How the model will be materialized in the database
-        sql_header: SQL statement to inject above create table/view as
+        sql_header: SQL statement to run before table/view creation. Currently implemented as a pre-hook.
         unique_key: List of columns that define row uniqueness for the model
         partition_by: List of partition columns or dictionary of bigquery partition by parameters ([dbt bigquery config](https://docs.getdbt.com/reference/resource-configs/bigquery-configs)).
     """
@@ -106,6 +106,9 @@ def _validate_partition_by(cls, v: t.Any) -> t.Union[t.List[str], t.Dict[str, t.
         if isinstance(v, dict):
             if not v.get("field"):
                 raise ConfigError("'field' key required for partition_by.")
+            if "granularity" in v and v["granularity"] not in GRANULARITY_TO_PARTITION_FORMAT:
+                granularity = v["granularity"]
+                raise ConfigError(f"Unexpected granularity '{granularity}' in partition_by '{v}'.")
             return {"data_type": "date", "granularity": "day", **v}
         raise ConfigError(f"Invalid format for partition_by '{v}'")
 
@@ -268,12 +271,59 @@ def to_sqlmesh(self, context: DbtContext) -> Model:
         if not context.target:
             raise ConfigError(f"Target required to load '{self.sql_name}' into SQLMesh.")
 
+        model_kwargs = self.sqlmesh_model_kwargs(context)
+        if self.sql_header:
+            model_kwargs["pre_statements"].insert(0, d.jinja_statement(self.sql_header))
+
+        if context.target.type == "bigquery":
+            dbt_max_partition_blob = self._dbt_max_partition_blob()
+            if dbt_max_partition_blob:
+                model_kwargs["pre_statements"].append(d.jinja_statement(dbt_max_partition_blob))
+
         return create_sql_model(
             self.sql_name,
             query,
             dialect=dialect,
             kind=self.model_kind(context.target),
             start=self.start,
             **optional_kwargs,
-            **self.sqlmesh_model_kwargs(context),
+            **model_kwargs,
         )
+
+    def _dbt_max_partition_blob(self) -> t.Optional[str]:
+        """Returns a SQL blob which declares the _dbt_max_partition variable. Only applicable to BigQuery."""
+        if (
+            not isinstance(self.partition_by, dict)
+            or self.model_materialization != Materialization.INCREMENTAL
+        ):
+            return None
+
+        data_type = self.partition_by["data_type"]
+        granularity = self.partition_by["granularity"]
+
+        parse_fun = f"parse_{data_type}" if data_type in ("date", "datetime", "timestamp") else None
+        if parse_fun:
+            parse_format = GRANULARITY_TO_PARTITION_FORMAT[granularity]
+            partition_exp = f"{parse_fun}('{parse_format}', partition_id)"
+        else:
+            partition_exp = "CAST(partition_id AS INT64)"
+
+        return f"""
+{{% if is_incremental() %}}
+  DECLARE _dbt_max_partition {data_type.upper()} DEFAULT (
+    SELECT MAX({partition_exp})
+    FROM `{{{{ target.database }}}}`.`{{{{ adapter.resolve_schema(this) }}}}`.INFORMATION_SCHEMA.PARTITIONS
+    WHERE table_name = '{{{{ adapter.resolve_identifier(this) }}}}'
+      AND partition_id IS NOT NULL
+      AND partition_id != '__NULL__'
+  );
+{{% endif %}}
+"""
+
+
+GRANULARITY_TO_PARTITION_FORMAT = {
+    "day": "%Y%m%d",
+    "month": "%Y%m",
+    "year": "%Y",
+    "hour": "%Y%m%d%H",
+}

tests/dbt/test_adapter.py

@@ -3,6 +3,7 @@
 import typing as t
 
 import pytest
+from dbt.adapters.base import BaseRelation
 from dbt.adapters.base.column import Column
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp
@@ -80,7 +81,11 @@ def test_adapter_map_snapshot_tables(
     assert context.target
     engine_adapter = context.target.to_sqlmesh().create_engine_adapter()
     renderer = runtime_renderer(
-        context, engine_adapter=engine_adapter, snapshots={"test_db.test_model": snapshot_mock}
+        context,
+        engine_adapter=engine_adapter,
+        snapshots={"test_db.test_model": snapshot_mock},
+        test_model=BaseRelation.create(schema="test_db", identifier="test_model"),
+        foo_bar=BaseRelation.create(schema="foo", identifier="bar"),
     )
 
     engine_adapter.create_schema("foo")
@@ -106,3 +111,9 @@ def test_adapter_map_snapshot_tables(
         renderer("{{ adapter.get_relation(database=none, schema='foo', identifier='bar') }}")
         == '"foo"."bar"'
     )
+
+    assert renderer("{{ adapter.resolve_schema(test_model) }}") == "sqlmesh"
+    assert renderer("{{ adapter.resolve_identifier(test_model) }}") == "test_db__test_model"
+
+    assert renderer("{{ adapter.resolve_schema(foo_bar) }}") == "foo"
+    assert renderer("{{ adapter.resolve_identifier(foo_bar) }}") == "bar"

tests/dbt/test_transformation.py

@@ -6,6 +6,7 @@
 from dbt.adapters.base import BaseRelation
 from dbt.contracts.relation import Policy
 from dbt.exceptions import CompilationError
+from pytest_mock.plugin import MockerFixture
 from sqlglot import exp, parse_one
 
 from sqlmesh.core.context import Context
@@ -570,3 +571,39 @@ def test_is_incremental(sushi_test_project: Project, assert_exp_eq):
         model_config.to_sqlmesh(context).render_query_or_raise(has_intervals=True).sql(),
         'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a" WHERE "ds" > (SELECT MAX("ds") FROM "model" AS "model")',
     )
+
+
+def test_dbt_max_partition(sushi_test_project: Project, assert_exp_eq, mocker: MockerFixture):
+    model_config = ModelConfig(
+        name="model",
+        package_name="package",
+        schema="sushi",
+        partition_by={"field": "`ds`", "data_type": "datetime", "granularity": "month"},
+        materialized=Materialization.INCREMENTAL,
+        sql="""
+        SELECT 1 AS one FROM tbl_a
+        {% if is_incremental() %}
+        WHERE ds > _dbt_max_partition
+        {% endif %}
+        """,
+    )
+    context = sushi_test_project.context
+    context.target = BigQueryConfig(
+        name="test_target", schema="test_schema", database="test-project"
+    )
+
+    assert (
+        model_config.to_sqlmesh(context).pre_statements[-1].sql().strip()  # type: ignore
+        == """
+JINJA_STATEMENT_BEGIN;
+{% if is_incremental() %}
+  DECLARE _dbt_max_partition DATETIME DEFAULT (
+    SELECT MAX(parse_datetime('%Y%m', partition_id))
+    FROM `{{ target.database }}`.`{{ adapter.resolve_schema(this) }}`.INFORMATION_SCHEMA.PARTITIONS
+    WHERE table_name = '{{ adapter.resolve_identifier(this) }}'
+      AND partition_id IS NOT NULL
+      AND partition_id != '__NULL__'
+  );
+{% endif %}
+JINJA_END;""".strip()
+    )