Repository: TobikoData/sqlmesh
Pull request #2878 (2024-07-08 16:52:47)

Docs: Clarify the scheduler type value for each support scheduler configuration


Patch:
diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md
index 2da2cfad..e8c1d0ac 100644
--- a/docs/reference/configuration.md
+++ b/docs/reference/configuration.md
@@ -193,10 +193,14 @@ Below is the list of configuration options specific to each corresponding schedu
 
 #### Builtin
 
+**Type:** `builtin`
+
 No configuration options are supported by this scheduler type.
 
 #### Airflow
 
+**Type:** `airflow`
+
 See [Airflow Integration Guide](../integrations/airflow.md) for information about how to integrate Airflow with SQLMesh.
 
 | Option                            | Description                                                                                                                                                                                                                                                                                                     |  Type   | Required |
@@ -216,6 +220,8 @@ See [Airflow Integration Guide](../integrations/airflow.md) for information abou
 
 #### Cloud Composer
 
+**Type:** `cloud_composer`
+
 The Google Cloud Composer scheduler type shares the same configuration options as the `airflow` type, except for `username` and `password`. Cloud Composer relies on `gcloud` authentication, so the `username` and `password` options are not required.
 
 ## Gateway/connection defaults
diff --git a/examples/sushi_dbt/models/schema.yml b/examples/sushi_dbt/models/schema.yml
index 3bb0c047..eb5d288d 100644
--- a/examples/sushi_dbt/models/schema.yml
+++ b/examples/sushi_dbt/models/schema.yml
@@ -77,3 +77,6 @@ models:
               amount: 10000
       - name: ds
         description: Date
+    versions:
+      - v: 1
+      - v: 2
diff --git a/examples/sushi_dbt/models/top_waiters.sql b/examples/sushi_dbt/models/top_waiters.sql
index b08eca97..f839b31d 100644
--- a/examples/sushi_dbt/models/top_waiters.sql
+++ b/examples/sushi_dbt/models/top_waiters.sql
@@ -7,12 +7,12 @@
 SELECT
   waiter_id::INT AS waiter_id,
   revenue::DOUBLE AS revenue
-FROM {{ ref('waiter_revenue_by_day') }}
+FROM {{ ref('waiter_revenue_by_day', version=1) }}
 WHERE
   ds = (
     SELECT
       MAX(ds)
-    FROM {{ ref('waiter_revenue_by_day') }}
+    FROM {{ ref('waiter_revenue_by_day', v=1) }}
   )
 ORDER BY
   revenue DESC
diff --git a/examples/sushi_dbt/models/waiter_revenue_by_day_v1.sql b/examples/sushi_dbt/models/waiter_revenue_by_day_v1.sql
new file mode 100644
index 00000000..d430c612
--- /dev/null
+++ b/examples/sushi_dbt/models/waiter_revenue_by_day_v1.sql
@@ -0,0 +1,29 @@
+{{
+  config(
+    materialized='incremental',
+    incremental_strategy='delete+insert',
+    cluster_by=['ds'],
+    time_column='ds',
+  )
+}}
+
+SELECT
+  o.waiter_id::INT AS waiter_id, /* Waiter id */
+  SUM(oi.quantity * i.price)::DOUBLE AS revenue, /* Revenue from orders taken by this waiter */
+  o.ds::TEXT AS ds /* Date */
+FROM {{ ref('orders') }} AS o
+LEFT JOIN {{ ref('order_items') }} AS oi
+  ON o.id = oi.order_id AND o.ds = oi.ds
+LEFT JOIN {{ ref('items') }} AS i
+  ON oi.item_id = i.id AND oi.ds = i.ds
+{% if is_incremental() %}
+  WHERE
+    o.ds > (select max(ds) from {{ this }})
+{% endif %}
+{% if sqlmesh_incremental is defined %}
+  WHERE
+      o.ds BETWEEN '{{ start_ds }}' AND '{{ end_ds }}'
+{% endif %}
+GROUP BY
+  o.waiter_id,
+  o.ds
diff --git a/sqlmesh/core/table_diff.py b/sqlmesh/core/table_diff.py
index e68fa3ce..44d9afd7 100644
--- a/sqlmesh/core/table_diff.py
+++ b/sqlmesh/core/table_diff.py
@@ -221,6 +221,8 @@ class TableDiff:
                 elif col.table == "t":
                     t_index.append(col)
             index_cols = list(dict.fromkeys(index_cols))
