Repository: TobikoData/sqlmesh
Pull request #3648 (2025-01-17 14:41:14)

Fix!: Support full precision for TSQL date types
This update aims to enhance the support for MSSQL/T-SQL full precision when building the `end_dt` macro and date ranges, specifically for [date types](https://learn.microsoft.com/en-us/sql/t-sql/data-types/datetime2-transact-sql?view=sql-server-ver16) with up to 100 nano seconds precision, fixes: #3624 

tobymao: do we need a migration for this? cc @georgesittas 
georgesittas: > do we need a migration for this? cc @georgesittas

There's definitely a diff between main and this PR. E.g. if I change the filter in the example incremental model to:

```sql
-- instead of `@start_date` and `@end_date`
event_date BETWEEN @start_dt AND @end_dt
```

Then I see this after running first plan, switching to this branch, wiping out cache and re-running plan:

```
Models:
├── Directly Modified:
│   └── sqlmesh_example.incremental_model
└── Indirectly Modified:
    └── sqlmesh_example.full_model

Directly Modified: sqlmesh_example.incremental_model (Non-breaking)
└── Indirectly Modified Children:
    └── sqlmesh_example.full_model (Indirect Non-breaking)
Models needing backfill (missing dates):
└── sqlmesh_example.incremental_model: 2020-01-01 - 2025-01-21
```

Patch:
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 49906933..3f4b1b74 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1248,7 +1248,9 @@ class EngineAdapter:
         )
         if not columns_to_types or not columns_to_types_all_known(columns_to_types):
             columns_to_types = self.columns(table_name)
-        low, high = [time_formatter(dt, columns_to_types) for dt in make_inclusive(start, end)]
+        low, high = [
+            time_formatter(dt, columns_to_types) for dt in make_inclusive(start, end, self.dialect)
+        ]
         if isinstance(time_column, TimeColumn):
             time_column = time_column.column
         where = exp.Between(
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index fe890b11..69141e66 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -1622,7 +1622,7 @@ class PythonModel(_Model):
         **kwargs: t.Any,
     ) -> t.Iterator[QueryOrDF]:
         env = prepare_env(self.python_env)
-        start, end = make_inclusive(start or c.EPOCH, end or c.EPOCH)
+        start, end = make_inclusive(start or c.EPOCH, end or c.EPOCH, self.dialect)
         execution_time = to_datetime(execution_time or c.EPOCH)
 
         variables = env.get(c.SQLMESH_VARS, {})
diff --git a/sqlmesh/core/renderer.py b/sqlmesh/core/renderer.py
index 893ef174..ed8d6785 100644
--- a/sqlmesh/core/renderer.py
+++ b/sqlmesh/core/renderer.py
@@ -15,7 +15,7 @@ from sqlglot.optimizer.simplify import simplify
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
 from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
