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            if not file.stem.startswith(self._cache_version) or file.stat().st_atime < threshold:
 67                file.unlink(missing_ok=True)
 68
 69    def get_or_load(self, name: str, entry_id: str = "", *, loader: t.Callable[[], T]) -> T:
 70        """Returns an existing cached entry or loads and caches a new one.
 71
 72        Args:
 73            name: The name of the entry.
 74            entry_id: The unique entry identifier. Used for cache invalidation.
 75            loader: Used to load a new entry when no cached instance was found.
 76
 77        Returns:
 78            The entry.
 79        """
 80        cached_entry = self.get(name, entry_id)
 81        if cached_entry is not None:
 82            return cached_entry
 83
 84        loaded_entry = loader()
 85        self.put(name, entry_id, value=loaded_entry)
 86        return loaded_entry
 87
 88    def get(self, name: str, entry_id: str = "") -> t.Optional[T]:
 89        """Returns a cached entry if exists.
 90
 91        Args:
 92            name: The name of the entry.
 93            entry_id: The unique entry identifier. Used for cache invalidation.
 94
 95        Returns:
 96            The entry or None if no entry was found in the cache.
 97        """
 98        cache_entry_path = self._cache_entry_path(name, entry_id)
 99        if cache_entry_path.exists():
100            with gzip.open(cache_entry_path, "rb") as fd:
101                try:
102                    return pickle.load(fd)
103                except Exception as ex:
104                    logger.warning("Failed to load a cache entry '%s': %s", name, ex)
105
106        return None
107
108    def put(self, name: str, entry_id: str = "", *, value: T) -> None:
109        """Stores the given value in the cache.
110
111        Args:
112            name: The name of the entry.
113            entry_id: The unique entry identifier. Used for cache invalidation.
114            value: The value to store in the cache.
115        """
116        self._path.mkdir(parents=True, exist_ok=True)
117        if not self._path.is_dir():
118            raise SQLMeshError(f"Cache path '{self._path}' is not a directory.")
119
120        with gzip.open(self._cache_entry_path(name, entry_id), "wb", compresslevel=1) as fd:
121            pickle.dump(value, fd)
122
123    def exists(self, name: str, entry_id: str = "") -> bool:
124        """Returns true if the cache entry with the given name and ID exists, false otherwise.
125
126        Args:
127            name: The name of the entry.
128            entry_id: The unique entry identifier. Used for cache invalidation.
129        """
130        return self._cache_entry_path(name, entry_id).exists()
131
132    def clear(self) -> None:
133        try:
134            shutil.rmtree(str(self._path.absolute()))
135        except Exception:
136            pass
137
138    def _cache_entry_path(self, name: str, entry_id: str = "") -> Path:
139        entry_file_name = "__".join(p for p in (self._cache_version, name, entry_id) if p)
140        full_path = self._path / sanitize_name(entry_file_name, include_unicode=True)
141        if IS_WINDOWS:
142            # handle paths longer than 260 chars
143            full_path = fix_windows_path(full_path)
144        return full_path