+            s_index = list(dict.fromkeys(s_index))
+            t_index = list(dict.fromkeys(t_index))
 
             matched_columns = {
                 c: t for c, t in self.source_schema.items() if t == self.target_schema.get(c)
diff --git a/sqlmesh/dbt/basemodel.py b/sqlmesh/dbt/basemodel.py
index 0e670531..40a07ea9 100644
--- a/sqlmesh/dbt/basemodel.py
+++ b/sqlmesh/dbt/basemodel.py
@@ -120,6 +120,9 @@ class BaseModelConfig(GeneralConfig):
     columns: t.Dict[str, ColumnConfig] = {}
     quoting: t.Dict[str, t.Optional[bool]] = {}
 
+    version: t.Optional[int] = None
+    latest_version: t.Optional[int] = None
+
     _canonical_name: t.Optional[str] = None
 
     @field_validator("pre_hook", "post_hook", mode="before")
@@ -184,9 +187,9 @@ class BaseModelConfig(GeneralConfig):
     @property
     def config_name(self) -> str:
         """
-        Get the model's config name (package_name.name)
+        Get the model's config name (package_name.alias)
         """
-        return f"{self.package_name}.{self.name}"
+        return f"{self.package_name}.{self.alias}"
 
     def dialect(self, context: DbtContext) -> str:
         return self.dialect_ or context.default_dialect
diff --git a/sqlmesh/dbt/builtin.py b/sqlmesh/dbt/builtin.py
index d6d62db4..00cd897a 100644
--- a/sqlmesh/dbt/builtin.py
+++ b/sqlmesh/dbt/builtin.py
@@ -171,11 +171,33 @@ def log(msg: str, info: bool = False) -> str:
 
 
 def generate_ref(refs: t.Dict[str, t.Any], api: Api) -> t.Callable:
-    def ref(package: str, name: t.Optional[str] = None) -> t.Optional[BaseRelation]:
+    def ref(
+        package: str, name: t.Optional[str] = None, **kwargs: t.Any
+    ) -> t.Optional[BaseRelation]:
+        version = kwargs.get("version", kwargs.get("v"))
         ref_name = f"{package}.{name}" if name else package
-        relation_info = refs.get(ref_name)
+
+        if version is not None:
+            relation_info = refs.get(f"{ref_name}_v{version}")
+            if relation_info is None:
+                logger.warning(
+                    "Could not resolve ref '%s' with version '%s'. Falling back to unversioned reference",
+                    ref_name,
+                    version,
+                )
+                relation_info = refs.get(ref_name)
+        else:
+            relation_info = refs.get(ref_name)
+            if not relation_info:
+                versioned_infos = sorted(
+                    [(r, info) for r, info in refs.items() if r.startswith(f"{ref_name}_v")],
+                    key=lambda i: i[0],
+                )
+                if versioned_infos:
+                    relation_info = versioned_infos[-1][1]
+
         if relation_info is None:
-            logger.debug("Could not resolve ref '%s'", ref_name)
+            logger.debug("Could not resolve ref '%s', version '%s'", ref_name, version)
             return None
 
         return _relation_info_to_relation(relation_info, api.Relation, api.quote_policy)
diff --git a/sqlmesh/dbt/context.py b/sqlmesh/dbt/context.py
index d6a41286..14e14efa 100644
--- a/sqlmesh/dbt/context.py
+++ b/sqlmesh/dbt/context.py
@@ -185,8 +185,11 @@ class DbtContext:
             for model in t.cast(
                 t.Dict[str, t.Union[ModelConfig, SeedConfig]], {**self._seeds, **self._models}
             ).values():
-                self._refs[model.name] = model
+                self._refs[model.alias or model.name] = model
                 self._refs[model.config_name] = model
+                if model.version is not None and model.version == model.latest_version:
+                    self._refs[model.name] = model
+                    self._refs[f"{model.package_name}.{model.name}"] = model
         return self._refs
 
     @property
@@ -270,7 +273,6 @@ class DbtContext:
         dependency_context.seeds = seeds
         dependency_context.models = models
         dependency_context.variables = variables
-        dependency_context._refs = {**dependency_context._seeds, **dependency_context._models}  # type: ignore
 
         return dependency_context
 
diff --git a/sqlmesh/dbt/manifest.py b/sqlmesh/dbt/manifest.py
index 0bac148a..f62dbc38 100644
--- a/sqlmesh/dbt/manifest.py
+++ b/sqlmesh/dbt/manifest.py
@@ -223,14 +223,15 @@ class ManifestHelper:
                     self._flatten_dependencies_from_macros(dependencies.macros, node.package_name)
                 )
 