-from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive_end, to_datetime
+from sqlmesh.utils.date import TimeLike, date_dict, make_inclusive, to_datetime
 from sqlmesh.utils.errors import (
     ConfigError,
     MacroEvalError,
@@ -119,11 +119,17 @@ class BaseExpressionRenderer:
 
         expressions = [self._expression]
 
+        start_time, end_time = (
+            make_inclusive(start or c.EPOCH, end or c.EPOCH, self._dialect)
+            if not self._only_execution_time
+            else (None, None)
+        )
+
         render_kwargs = {
             **date_dict(
                 to_datetime(execution_time or c.EPOCH),
-                to_datetime(start or c.EPOCH) if not self._only_execution_time else None,
-                make_inclusive_end(end or c.EPOCH) if not self._only_execution_time else None,
+                start_time,
+                end_time,
             ),
             **kwargs,
         }
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index b5549aaa..85f11802 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -25,8 +25,8 @@ from sqlmesh.utils.date import (
     TimeLike,
     is_date,
     make_inclusive,
-    make_inclusive_end,
     make_exclusive,
+    make_inclusive_end,
     now,
     now_timestamp,
     time_like_to_str,
diff --git a/sqlmesh/migrations/v0067_add_tsql_date_full_precision.py b/sqlmesh/migrations/v0067_add_tsql_date_full_precision.py
new file mode 100644
index 00000000..1178ee01
--- /dev/null
+++ b/sqlmesh/migrations/v0067_add_tsql_date_full_precision.py
@@ -0,0 +1,5 @@
+"""Add full precision for tsql to support nanoseconds."""
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    pass
diff --git a/sqlmesh/utils/date.py b/sqlmesh/utils/date.py
index 20dae649..783541e4 100644
--- a/sqlmesh/utils/date.py
+++ b/sqlmesh/utils/date.py
@@ -17,6 +17,9 @@ from sqlglot import exp
 
 from sqlmesh.utils import ttl_cache
 
+if t.TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
 UTC = timezone.utc
 TimeLike = t.Union[date, datetime, str, int, float]
 DatetimeRange = t.Tuple[datetime, datetime]
@@ -216,8 +219,10 @@ def to_date(value: TimeLike, relative_base: t.Optional[datetime] = None) -> date
 
 
 def date_dict(
-    execution_time: TimeLike, start: t.Optional[TimeLike], end: t.Optional[TimeLike]
-) -> t.Dict[str, t.Union[str, datetime, date, float, int]]:
+    execution_time: TimeLike,
+    start: t.Optional[TimeLike],
+    end: t.Optional[TimeLike],
+) -> t.Dict[str, TimeLike]:
     """Creates a kwarg dictionary of datetime variables for use in SQL Contexts.
 
     Keys are like start_date, start_ds, end_date, end_ds...
@@ -284,7 +289,9 @@ def is_date(obj: TimeLike) -> bool:
         return False
 
 
-def make_inclusive(start: TimeLike, end: TimeLike) -> DatetimeRange:
+def make_inclusive(
+    start: TimeLike, end: TimeLike, dialect: t.Optional[DialectType] = ""
+) -> DatetimeRange:
     """Adjust start and end times to to become inclusive datetimes.
 
     SQLMesh treats start and end times as inclusive so that filters can be written as
@@ -295,7 +302,8 @@ def make_inclusive(start: TimeLike, end: TimeLike) -> DatetimeRange:
     In the ds ('2020-01-01') case, because start_ds and end_ds are categorical, between works even if
     start_ds and end_ds are equivalent. However, when we move to ts ('2022-01-01 12:00:00'), because timestamps
     are numeric, using simple equality doesn't make sense. When the end is not a categorical date, then it is
-    treated as an exclusive range and converted to inclusive by subtracting 1 microsecond.
+    treated as an exclusive range and converted to inclusive by subtracting 1 microsecond. If the dialect is
+    T-SQL then 1 nanoseconds is subtracted to account for the increased precision.
 
     Args:
         start: Start timelike object.
@@ -308,11 +316,14 @@ def make_inclusive(start: TimeLike, end: TimeLike) -> DatetimeRange:
     Returns:
         A tuple of inclusive datetime objects.
     """
-    return (to_datetime(start), make_inclusive_end(end))
+    return (to_datetime(start), make_inclusive_end(end, dialect=dialect))
 
 
-def make_inclusive_end(end: TimeLike) -> datetime:
-    return make_exclusive(end) - timedelta(microseconds=1)
+def make_inclusive_end(end: TimeLike, dialect: t.Optional[DialectType] = "") -> datetime:
+    exclusive_end = make_exclusive(end)
+    if dialect == "tsql":
+        return to_utc_timestamp(exclusive_end) - pd.Timedelta(1, unit="ns")
+    return exclusive_end - timedelta(microseconds=1)
 
 
 def make_exclusive(time: TimeLike) -> datetime:
@@ -322,6 +333,12 @@ def make_exclusive(time: TimeLike) -> datetime:
     return dt
 
 
+def to_utc_timestamp(time: datetime) -> pd.Timestamp:
+    if time.tzinfo is not None:
+        return pd.Timestamp(time).tz_convert("utc")
+    return pd.Timestamp(time, tz="utc")
+
+
 def validate_date_range(
     start: t.Optional[TimeLike],
     end: t.Optional[TimeLike],


Test patch:
diff --git a/tests/core/engine_adapter/test_mssql.py b/tests/core/engine_adapter/test_mssql.py
index 47d08475..31701cc1 100644
--- a/tests/core/engine_adapter/test_mssql.py
+++ b/tests/core/engine_adapter/test_mssql.py
@@ -10,6 +10,9 @@ from sqlglot import expressions as exp
 from sqlglot import parse_one
 
 from sqlmesh.core.engine_adapter.mssql import MSSQLEngineAdapter
+from sqlmesh.core.snapshot import SnapshotEvaluator, SnapshotChangeCategory
+from sqlmesh.core.model import load_sql_based_model
+from sqlmesh.core import dialect as d
 from sqlmesh.core.engine_adapter.shared import (
     DataObject,
     DataObjectType,
@@ -152,6 +155,113 @@ def test_table_exists(make_mocked_engine_adapter: t.Callable):
     assert not resp
 
 
+@pytest.mark.parametrize(
+    "select_expr, input_time, expected_sql",
+    [
+        # Respect the user's precision for datetimeoffset, time, datetime2
+        (
+            "SELECT ds::datetime2",
+            pd.Timestamp("2022-01-01 00:00:00.1234567"),
+            "CAST('2022-01-01 00:00:00.123456700' AS DATETIME2)",
+        ),
+        (
+            "SELECT ds::datetimeoffset(4)",
+            pd.Timestamp("2022-01-01 00:00:00.1234567"),
+            "CAST('2022-01-01 00:00:00.123456700+00:00' AS DATETIMEOFFSET(4))",
+        ),
+        (
+            "SELECT ds::time",
+            pd.Timestamp("2022-01-01 00:00:00.1234567"),
+            "CAST('2022-01-01 00:00:00.123456700' AS TIME)",
+        ),
+        # Respecting precision in datetimeoffset with time zone offsets
+        (
+            "SELECT ds::time(7)",
+            pd.Timestamp("2022-01-01 00:00:00.1234567+00:00"),
+            "CAST('2022-01-01 00:00:00.123456700' AS TIME(7))",
+        ),
+        (
+            "SELECT ds::datetimeoffset(6)",
+            pd.Timestamp("2022-01-01 00:00:00.1234567+02:00"),
+            "CAST('2021-12-31 22:00:00.123456700+00:00' AS DATETIMEOFFSET(6))",
+        ),
+        # For date types without nano-second precision, truncate as usual
+        (
+            "SELECT ds::datetime",
+            "2022-01-01 00:00:00.1234567+01:00",
+            "CAST('2021-12-31 23:00:00.123456' AS DATETIME)",
+        ),
+        (
+            "SELECT ds::smalldatetime",
+            "2022-01-01 00:00:00.1234567+00:00",
+            "CAST('2022-01-01 00:00:00.123456' AS SMALLDATETIME)",
+        ),
+        ("SELECT ds::date", "2022-01-01 00:00:00.001", "CAST('2022-01-01' AS DATE)"),
+    ],
+)
+def test_to_time_column(select_expr, input_time, expected_sql):
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name db.table,
+            kind INCREMENTAL_BY_TIME_RANGE(
+                time_column (ds)
+            ),
+            dialect tsql
+        );
+
+        {select_expr}
+        """
+    )
+    model = load_sql_based_model(expressions)
+    assert model.convert_to_time_column(input_time).sql("tsql") == expected_sql
+
+
+def test_incremental_by_time_datetimeoffset_precision(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_snapshot
+):
+    adapter = make_mocked_engine_adapter(MSSQLEngineAdapter)
+    adapter.get_current_catalog = mocker.MagicMock(return_value="other_catalog")
+
+    evaluator = SnapshotEvaluator(adapter)
+    parsed = d.parse(  # type: ignore
+        """
+            MODEL (
+                name test_schema.test_model,
+                kind INCREMENTAL_BY_TIME_RANGE (time_column ds),
+            );
+
+            SELECT a::int, ds::datetimeoffset FROM tbl as t WHERE t.ds BETWEEN @start_dt and @end_dt;
+            """,
+    )
+
+    model = load_sql_based_model(parsed, dialect="tsql")
+
+    snapshot = make_snapshot(model)
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    evaluator.evaluate(
+        snapshot,
+        start="2020-01-01",
+        end="2020-01-02",
+        execution_time="2020-01-02",
+        snapshots={},
+    )
+
+    assert adapter.cursor.execute.call_args_list[0][0][0] == (
+        f"MERGE INTO [sqlmesh__test_schema].[test_schema__test_model__{snapshot.version}] AS [__MERGE_TARGET__] USING "
+        "(SELECT [a] AS [a], [ds] AS [ds] FROM (SELECT CAST([a] AS INTEGER) AS [a], "
+        "CAST([ds] AS DATETIMEOFFSET) AS [ds] FROM [tbl] AS [t] WHERE [t].[ds] BETWEEN "
+        "CAST('2020-01-01 00:00:00+00:00' AS DATETIMEOFFSET) AT TIME ZONE 'UTC' AND "
+        "CAST('2020-01-02 23:59:59.999999999+00:00' AS DATETIMEOFFSET) AT TIME ZONE 'UTC') AS [_subquery] "
+        "WHERE [ds] BETWEEN CAST('2020-01-01 00:00:00+00:00' AS DATETIMEOFFSET) AND "
+        "CAST('2020-01-02 23:59:59.999999999+00:00' AS DATETIMEOFFSET)) AS [__MERGE_SOURCE__] ON (1 = 0) WHEN NOT "
+        "MATCHED BY SOURCE AND [ds] BETWEEN CAST('2020-01-01 00:00:00+00:00' AS DATETIMEOFFSET) AND "
+        "CAST('2020-01-02 23:59:59.999999999+00:00' AS DATETIMEOFFSET) THEN DELETE WHEN NOT MATCHED THEN INSERT "
+        "([a], [ds]) VALUES ([a], [ds]);"
+    )
+
+
 def test_insert_overwrite_by_time_partition_supports_insert_overwrite_pandas_not_exists(
     make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
 ):
