Repository: TobikoData/sqlmesh
Pull request #3057 (2024-08-27 16:15:08)

fix: cache inclusive exclusive for performance


Patch:
diff --git a/sqlmesh/core/snapshot/definition.py b/sqlmesh/core/snapshot/definition.py
index 90514943..31ab318d 100644
--- a/sqlmesh/core/snapshot/definition.py
+++ b/sqlmesh/core/snapshot/definition.py
@@ -769,25 +769,14 @@ class Snapshot(PydanticModel, SnapshotInfoMixin):
         Returns:
             A [start, end) pair.
         """
-        interval_unit = self.node.interval_unit
-        start_ts = to_timestamp(interval_unit.cron_floor(start))
-        if start_ts < to_timestamp(start) and not self.model.allow_partials:
-            start_ts = to_timestamp(interval_unit.cron_next(start_ts))
-
-        if is_date(end):
-            end = to_datetime(end) + timedelta(days=1)
-        end_ts = to_timestamp(interval_unit.cron_floor(end) if not allow_partial else end)
-        if end_ts < start_ts and to_timestamp(end) > to_timestamp(start) and not strict:
-            # This can happen when the interval unit is coarser than the size of the input interval.
-            # For example, if the interval unit is monthly, but the input interval is only 1 hour long.
-            return (start_ts, end_ts)
-
-        if (strict and start_ts >= end_ts) or (start_ts > end_ts):
-            raise ValueError(
-                f"`end` ({to_datetime(end_ts)}) must be greater than `start` ({to_datetime(start_ts)})"
-            )
-
-        return (start_ts, end_ts)
+        return inclusive_exclusive(
+            start,
+            end,
+            self.node.interval_unit,
+            model_allow_partials=self.is_model and self.model.allow_partials,
+            strict=strict,
+            allow_partial=allow_partial,
+        )
 
     def merge_intervals(self, other: t.Union[Snapshot, SnapshotIntervals]) -> None:
         """Inherits intervals from the target snapshot.