-                self._models_per_package[node.package_name][node.name] = ModelConfig(
+                # Using the alias instead of the name because the alias captures the version of the model.
+                self._models_per_package[node.package_name][node.alias] = ModelConfig(
                     sql=sql,
                     dependencies=dependencies,
                     tests=tests,
                     **node_config,
                 )
             else:
-                self._seeds_per_package[node.package_name][node.name] = SeedConfig(
+                self._seeds_per_package[node.package_name][node.alias] = SeedConfig(
                     dependencies=Dependencies(macros=macro_references),
                     tests=tests,
                     **node_config,
@@ -451,7 +452,13 @@ def _macro_references(
 
 def _refs(node: ManifestNode) -> t.Set[str]:
     if DBT_VERSION >= (1, 5):
-        return {f"{r.package}.{r.name}" if r.package else r.name for r in node.refs}  # type: ignore
+        result = set()
+        for r in node.refs:
+            ref_name = f"{r.package}.{r.name}" if r.package else r.name
+            if getattr(r, "version", None):
+                ref_name = f"{ref_name}_v{r.version}"
+            result.add(ref_name)
+        return result
     else:
         return {".".join(r) for r in node.refs}  # type: ignore
 


Test patch:
diff --git a/tests/core/test_table_diff.py b/tests/core/test_table_diff.py
index 598609a9..99b933ed 100644
--- a/tests/core/test_table_diff.py
+++ b/tests/core/test_table_diff.py
@@ -1,8 +1,9 @@
 import pytest
+from pytest_mock.plugin import MockerFixture
 import pandas as pd
 from sqlglot import exp
 from sqlmesh.core import dialect as d
-
+from sqlmesh.core.context import Context
 from sqlmesh.core.config import AutoCategorizationMode, CategorizerConfig
 from sqlmesh.core.model import SqlModel, load_sql_based_model
 
@@ -217,3 +218,45 @@ def test_grain_check(sushi_context_fixed_date):
     assert row_diff.stats["distinct_count_t"] == 10
     assert row_diff.s_sample.shape == (0, 3)
     assert row_diff.t_sample.shape == (3, 3)
+
+
+@pytest.mark.slow
+def test_generated_sql(sushi_context_fixed_date: Context, mocker: MockerFixture):
+    engine_adapter = sushi_context_fixed_date.engine_adapter
+
+    engine_adapter.ctas(
+        "table_diff_source",
+        pd.DataFrame(
+            {
+                "key": [1, 2, 3],
+                "value": [1.0, 4.2, 4.1],
+            }
+        ),
+    )
+
+    engine_adapter.ctas(
+        "table_diff_target",
+        pd.DataFrame(
+            {
+                "key": [1, 2, 6],
+                "value": [1.0, 3.0, -2.2],
+            }
+        ),
+    )
+
+    query_sql = 'CREATE TABLE IF NOT EXISTS "sqlmesh_temp"."__temp_diff_abcdefgh" AS SELECT *, CASE WHEN "key_matches" = 1 AND "value_matches" = 1 THEN 1 ELSE 0 END AS "row_full_match" FROM (SELECT "s"."key" AS "s__key", "s"."value" AS "s__value", "t"."key" AS "t__key", "t"."value" AS "t__value", CASE WHEN NOT "s"."key" IS NULL THEN 1 ELSE 0 END AS "s_exists", CASE WHEN NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "t_exists", CASE WHEN "s"."key" = "t"."key" AND NOT "s"."key" IS NULL AND NOT "t"."key" IS NULL THEN 1 ELSE 0 END AS "row_joined", CASE WHEN "s"."key" IS NULL AND "t"."key" IS NULL THEN 1 ELSE 0 END AS "null_grain", CASE WHEN "s"."key" = "t"."key" THEN 1 WHEN ("s"."key" IS NULL) AND ("t"."key" IS NULL) THEN 1 WHEN ("s"."key" IS NULL) OR ("t"."key" IS NULL) THEN 0 ELSE 0 END AS "key_matches", CASE WHEN ROUND("s"."value", 3) = ROUND("t"."value", 3) THEN 1 WHEN ("s"."value" IS NULL) AND ("t"."value" IS NULL) THEN 1 WHEN ("s"."value" IS NULL) OR ("t"."value" IS NULL) THEN 0 ELSE 0 END AS "value_matches" FROM "table_diff_source" AS "s" FULL JOIN "table_diff_target" AS "t" ON ("s"."key" = "t"."key") OR (("s"."key" IS NULL) AND ("t"."key" IS NULL))) AS "stats"'
+    summary_query_sql = 'SELECT SUM("s_exists") AS "s_count", SUM("t_exists") AS "t_count", SUM("row_joined") AS "join_count", SUM("null_grain") AS "null_grain_count", SUM("row_full_match") AS "full_match_count", SUM("key_matches") AS "key_matches", SUM("value_matches") AS "value_matches", COUNT(DISTINCT ("s__key")) AS "distinct_count_s", COUNT(DISTINCT ("t__key")) AS "distinct_count_t" FROM "sqlmesh_temp"."__temp_diff_abcdefgh"'
+    sample_query_sql = 'SELECT "s_exists", "t_exists", "row_joined", "row_full_match", "s__key", "s__value", "t__key", "t__value" FROM "sqlmesh_temp"."__temp_diff_abcdefgh" WHERE "key_matches" = 0 OR "value_matches" = 0 ORDER BY "s__key" NULLS FIRST, "t__key" NULLS FIRST LIMIT 20'
+
+    spy_execute = mocker.spy(engine_adapter, "_execute")
+    mocker.patch("sqlmesh.core.engine_adapter.base.random_id", return_value="abcdefgh")
+
+    sushi_context_fixed_date.table_diff(
+        source="table_diff_source",
+        target="table_diff_target",
+        on=["key"],
+    )
+
+    spy_execute.assert_any_call(query_sql)
+    spy_execute.assert_any_call(summary_query_sql)
+    spy_execute.assert_any_call(sample_query_sql)
diff --git a/tests/dbt/test_manifest.py b/tests/dbt/test_manifest.py
index bb81d193..59b4d0ff 100644
--- a/tests/dbt/test_manifest.py
+++ b/tests/dbt/test_manifest.py
@@ -25,38 +25,45 @@ def test_manifest_helper(caplog):
         variable_overrides={"start": "2020-01-01"},
     )
 
-    assert helper.models()["top_waiters"].dependencies == Dependencies(
+    models = helper.models()
+
+    assert models["top_waiters"].dependencies == Dependencies(
         refs={"sushi.waiter_revenue_by_day", "waiter_revenue_by_day"},
         variables={"top_waiters:revenue", "top_waiters:limit"},
         macros=[MacroReference(name="ref"), MacroReference(name="var")],
     )
-    assert helper.models()["top_waiters"].materialized == "view"
-    assert helper.models()["top_waiters"].dialect_ == "postgres"
+    assert models["top_waiters"].materialized == "view"
+    assert models["top_waiters"].dialect_ == "postgres"
 
-    assert helper.models()["waiters"].dependencies == Dependencies(
+    assert models["waiters"].dependencies == Dependencies(
         macros={MacroReference(name="incremental_by_time"), MacroReference(name="source")},
         sources={"streaming.orders"},
     )
-    assert helper.models()["waiters"].materialized == "ephemeral"
-    assert helper.models()["items_snapshot"].materialized == "snapshot"
-    assert helper.models()["items_snapshot"].updated_at == "ds"
-    assert helper.models()["items_snapshot"].unique_key == ["id"]
-    assert helper.models()["items_snapshot"].strategy == "timestamp"
-    assert helper.models()["items_snapshot"].table_schema == "snapshots"
-    assert helper.models()["items_snapshot"].invalidate_hard_deletes is True
-    assert helper.models()["items_check_snapshot"].materialized == "snapshot"
-    assert helper.models()["items_check_snapshot"].check_cols == ["ds"]
-    assert helper.models()["items_check_snapshot"].unique_key == ["id"]
-    assert helper.models()["items_check_snapshot"].strategy == "check"
-    assert helper.models()["items_check_snapshot"].table_schema == "snapshots"
-    assert helper.models()["items_check_snapshot"].invalidate_hard_deletes is True
-    assert helper.models()["items_no_hard_delete_snapshot"].materialized == "snapshot"
-    assert helper.models()["items_no_hard_delete_snapshot"].unique_key == ["id"]
-    assert helper.models()["items_no_hard_delete_snapshot"].strategy == "timestamp"
-    assert helper.models()["items_no_hard_delete_snapshot"].table_schema == "snapshots"
-    assert helper.models()["items_no_hard_delete_snapshot"].invalidate_hard_deletes is False
-
-    waiter_as_customer_by_day_config = helper.models()["waiter_as_customer_by_day"]
+    assert models["waiters"].materialized == "ephemeral"
+    assert models["items_snapshot"].materialized == "snapshot"
+    assert models["items_snapshot"].updated_at == "ds"
+    assert models["items_snapshot"].unique_key == ["id"]
+    assert models["items_snapshot"].strategy == "timestamp"
+    assert models["items_snapshot"].table_schema == "snapshots"
+    assert models["items_snapshot"].invalidate_hard_deletes is True
+    assert models["items_check_snapshot"].materialized == "snapshot"
+    assert models["items_check_snapshot"].check_cols == ["ds"]
+    assert models["items_check_snapshot"].unique_key == ["id"]
+    assert models["items_check_snapshot"].strategy == "check"
+    assert models["items_check_snapshot"].table_schema == "snapshots"
+    assert models["items_check_snapshot"].invalidate_hard_deletes is True
+    assert models["items_no_hard_delete_snapshot"].materialized == "snapshot"
+    assert models["items_no_hard_delete_snapshot"].unique_key == ["id"]
+    assert models["items_no_hard_delete_snapshot"].strategy == "timestamp"
+    assert models["items_no_hard_delete_snapshot"].table_schema == "snapshots"
+    assert models["items_no_hard_delete_snapshot"].invalidate_hard_deletes is False
+
+    # Test versioned models
+    assert models["waiter_revenue_by_day_v1"].version == 1
+    assert models["waiter_revenue_by_day_v2"].version == 2
+    assert "waiter_revenue_by_day" not in models
+
+    waiter_as_customer_by_day_config = models["waiter_as_customer_by_day"]
     assert waiter_as_customer_by_day_config.dependencies == Dependencies(
         refs={"waiters", "waiter_names", "customers"},
         macros=[MacroReference(name="ref")],
@@ -66,7 +73,7 @@ def test_manifest_helper(caplog):
     assert waiter_as_customer_by_day_config.cluster_by == ["ds"]
     assert waiter_as_customer_by_day_config.time_column == "ds"
 
-    waiter_revenue_by_day_config = helper.models()["waiter_revenue_by_day"]
+    waiter_revenue_by_day_config = models["waiter_revenue_by_day_v2"]
     assert waiter_revenue_by_day_config.dependencies == Dependencies(
         macros={
             MacroReference(name="log_value"),
@@ -96,12 +103,14 @@ def test_manifest_helper(caplog):
 
     assert helper.seeds()["waiter_names"].path == Path("seeds/waiter_names.csv")
 
-    assert helper.sources()["streaming.items"].table_name == "items"
-    assert helper.sources()["streaming.items"].schema_ == "raw"
-    assert helper.sources()["streaming.orders"].table_name == "orders"
-    assert helper.sources()["streaming.orders"].schema_ == "raw"
-    assert helper.sources()["streaming.order_items"].table_name == "order_items"
-    assert helper.sources()["streaming.order_items"].schema_ == "raw"
+    sources = helper.sources()
+
+    assert sources["streaming.items"].table_name == "items"
+    assert sources["streaming.items"].schema_ == "raw"
+    assert sources["streaming.orders"].table_name == "orders"
+    assert sources["streaming.orders"].schema_ == "raw"
+    assert sources["streaming.order_items"].table_name == "order_items"
+    assert sources["streaming.order_items"].schema_ == "raw"
 
 
 @pytest.mark.xdist_group("dbt_manifest")
diff --git a/tests/fixtures/dbt/sushi_test/models/schema.yml b/tests/fixtures/dbt/sushi_test/models/schema.yml
index d6119c53..78d6c3e9 100644
--- a/tests/fixtures/dbt/sushi_test/models/schema.yml
+++ b/tests/fixtures/dbt/sushi_test/models/schema.yml
@@ -7,6 +7,9 @@ models:
   - name: waiters
   - name: waiter_as_customer_by_day
   - name: waiter_revenue_by_day
+    versions:
+      - v: 1
+      - v: 2
   - name: disabled_model
     columns:
       - name: one
diff --git a/tests/fixtures/dbt/sushi_test/models/waiter_revenue_by_day_v1.sql b/tests/fixtures/dbt/sushi_test/models/waiter_revenue_by_day_v1.sql
new file mode 100644
index 00000000..335e7ab7
--- /dev/null
+++ b/tests/fixtures/dbt/sushi_test/models/waiter_revenue_by_day_v1.sql
@@ -0,0 +1,40 @@
+{{
+    config(
+        materialized='incremental',
+        incremental_strategy='delete+insert',
+        cluster_by=['ds'],
+        time_column='ds',
+        dialect="bigquery"
+    )
+}}
+
+{{ log_value(5) }}
+{{ log(adapter.dispatch('current_engine', 'customers')()) }}
+
+{{ test_dependencies() }}
+
+{% set results = run_query('select 1 as constant') %}
+
+SELECT
+  o.waiter_id::INT AS waiter_id, /* Waiter id */
+  SUM(oi.quantity * i.price)::DOUBLE AS revenue, /* Revenue from orders taken by this waiter */
+  o.ds::TEXT AS ds, /* Date */
+  {% if execute %}
+    {{ results.columns[0].values()[0] }}::INT AS constant /* Constant */
+  {% endif %}
+FROM {{ source('streaming', 'orders') }} AS o
+LEFT JOIN {{ source('streaming', 'order_items') }} AS oi
+  ON o.id = oi.order_id AND o.ds = oi.ds
+LEFT JOIN {{ source('streaming', 'items') }} AS i
+  ON oi.item_id = i.id AND oi.ds = i.ds
+{% if is_incremental() %}
+  WHERE
+    o.ds > (select max(ds) from {{ this }})
+{% endif %}
+{% if sqlmesh_incremental is defined %}
+  WHERE
+      o.ds BETWEEN '{{ start_ds }}' AND '{{ end_ds }}'
+{% endif %}
+GROUP BY
+  o.waiter_id,
+  o.ds