diff --git a/tests/utils/test_date.py b/tests/utils/test_date.py
index 89723fc0..62a1b2a1 100644
--- a/tests/utils/test_date.py
+++ b/tests/utils/test_date.py
@@ -4,6 +4,7 @@ from datetime import date, datetime
 import pytest
 import time_machine
 from sqlglot import exp
+import pandas as pd
 
 from sqlmesh.utils.date import (
     UTC,
@@ -69,25 +70,25 @@ def test_to_timestamp() -> None:
 @pytest.mark.parametrize(
     "start_in, end_in, start_out, end_out",
     [
-        ("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-01 23:59:59.999999"),
-        ("2020-01-01", date(2020, 1, 1), "2020-01-01", "2020-01-01 23:59:59.999999"),
+        ("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-01 23:59:59.999999+00:00"),
+        ("2020-01-01", date(2020, 1, 1), "2020-01-01", "2020-01-01 23:59:59.999999+00:00"),
         (
             date(2020, 1, 1),
             date(2020, 1, 1),
             "2020-01-01",
-            "2020-01-01 23:59:59.999999",
+            "2020-01-01 23:59:59.999999+00:00",
         ),
         (
             "2020-01-01",
             "2020-01-01 12:00:00",
             "2020-01-01",
-            "2020-01-01 11:59:59.999999",
+            "2020-01-01 11:59:59.999999+00:00",
         ),
         (
             "2020-01-01",
             to_datetime("2020-01-02"),
             "2020-01-01",
-            "2020-01-01 23:59:59.999999",
+            "2020-01-01 23:59:59.999999+00:00",
         ),
     ],
 )
@@ -98,6 +99,47 @@ def test_make_inclusive(start_in, end_in, start_out, end_out) -> None:
     )
 
 