@@ -1740,6 +1729,48 @@ def compute_missing_intervals(
     return sorted(missing)
 
 
+@lru_cache(maxsize=None)
+def inclusive_exclusive(
+    start: TimeLike,
+    end: TimeLike,
+    interval_unit: IntervalUnit,
+    model_allow_partials: bool,
+    strict: bool = True,
+    allow_partial: bool = False,
+) -> Interval:
+    """Transform the inclusive start and end into a [start, end) pair.
+
+    Args:
+        start: The start date/time of the interval (inclusive)
+        end: The end date/time of the interval (inclusive)
+        interval_unit: The interval unit.
+        model_allow_partials: Whether or not the model allows partials.
+        strict: Whether to fail when the inclusive start is the same as the exclusive end.
+        allow_partial: Whether the interval can be partial or not.
+
+    Returns:
+        A [start, end) pair.
+    """
+    start_ts = to_timestamp(interval_unit.cron_floor(start))
+    if start_ts < to_timestamp(start) and not model_allow_partials:
+        start_ts = to_timestamp(interval_unit.cron_next(start_ts))
+
+    if is_date(end):
+        end = to_datetime(end) + timedelta(days=1)
+    end_ts = to_timestamp(interval_unit.cron_floor(end) if not allow_partial else end)
+    if end_ts < start_ts and to_timestamp(end) > to_timestamp(start) and not strict:
+        # This can happen when the interval unit is coarser than the size of the input interval.
+        # For example, if the interval unit is monthly, but the input interval is only 1 hour long.
+        return (start_ts, end_ts)
+
+    if (strict and start_ts >= end_ts) or (start_ts > end_ts):
+        raise ValueError(
+            f"`end` ({to_datetime(end_ts)}) must be greater than `start` ({to_datetime(start_ts)})"
+        )
+
+    return (start_ts, end_ts)
+
+
 def earliest_start_date(
     snapshots: t.Collection[Snapshot],
     cache: t.Optional[t.Dict[str, datetime]] = None,
diff --git a/sqlmesh/utils/metaprogramming.py b/sqlmesh/utils/metaprogramming.py
index 31743c7a..2da34688 100644
--- a/sqlmesh/utils/metaprogramming.py
+++ b/sqlmesh/utils/metaprogramming.py
@@ -359,11 +359,12 @@ def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable
             name = v.__name__
             name = k if name == "<lambda>" else name
 
-            # We can't call getfile on built-in callables
+            # getfile raises a `TypeError` for built-in modules, classes, or functions
             # https://docs.python.org/3/library/inspect.html#inspect.getfile
-            file_path = (
-                Path(inspect.getfile(inspect.unwrap(v))) if not inspect.isbuiltin(v) else None
-            )
+            try:
+                file_path = Path(inspect.getfile(v))
+            except TypeError:
+                file_path = None
 
             if _is_relative_to(file_path, path):
                 serialized[k] = Executable(


Test patch:
diff --git a/tests/utils/test_metaprogramming.py b/tests/utils/test_metaprogramming.py
index 8537851a..be360f9e 100644
--- a/tests/utils/test_metaprogramming.py
+++ b/tests/utils/test_metaprogramming.py
@@ -1,4 +1,5 @@
 import typing as t
+from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -45,7 +46,7 @@ def test_print_exception(mocker: MockerFixture):
 
     expected_message = f"""Traceback (most recent call last):
 
-  File "{__file__}", line 42, in test_print_exception
+  File "{__file__}", line 43, in test_print_exception
     eval("test_fun()", env)
 
   File "<string>", line 1, in <module>
@@ -106,6 +107,11 @@ def noop_metadata() -> None:
 setattr(noop_metadata, c.SQLMESH_METADATA, True)
 
 
+@contextmanager
+def test_context_manager():
+    yield
+
+
 def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2) -> int:
     """DOC STRING"""
     sqlglot.parse_one("1")
@@ -117,6 +123,9 @@ def main_func(y: int, foo=exp.true(), *, bar=expressions.Literal.number(1) + 2)
     def closure(z: int) -> int:
         return z + Z
 
+    with test_context_manager():
+        pass
+
     return closure(y) + other_func(Y)
 
 
@@ -132,6 +141,7 @@ def test_func_globals() -> None:
         "sqlglot": sqlglot,
         "exp": exp,
         "expressions": exp,
+        "test_context_manager": test_context_manager,
     }
     assert func_globals(other_func) == {
         "X": 1,
@@ -167,6 +177,8 @@ def test_normalize_source() -> None:
 
     def closure(z: int):
         return z + Z
+    with test_context_manager():
+        pass
     return closure(y) + other_func(Y)"""
     )
 
@@ -194,6 +206,7 @@ def test_serialize_env() -> None:
     build_env(main_func, env=env, name="MAIN", path=path)
     env = serialize_env(env, path=path)  # type: ignore
 
+    assert prepare_env(env)
     assert env == {
         "MAIN": Executable(
             name="main_func",
@@ -209,11 +222,19 @@ def test_serialize_env() -> None:
 
     def closure(z: int):
         return z + Z
+    with test_context_manager():
+        pass
     return closure(y) + other_func(Y)""",
         ),
         "X": Executable(payload="1", kind=ExecutableKind.VALUE),
         "Y": Executable(payload="2", kind=ExecutableKind.VALUE),
         "Z": Executable(payload="3", kind=ExecutableKind.VALUE),
+        "_GeneratorContextManager": Executable(
+            payload="from contextlib import _GeneratorContextManager", kind=ExecutableKind.IMPORT
+        ),
+        "contextmanager": Executable(
+            payload="from contextlib import contextmanager", kind=ExecutableKind.IMPORT
+        ),
         "KLASS_X": Executable(payload="1", kind=ExecutableKind.VALUE),
         "KLASS_Y": Executable(payload="2", kind=ExecutableKind.VALUE),
         "KLASS_Z": Executable(payload="3", kind=ExecutableKind.VALUE),
@@ -255,6 +276,14 @@ class DataClass:
         "expressions": Executable(
             kind=ExecutableKind.IMPORT, payload="import sqlglot.expressions as expressions"
         ),
+        "func": Executable(
+            payload="""@contextmanager
+def test_context_manager():
+    yield""",
+            name="test_context_manager",
+            path="test_metaprogramming.py",
+            alias="func",
+        ),
         "my_lambda": Executable(
             name="my_lambda",
             path="test_metaprogramming.py",
@@ -282,4 +311,9 @@ class DataClass:
     my_lambda()
     return X + a""",
         ),
+        "test_context_manager": Executable(
+            payload="from tests.utils.test_metaprogramming import test_context_manager",
+            kind=ExecutableKind.IMPORT,
+        ),
+        "wraps": Executable(payload="from functools import wraps", kind=ExecutableKind.IMPORT),
     }