Fix: Only apply partition interval unit if partitioned_by is not set by the user explicitly (#3636)

Author: izeigerman

sqlmesh/core/model/definition.py

@@ -23,6 +23,7 @@
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
 from sqlmesh.core.audit import Audit, ModelAudit
+from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.macros import MacroRegistry, macro
 from sqlmesh.core.model.common import (
     expression_validator,
@@ -1044,9 +1045,7 @@ def full_depends_on(self) -> t.Set[str]:
     @property
     def partitioned_by(self) -> t.List[exp.Expression]:
         """Columns to partition the model by, including the time column if it is not already included."""
-        if self.time_column and self.time_column.column not in {
-            col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)
-        }:
+        if self.time_column and not self._is_time_column_in_partitioned_by:
             return [
                 TIME_COL_PARTITION_FUNC.get(self.dialect, lambda x, y: x)(
                     self.time_column.column, self.columns_to_types
@@ -1055,6 +1054,16 @@ def partitioned_by(self) -> t.List[exp.Expression]:
             ]
         return self.partitioned_by_
 
+    @property
+    def partition_interval_unit(self) -> t.Optional[IntervalUnit]:
+        """The interval unit to use for partitioning if applicable."""
+        # Only return the interval unit for partitioning if the partitioning
+        # wasn't explicitly set by the user. Otherwise, the user-provided
+        # value should always take precedence.
+        if self.time_column and not self._is_time_column_in_partitioned_by:
+            return self.interval_unit
+        return None
+
     @property
     def audits_with_args(self) -> t.List[t.Tuple[Audit, t.Dict[str, exp.Expression]]]:
         from sqlmesh.core.audit.builtin import BUILT_IN_AUDITS
@@ -1071,6 +1080,12 @@ def audits_with_args(self) -> t.List[t.Tuple[Audit, t.Dict[str, exp.Expression]]
 
         return list(audits_with_args.values())
 
+    @property
+    def _is_time_column_in_partitioned_by(self) -> bool:
+        return self.time_column is not None and self.time_column.column in {
+            col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)
+        }
+
 
 class SqlModel(_Model):
     """The model definition which relies on a SQL query to fetch the data.

sqlmesh/core/snapshot/evaluator.py

@@ -1156,7 +1156,7 @@ def _replace_query_for_model(self, model: Model, name: str, query_or_df: QueryOr
             table_format=model.table_format,
             storage_format=model.storage_format,
             partitioned_by=model.partitioned_by,
-            partition_interval_unit=model.interval_unit,
+            partition_interval_unit=model.partition_interval_unit,
             clustered_by=model.clustered_by,
             table_properties=model.physical_properties,
             table_description=model.description,
@@ -1286,7 +1286,7 @@ def create(
                 table_format=model.table_format,
                 storage_format=model.storage_format,
                 partitioned_by=model.partitioned_by,
-                partition_interval_unit=model.interval_unit,
+                partition_interval_unit=model.partition_interval_unit,
                 clustered_by=model.clustered_by,
                 table_properties=model.physical_properties,
                 table_description=model.description if is_table_deployable else None,
@@ -1310,7 +1310,7 @@ def create(
                 table_format=model.table_format,
                 storage_format=model.storage_format,
                 partitioned_by=model.partitioned_by,
-                partition_interval_unit=model.interval_unit,
+                partition_interval_unit=model.partition_interval_unit,
                 clustered_by=model.clustered_by,
                 table_properties=model.physical_properties,
                 table_description=model.description if is_table_deployable else None,
@@ -1519,7 +1519,7 @@ def create(
                 table_format=model.table_format,
                 storage_format=model.storage_format,
                 partitioned_by=model.partitioned_by,
-                partition_interval_unit=model.interval_unit,
+                partition_interval_unit=model.partition_interval_unit,
                 clustered_by=model.clustered_by,
                 table_properties=model.physical_properties,
                 table_description=model.description if is_table_deployable else None,
@@ -1712,7 +1712,7 @@ def create(
             materialized_properties = {
                 "partitioned_by": model.partitioned_by,
                 "clustered_by": model.clustered_by,
-                "partition_interval_unit": model.interval_unit,
+                "partition_interval_unit": model.partition_interval_unit,
             }
         self.adapter.create_view(
             table_name,

tests/core/test_model.py

@@ -6677,3 +6677,37 @@ def dummy_model_entry(evaluator: MacroEvaluator) -> exp.Select:
     )
     assert isinstance(context._get_engine_adapter("duckdb"), DuckDBEngineAdapter)
     assert len(context._engine_adapters) == 2
+
+
+def test_partition_interval_unit():
+    expressions = d.parse(
+        """
+        MODEL (
+            name test,
+            kind INCREMENTAL_BY_TIME_RANGE(
+                time_column ds,
+            ),
+            cron '0 0 1 * *'
+        );
+        SELECT '2024-01-01' AS ds;
+        """
+    )
+    model = load_sql_based_model(expressions)
+    assert model.partition_interval_unit == IntervalUnit.MONTH
+
+    # Partitioning was explicitly set by the user
+    expressions = d.parse(
+        """
+        MODEL (
+            name test,
+            kind INCREMENTAL_BY_TIME_RANGE(
+                time_column ds,
+            ),
+            cron '0 0 1 * *',
+            partitioned_by (ds)
+        );
+        SELECT '2024-01-01' AS ds;
+        """
+    )
+    model = load_sql_based_model(expressions)
+    assert model.partition_interval_unit is None

tests/core/test_snapshot_evaluator.py

@@ -679,7 +679,7 @@ def test_evaluate_incremental_unmanaged_no_intervals(
         clustered_by=[],
         column_descriptions={},
         columns_to_types=table_columns,
-        partition_interval_unit=model.interval_unit,
+        partition_interval_unit=model.partition_interval_unit,
         partitioned_by=model.partitioned_by,
         table_format=None,
         storage_format=None,
@@ -891,7 +891,7 @@ def test_create_prod_table_exists_forward_only(mocker: MockerFixture, adapter_mo
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -965,7 +965,7 @@ def test_create_materialized_view(mocker: MockerFixture, adapter_mock, make_snap
         materialized=True,
         materialized_properties={
             "clustered_by": [],
-            "partition_interval_unit": IntervalUnit.DAY,
+            "partition_interval_unit": None,
             "partitioned_by": [],
         },
         view_properties={},
@@ -1014,7 +1014,7 @@ def test_create_view_with_properties(mocker: MockerFixture, adapter_mock, make_s
         },
         materialized_properties={
             "clustered_by": [],
-            "partition_interval_unit": IntervalUnit.DAY,
+            "partition_interval_unit": None,
             "partitioned_by": [],
         },
         table_description=None,
@@ -1667,7 +1667,7 @@ def test_create_scd_type_2_by_time(adapter_mock, make_snapshot):
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -1722,7 +1722,7 @@ def test_create_ctas_scd_type_2_by_time(adapter_mock, make_snapshot):
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -1843,7 +1843,7 @@ def test_create_scd_type_2_by_column(adapter_mock, make_snapshot):
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -1892,7 +1892,7 @@ def test_create_ctas_scd_type_2_by_column(adapter_mock, make_snapshot):
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -2166,7 +2166,7 @@ def test_create_incremental_by_unique_no_intervals(adapter_mock, make_snapshot):
         clustered_by=[],
         column_descriptions={},
         columns_to_types=table_columns,
-        partition_interval_unit=model.interval_unit,
+        partition_interval_unit=model.partition_interval_unit,
         partitioned_by=model.partitioned_by,
         table_format=None,
         storage_format=None,
@@ -2271,7 +2271,7 @@ def test_create_seed(mocker: MockerFixture, adapter_mock, make_snapshot):
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -2346,7 +2346,7 @@ def test_create_seed_on_error(mocker: MockerFixture, adapter_mock, make_snapshot
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -2402,7 +2402,7 @@ def test_create_seed_no_intervals(mocker: MockerFixture, adapter_mock, make_snap
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -2808,7 +2808,7 @@ def test_create_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
         table_format=model.table_format,
         storage_format=model.storage_format,
         partitioned_by=model.partitioned_by,
-        partition_interval_unit=model.interval_unit,
+        partition_interval_unit=model.partition_interval_unit,
         clustered_by=model.clustered_by,
         table_properties=model.physical_properties,
         table_description=None,
@@ -2894,7 +2894,7 @@ def test_evaluate_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
         table_format=model.table_format,
         storage_format=model.storage_format,
         partitioned_by=model.partitioned_by,
-        partition_interval_unit=model.interval_unit,
+        partition_interval_unit=model.partition_interval_unit,
         clustered_by=model.clustered_by,
         table_properties=model.physical_properties,
         table_description=model.description,
@@ -3061,7 +3061,7 @@ def test_create_snapshot(
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,
@@ -3108,7 +3108,7 @@ def test_migrate_snapshot(snapshot: Snapshot, mocker: MockerFixture, adapter_moc
         table_format=None,
         storage_format=None,
         partitioned_by=[],
-        partition_interval_unit=IntervalUnit.DAY,
+        partition_interval_unit=None,
         clustered_by=[],
         table_properties={},
         table_description=None,