+@pytest.mark.parametrize(
+    "start_in, end_in, start_out, end_out, dialect",
+    [
+        ("2020-01-01", "2020-01-01", "2020-01-01", "2020-01-01 23:59:59.999999999+00:00", "tsql"),
+        (
+            "2020-01-01",
+            date(2020, 1, 1),
+            "2020-01-01",
+            "2020-01-01 23:59:59.999999999+00:00",
+            "tsql",
+        ),
+        (
+            date(2020, 1, 1),
+            date(2020, 1, 1),
+            "2020-01-01",
+            "2020-01-01 23:59:59.999999999+00:00",
+            "tsql",
+        ),
+        (
+            "2020-01-01",
+            "2020-01-01 12:00:00",
+            "2020-01-01",
+            "2020-01-01 11:59:59.999999999+00:00",
+            "tsql",
+        ),
+        (
+            "2020-01-01",
+            to_datetime("2020-01-02"),
+            "2020-01-01",
+            "2020-01-01 23:59:59.999999999+00:00",
+            "tsql",
+        ),
+    ],
+)
+def test_make_inclusive_tsql(start_in, end_in, start_out, end_out, dialect) -> None:
+    assert make_inclusive(start_in, end_in, "tsql") == (
+        to_datetime(start_out),
+        pd.Timestamp(end_out),
+    )
+
+
 @pytest.mark.parametrize(
     "expression, result",
     [
@@ -179,6 +221,20 @@ def test_to_tstz():
             "%Y%m%d",
             "20200101",
         ),
+        (
+            "2020-01-01 00:00:00+00:00",
+            exp.DataType.build("TIMESTAMPTZ"),
+            "tsql",
+            "%Y%m%d",
+            "CAST('2020-01-01 00:00:00+00:00' AS DATETIMEOFFSET)",
+        ),
+        (
+            pd.Timestamp("2020-01-01 00:00:00.1234567+00:00"),
+            exp.DataType.build("DATETIME2", dialect="tsql"),
+            "tsql",
+            None,
+            "CAST('2020-01-01 00:00:00.123456700' AS DATETIME2)",
+        ),
     ],
 )
 def test_to_time_column(
@@ -189,7 +245,8 @@ def test_to_time_column(
     result: str,
 ):
     assert (
-        to_time_column(time_column, time_column_type, dialect, time_column_format).sql() == result
+        to_time_column(time_column, time_column_type, dialect, time_column_format).sql(dialect)
+        == result
     )
 
 
@@ -231,6 +288,38 @@ def test_date_dict():
     }
 
 
