Skip to content

Commit 5da32f2

Browse files
izeigermanizeigerman
authored
Fix: Only apply partition interval unit if partitioned_by is not set by the user explicitly (#3636)
1 parent 42396fb commit 5da32f2

File tree

4 files changed

+73
-24
lines changed

4 files changed

+73
-24
lines changed

‎sqlmesh/core/model/definition.py‎

Lines changed: 18 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
from sqlmesh.core import constants as c
2424
from sqlmesh.core import dialect as d
2525
from sqlmesh.core.audit import Audit, ModelAudit
26+
from sqlmesh.core.node import IntervalUnit
2627
from sqlmesh.core.macros import MacroRegistry, macro
2728
from sqlmesh.core.model.common import (
2829
expression_validator,
@@ -1113,9 +1114,7 @@ def full_depends_on(self) -> t.Set[str]:
11131114
@property
11141115
def partitioned_by(self) -> t.List[exp.Expression]:
11151116
"""Columns to partition the model by, including the time column if it is not already included."""
1116-
if self.time_column and self.time_column.column not in {
1117-
col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)
1118-
}:
1117+
if self.time_column and not self._is_time_column_in_partitioned_by:
11191118
return [
11201119
TIME_COL_PARTITION_FUNC.get(self.dialect, lambda x, y: x)(
11211120
self.time_column.column, self.columns_to_types
@@ -1124,6 +1123,16 @@ def partitioned_by(self) -> t.List[exp.Expression]:
11241123
]
11251124
return self.partitioned_by_
11261125

1126+
@property
1127+
def partition_interval_unit(self) -> t.Optional[IntervalUnit]:
1128+
"""The interval unit to use for partitioning if applicable."""
1129+
# Only return the interval unit for partitioning if the partitioning
1130+
# wasn't explicitly set by the user. Otherwise, the user-provided
1131+
# value should always take precedence.
1132+
if self.time_column and not self._is_time_column_in_partitioned_by:
1133+
return self.interval_unit
1134+
return None
1135+
11271136
@property
11281137
def audits_with_args(self) -> t.List[t.Tuple[Audit, t.Dict[str, exp.Expression]]]:
11291138
from sqlmesh.core.audit.builtin import BUILT_IN_AUDITS
@@ -1140,6 +1149,12 @@ def audits_with_args(self) -> t.List[t.Tuple[Audit, t.Dict[str, exp.Expression]]
11401149

11411150
return list(audits_with_args.values())
11421151

1152+
@property
1153+
def _is_time_column_in_partitioned_by(self) -> bool:
1154+
return self.time_column is not None and self.time_column.column in {
1155+
col for expr in self.partitioned_by_ for col in expr.find_all(exp.Column)
1156+
}
1157+
11431158

11441159
class SqlModel(_Model):
11451160
"""The model definition which relies on a SQL query to fetch the data.

‎sqlmesh/core/snapshot/evaluator.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1177,7 +1177,7 @@ def _replace_query_for_model(self, model: Model, name: str, query_or_df: QueryOr
11771177
table_format=model.table_format,
11781178
storage_format=model.storage_format,
11791179
partitioned_by=model.partitioned_by,
1180-
partition_interval_unit=model.interval_unit,
1180+
partition_interval_unit=model.partition_interval_unit,
11811181
clustered_by=model.clustered_by,
11821182
table_properties=model.physical_properties,
11831183
table_description=model.description,
@@ -1307,7 +1307,7 @@ def create(
13071307
table_format=model.table_format,
13081308
storage_format=model.storage_format,
13091309
partitioned_by=model.partitioned_by,
1310-
partition_interval_unit=model.interval_unit,
1310+
partition_interval_unit=model.partition_interval_unit,
13111311
clustered_by=model.clustered_by,
13121312
table_properties=model.physical_properties,
13131313
table_description=model.description if is_table_deployable else None,
@@ -1331,7 +1331,7 @@ def create(
13311331
table_format=model.table_format,
13321332
storage_format=model.storage_format,
13331333
partitioned_by=model.partitioned_by,
1334-
partition_interval_unit=model.interval_unit,
1334+
partition_interval_unit=model.partition_interval_unit,
13351335
clustered_by=model.clustered_by,
13361336
table_properties=model.physical_properties,
13371337
table_description=model.description if is_table_deployable else None,
@@ -1548,7 +1548,7 @@ def create(
15481548
table_format=model.table_format,
15491549
storage_format=model.storage_format,
15501550
partitioned_by=model.partitioned_by,
1551-
partition_interval_unit=model.interval_unit,
1551+
partition_interval_unit=model.partition_interval_unit,
15521552
clustered_by=model.clustered_by,
15531553
table_properties=model.physical_properties,
15541554
table_description=model.description if is_table_deployable else None,
@@ -1742,7 +1742,7 @@ def create(
17421742
materialized_properties = {
17431743
"partitioned_by": model.partitioned_by,
17441744
"clustered_by": model.clustered_by,
1745-
"partition_interval_unit": model.interval_unit,
1745+
"partition_interval_unit": model.partition_interval_unit,
17461746
}
17471747
self.adapter.create_view(
17481748
table_name,

‎tests/core/test_model.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7031,3 +7031,37 @@ def test_compile_time_checks(tmp_path: Path, assert_exp_eq):
70317031
model = create_seed_model("test_db.test_seed_model", model_kind, validate_query=False)
70327032
context.upsert_model(model)
70337033
context.plan(auto_apply=True, no_prompts=True)
7034+
7035+
7036+
def test_partition_interval_unit():
7037+
expressions = d.parse(
7038+
"""
7039+
MODEL (
7040+
name test,
7041+
kind INCREMENTAL_BY_TIME_RANGE(
7042+
time_column ds,
7043+
),
7044+
cron '0 0 1 * *'
7045+
);
7046+
SELECT '2024-01-01' AS ds;
7047+
"""
7048+
)
7049+
model = load_sql_based_model(expressions)
7050+
assert model.partition_interval_unit == IntervalUnit.MONTH
7051+
7052+
# Partitioning was explicitly set by the user
7053+
expressions = d.parse(
7054+
"""
7055+
MODEL (
7056+
name test,
7057+
kind INCREMENTAL_BY_TIME_RANGE(
7058+
time_column ds,
7059+
),
7060+
cron '0 0 1 * *',
7061+
partitioned_by (ds)
7062+
);
7063+
SELECT '2024-01-01' AS ds;
7064+
"""
7065+
)
7066+
model = load_sql_based_model(expressions)
7067+
assert model.partition_interval_unit is None

‎tests/core/test_snapshot_evaluator.py‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -680,7 +680,7 @@ def test_evaluate_incremental_unmanaged_no_intervals(
680680
clustered_by=[],
681681
column_descriptions={},
682682
columns_to_types=table_columns,
683-
partition_interval_unit=model.interval_unit,
683+
partition_interval_unit=model.partition_interval_unit,
684684
partitioned_by=model.partitioned_by,
685685
table_format=None,
686686
storage_format=None,
@@ -892,7 +892,7 @@ def test_create_prod_table_exists_forward_only(mocker: MockerFixture, adapter_mo
892892
table_format=None,
893893
storage_format=None,
894894
partitioned_by=[],
895-
partition_interval_unit=IntervalUnit.DAY,
895+
partition_interval_unit=None,
896896
clustered_by=[],
897897
table_properties={},
898898
table_description=None,
@@ -966,7 +966,7 @@ def test_create_materialized_view(mocker: MockerFixture, adapter_mock, make_snap
966966
materialized=True,
967967
materialized_properties={
968968
"clustered_by": [],
969-
"partition_interval_unit": IntervalUnit.DAY,
969+
"partition_interval_unit": None,
970970
"partitioned_by": [],
971971
},
972972
view_properties={},
@@ -1015,7 +1015,7 @@ def test_create_view_with_properties(mocker: MockerFixture, adapter_mock, make_s
10151015
},
10161016
materialized_properties={
10171017
"clustered_by": [],
1018-
"partition_interval_unit": IntervalUnit.DAY,
1018+
"partition_interval_unit": None,
10191019
"partitioned_by": [],
10201020
},
10211021
table_description=None,
@@ -1668,7 +1668,7 @@ def test_create_scd_type_2_by_time(adapter_mock, make_snapshot):
16681668
table_format=None,
16691669
storage_format=None,
16701670
partitioned_by=[],
1671-
partition_interval_unit=IntervalUnit.DAY,
1671+
partition_interval_unit=None,
16721672
clustered_by=[],
16731673
table_properties={},
16741674
table_description=None,
@@ -1723,7 +1723,7 @@ def test_create_ctas_scd_type_2_by_time(adapter_mock, make_snapshot):
17231723
table_format=None,
17241724
storage_format=None,
17251725
partitioned_by=[],
1726-
partition_interval_unit=IntervalUnit.DAY,
1726+
partition_interval_unit=None,
17271727
clustered_by=[],
17281728
table_properties={},
17291729
table_description=None,
@@ -1844,7 +1844,7 @@ def test_create_scd_type_2_by_column(adapter_mock, make_snapshot):
18441844
table_format=None,
18451845
storage_format=None,
18461846
partitioned_by=[],
1847-
partition_interval_unit=IntervalUnit.DAY,
1847+
partition_interval_unit=None,
18481848
clustered_by=[],
18491849
table_properties={},
18501850
table_description=None,
@@ -1893,7 +1893,7 @@ def test_create_ctas_scd_type_2_by_column(adapter_mock, make_snapshot):
18931893
table_format=None,
18941894
storage_format=None,
18951895
partitioned_by=[],
1896-
partition_interval_unit=IntervalUnit.DAY,
1896+
partition_interval_unit=None,
18971897
clustered_by=[],
18981898
table_properties={},
18991899
table_description=None,
@@ -2167,7 +2167,7 @@ def test_create_incremental_by_unique_no_intervals(adapter_mock, make_snapshot):
21672167
clustered_by=[],
21682168
column_descriptions={},
21692169
columns_to_types=table_columns,
2170-
partition_interval_unit=model.interval_unit,
2170+
partition_interval_unit=model.partition_interval_unit,
21712171
partitioned_by=model.partitioned_by,
21722172
table_format=None,
21732173
storage_format=None,
@@ -2297,7 +2297,7 @@ def test_create_seed(mocker: MockerFixture, adapter_mock, make_snapshot):
22972297
table_format=None,
22982298
storage_format=None,
22992299
partitioned_by=[],
2300-
partition_interval_unit=IntervalUnit.DAY,
2300+
partition_interval_unit=None,
23012301
clustered_by=[],
23022302
table_properties={},
23032303
table_description=None,
@@ -2372,7 +2372,7 @@ def test_create_seed_on_error(mocker: MockerFixture, adapter_mock, make_snapshot
23722372
table_format=None,
23732373
storage_format=None,
23742374
partitioned_by=[],
2375-
partition_interval_unit=IntervalUnit.DAY,
2375+
partition_interval_unit=None,
23762376
clustered_by=[],
23772377
table_properties={},
23782378
table_description=None,
@@ -2428,7 +2428,7 @@ def test_create_seed_no_intervals(mocker: MockerFixture, adapter_mock, make_snap
24282428
table_format=None,
24292429
storage_format=None,
24302430
partitioned_by=[],
2431-
partition_interval_unit=IntervalUnit.DAY,
2431+
partition_interval_unit=None,
24322432
clustered_by=[],
24332433
table_properties={},
24342434
table_description=None,
@@ -2968,7 +2968,7 @@ def test_create_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
29682968
table_format=model.table_format,
29692969
storage_format=model.storage_format,
29702970
partitioned_by=model.partitioned_by,
2971-
partition_interval_unit=model.interval_unit,
2971+
partition_interval_unit=model.partition_interval_unit,
29722972
clustered_by=model.clustered_by,
29732973
table_properties=model.physical_properties,
29742974
table_description=None,
@@ -3054,7 +3054,7 @@ def test_evaluate_managed(adapter_mock, make_snapshot, mocker: MockerFixture):
30543054
table_format=model.table_format,
30553055
storage_format=model.storage_format,
30563056
partitioned_by=model.partitioned_by,
3057-
partition_interval_unit=model.interval_unit,
3057+
partition_interval_unit=model.partition_interval_unit,
30583058
clustered_by=model.clustered_by,
30593059
table_properties=model.physical_properties,
30603060
table_description=model.description,
@@ -3221,7 +3221,7 @@ def test_create_snapshot(
32213221
table_format=None,
32223222
storage_format=None,
32233223
partitioned_by=[],
3224-
partition_interval_unit=IntervalUnit.DAY,
3224+
partition_interval_unit=None,
32253225
clustered_by=[],
32263226
table_properties={},
32273227
table_description=None,
@@ -3268,7 +3268,7 @@ def test_migrate_snapshot(snapshot: Snapshot, mocker: MockerFixture, adapter_moc
32683268
table_format=None,
32693269
storage_format=None,
32703270
partitioned_by=[],
3271-
partition_interval_unit=IntervalUnit.DAY,
3271+
partition_interval_unit=None,
32723272
clustered_by=[],
32733273
table_properties={},
32743274
table_description=None,

0 commit comments

Comments
 (0)