sqlmesh.utils.cache

  1from __future__ import annotations
  2
  3import gzip
  4import logging
  5import pickle
  6import shutil
  7import typing as t
  8from pathlib import Path
  9
 10from sqlglot import __version__ as SQLGLOT_VERSION
 11
 12from sqlmesh.utils import sanitize_name
 13from sqlmesh.utils.date import to_datetime
 14from sqlmesh.utils.errors import SQLMeshError
 15from sqlmesh.utils.windows import IS_WINDOWS, fix_windows_path
 16
 17logger = logging.getLogger(__name__)
 18
 19T = t.TypeVar("T")
 20
 21
 22SQLGLOT_VERSION_TUPLE = tuple(SQLGLOT_VERSION.split("."))
 23SQLGLOT_MAJOR_VERSION = SQLGLOT_VERSION_TUPLE[0]
 24SQLGLOT_MINOR_VERSION = SQLGLOT_VERSION_TUPLE[1]
 25
 26
 27class FileCache(t.Generic[T]):
 28    """Generic file-based cache implementation.
 29
 30    Args:
 31        path: The path to the cache folder.
 32        entry_class: The type of cached entries.
 33        prefix: The prefix shared between all entries to distinguish them from other entries
 34            stored in the same cache folder.
 35    """
 36
 37    def __init__(self, path: Path, prefix: t.Optional[str] = None):
 38        self._path = path / prefix if prefix else path
 39
 40        from sqlmesh.core.state_sync.base import SCHEMA_VERSION
 41
 42        try:
 43            from sqlmesh._version import __version_tuple__
 44
 45            major, minor = __version_tuple__[0], __version_tuple__[1]
 46        except ImportError:
 47            major, minor = 0, 0
 48
 49        self._cache_version = "_".join(
 50            [
 51                str(major),
 52                str(minor),
 53                SQLGLOT_MAJOR_VERSION,
 54                SQLGLOT_MINOR_VERSION,
 55                str(SCHEMA_VERSION),
 56            ]
 57        )
 58
 59        threshold = to_datetime("1 week ago").timestamp()
 60        # delete all old cache files
 61        for file in self._path.glob("*"):
 62            if IS_WINDOWS:
 63                # the file.stat() call below will fail on windows if the :file name is longer than 260 chars
 64                file = fix_windows_path(file)
 65
 66            try:
 67                stat_result = file.stat()
 68                if (
 69                    not file.stem.startswith(self._cache_version)
 70                    or stat_result.st_atime < threshold
 71                ):
 72                    file.unlink(missing_ok=True)
 73            except FileNotFoundError:
 74                # File was deleted between glob() and stat() — skip stale cache entries gracefully
 75                continue
 76
 77    def get_or_load(self, name: str, entry_id: str = "", *, loader: t.Callable[[], T]) -> T:
 78        """Returns an existing cached entry or loads and caches a new one.
 79
 80        Args:
 81            name: The name of the entry.
 82            entry_id: The unique entry identifier. Used for cache invalidation.
 83            loader: Used to load a new entry when no cached instance was found.
 84
 85        Returns:
 86            The entry.
 87        """
 88        cached_entry = self.get(name, entry_id)
 89        if cached_entry is not None:
 90            return cached_entry
 91
 92        loaded_entry = loader()
 93        self.put(name, entry_id, value=loaded_entry)
 94        return loaded_entry
 95
 96    def get(self, name: str, entry_id: str = "") -> t.Optional[T]:
 97        """Returns a cached entry if exists.
 98
 99        Args:
100            name: The name of the entry.
101            entry_id: The unique entry identifier. Used for cache invalidation.
102
103        Returns:
104            The entry or None if no entry was found in the cache.
105        """
106        cache_entry_path = self._cache_entry_path(name, entry_id)
107        if cache_entry_path.exists():
108            with gzip.open(cache_entry_path, "rb") as fd:
109                try:
110                    return pickle.load(fd)
111                except Exception as ex:
112                    logger.warning("Failed to load a cache entry '%s': %s", name, ex)
113
114        return None
115
116    def put(self, name: str, entry_id: str = "", *, value: T) -> None:
117        """Stores the given value in the cache.
118
119        Args:
120            name: The name of the entry.
121            entry_id: The unique entry identifier. Used for cache invalidation.
122            value: The value to store in the cache.
123        """
124        self._path.mkdir(parents=True, exist_ok=True)
125        if not self._path.is_dir():
126            raise SQLMeshError(f"Cache path '{self._path}' is not a directory.")
127
128        with gzip.open(self._cache_entry_path(name, entry_id), "wb", compresslevel=1) as fd:
129            pickle.dump(value, fd)
130
131    def exists(self, name: str, entry_id: str = "") -> bool:
132        """Returns true if the cache entry with the given name and ID exists, false otherwise.
133
134        Args:
135            name: The name of the entry.
136            entry_id: The unique entry identifier. Used for cache invalidation.
137        """
138        return self._cache_entry_path(name, entry_id).exists()
139
140    def clear(self) -> None:
141        try:
142            shutil.rmtree(str(self._path.absolute()))
143        except Exception:
144            pass
145
146    def _cache_entry_path(self, name: str, entry_id: str = "") -> Path:
147        entry_file_name = "__".join(p for p in (self._cache_version, name, entry_id) if p)
148        full_path = self._path / sanitize_name(entry_file_name, include_unicode=True)
149        if IS_WINDOWS:
150            # handle paths longer than 260 chars
151            full_path = fix_windows_path(full_path)
152        return full_path