+@pytest.mark.parametrize(
+    "start, end, expected_start_dt, expected_end_dt",
+    [
+        (
+            "2020-01-01 00:00:00.1234567",
+            "2020-01-02",
+            to_datetime("2020-01-01 00:00:00.1234567+00:00"),
+            pd.Timestamp("2020-01-02 23:59:59.999999999+00:00"),
+        ),
+        (
+            "2020-01-01 00:00:00.1234567",
+            "2020-01-02 00:00:00.1234567",
+            to_datetime("2020-01-01 00:00:00.1234567+00:00"),
+            pd.Timestamp("2020-01-02 00:00:00.123455999+00:00"),
+        ),
+        (
+            "2020-01-01 00:00:00.1234567",
+            "2020-01-02 00:00:00",
+            to_datetime("2020-01-01 00:00:00.1234567+00:00"),
+            pd.Timestamp("2020-01-01 23:59:59.999999999+00:00"),
+        ),
+    ],
+)
+def test_tsql_date_dict(start, end, expected_start_dt, expected_end_dt):
+    resp = date_dict(
+        "2020-01-02 01:00:00",
+        *make_inclusive(start, end, "tsql"),
+    )
+    assert resp["start_dt"] == expected_start_dt
+    assert resp["end_dt"] == expected_end_dt
+
+
 def test_format_tz_datetime():
     test_datetime = to_datetime("2020-01-01 00:00:00")
     assert format_tz_datetime(test_datetime) == "2020-01-01 12:00AM UTC"