sqlmesh.utils

  1from __future__ import annotations
  2
  3import copy
  4import importlib
  5import logging
  6import os
  7import random
  8import re
  9import string
 10import sys
 11import time
 12import traceback
 13import types
 14import typing as t
 15import uuid
 16from dataclasses import dataclass
 17from collections import defaultdict
 18from contextlib import contextmanager
 19from copy import deepcopy
 20from enum import IntEnum, Enum
 21from functools import lru_cache, reduce, wraps
 22from pathlib import Path
 23
 24import unicodedata
 25from sqlglot import exp
 26from sqlglot.dialects.dialect import Dialects
 27
 28logger = logging.getLogger(__name__)
 29
 30T = t.TypeVar("T")
 31KEY = t.TypeVar("KEY", bound=t.Hashable)
 32VALUE = t.TypeVar("VALUE")
 33ITEM = t.TypeVar("ITEM")
 34GROUP = t.TypeVar("GROUP")
 35DECORATOR_RETURN_TYPE = t.TypeVar("DECORATOR_RETURN_TYPE")
 36
 37ALPHANUMERIC = string.ascii_lowercase + string.digits
 38
 39
 40def optional_import(name: str) -> t.Optional[types.ModuleType]:
 41    """Optionally import a module.
 42
 43    Args:
 44        name: The name of the module to import.
 45    Returns:
 46        The module if it is installed.
 47    """
 48    try:
 49        module = importlib.import_module(name)
 50    except ImportError:
 51        return None
 52    return module
 53
 54
 55def major_minor(version: str) -> t.Tuple[int, int]:
 56    """Returns a tuple of just the major.minor for a version string (major.minor.patch)."""
 57    return t.cast(t.Tuple[int, int], tuple(int(part) for part in version.split(".")[0:2]))
 58
 59
 60def unique(iterable: t.Iterable[T], by: t.Callable[[T], t.Any] = lambda i: i) -> t.List[T]:
 61    return list({by(i): None for i in iterable})
 62
 63
 64def random_id(short: bool = False) -> str:
 65    if short:
 66        return "".join(random.choices(ALPHANUMERIC, k=8))
 67
 68    return uuid.uuid4().hex
 69
 70
 71class UniqueKeyDict(t.Dict[KEY, VALUE]):
 72    """Dict that raises when a duplicate key is set."""
 73
 74    def __init__(self, name: str, *args: t.Dict[KEY, VALUE], **kwargs: VALUE) -> None:
 75        self.name = name
 76        super().__init__(*args, **kwargs)
 77
 78    def __setitem__(self, k: KEY, v: VALUE) -> None:
 79        if k in self:
 80            raise ValueError(
 81                f"Duplicate key '{k}' found in UniqueKeyDict<{self.name}>. Call dict.update(...) if this is intentional."
 82            )
 83        super().__setitem__(k, v)
 84
 85
 86class AttributeDict(dict, t.Mapping[KEY, VALUE]):
 87    def __getattr__(self, key: t.Any) -> t.Optional[VALUE]:
 88        if key.startswith("__") and not hasattr(self, key):
 89            raise AttributeError
 90        return self.get(key)
 91
 92    def set(self, field: str, value: t.Any) -> str:
 93        self[field] = value
 94        # Return an empty string, so that this method can be used within Jinja
 95        return ""
 96
 97    def __deepcopy__(self, memo: t.Dict[t.Any, AttributeDict]) -> AttributeDict:
 98        copy: AttributeDict = AttributeDict()
 99        memo[id(self)] = copy
100        for k, v in self.items():
101            copy[k] = deepcopy(v, memo)
102        return copy
103
104    def __call__(self, **kwargs: t.Dict[str, t.Any]) -> str:
105        self.update(**kwargs)
106        # Return an empty string, so that this method can be used within Jinja
107        return ""
108
109    def __getstate__(self) -> t.Optional[t.Dict[t.Any, t.Any]]:
110        return None
111
112
113class registry_decorator:
114    """A decorator that registers itself."""
115
116    registry_name = ""
117    _registry: t.Optional[UniqueKeyDict] = None
118
119    @classmethod
120    def registry(cls) -> UniqueKeyDict:
121        if cls._registry is None:
122            cls._registry = UniqueKeyDict(cls.registry_name)
123        return cls._registry
124
125    def __init__(self, name: str = "") -> None:
126        self.name = name
127
128    def __call__(
129        self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
130    ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
131        self.func = func
132
133        registry = self.registry()
134        func_name = (self.name or func.__name__).lower()
135
136        try:
137            registry[func_name] = self
138        except ValueError:
139            # No need to raise due to duplicate key if the functions are identical
140            if func.__code__.co_code != registry[func_name].func.__code__.co_code:
141                raise ValueError(f"Duplicate name: '{func_name}'.")
142
143        @wraps(func)
144        def wrapper(*args: t.Any, **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
145            return func(*args, **kwargs)
146
147        return wrapper
148
149    @classmethod
150    def get_registry(cls) -> UniqueKeyDict:
151        """Get a copy of the registry"""
152        return UniqueKeyDict(cls.registry_name, **(cls._registry or {}))
153
154    @classmethod
155    def set_registry(cls, registry: UniqueKeyDict) -> None:
156        """Set the registry."""
157        cls._registry = registry
158
159
160@contextmanager
161def sys_path(*paths: Path) -> t.Iterator[None]:
162    """A context manager to temporarily add a path to 'sys.path'."""
163    inserted = set()
164
165    for path in paths:
166        path_str = str(path.absolute())
167
168        if path_str not in sys.path:
169            sys.path.insert(0, path_str)
170            inserted.add(path_str)
171
172    try:
173        yield
174    finally:
175        for path_str in inserted:
176            sys.path.remove(path_str)
177
178
179def format_exception(exception: BaseException) -> t.List[str]:
180    if sys.version_info < (3, 10):
181        return traceback.format_exception(type(exception), exception, exception.__traceback__)  # type: ignore
182    return traceback.format_exception(exception)  # type: ignore
183
184
185def word_characters_only(s: str, replacement_char: str = "_") -> str:
186    """
187    Replace all non-word characters in string with the replacement character.
188    Reference SO: https://stackoverflow.com/questions/1276764/stripping-everything-but-alphanumeric-chars-from-a-string-in-python/70310018#70310018
189
190    >>> word_characters_only("Hello, world!")
191    'Hello__world_'
192    >>> word_characters_only("Hello, world! 123", '')
193    'Helloworld123'
194    """
195    return re.sub(r"\W", replacement_char, s)
196
197
198def str_to_bool(s: t.Optional[str]) -> bool:
199    """
200    Convert a string to a boolean. disutils is being deprecated and it is recommended to implement your own version:
201    https://peps.python.org/pep-0632/
202
203    Unlike disutils, this actually returns a bool and never raises. If a value cannot be determined to be true
204    then false is returned.
205    """
206    if not s:
207        return False
208    return s.lower() in ("true", "1", "t", "y", "yes", "on")
209
210
211_debug_mode_enabled: bool = False
212
213
214def enable_debug_mode() -> None:
215    global _debug_mode_enabled
216    _debug_mode_enabled = True
217
218
219def debug_mode_enabled() -> bool:
220    return _debug_mode_enabled or str_to_bool(os.environ.get("SQLMESH_DEBUG"))
221
222
223def ttl_cache(ttl: int = 60, maxsize: int = 128000) -> t.Callable:
224    """Caches a function that clears whenever the current epoch / ttl seconds changes.
225
226    TTL is not exact, it is used as a salt. So by default, at every minute mark, the cache will be cleared.
227    This is done for simplicity.
228
229    Args:
230        ttl: The number of seconds to hold the cache for.
231        maxsize: The maximum size of the cache.
232    """
233
234    def decorator(func: t.Callable) -> t.Any:
235        @lru_cache(maxsize=maxsize)
236        def cache(tick: int, *args: t.Any, **kwargs: t.Any) -> t.Any:
237            return func(*args, **kwargs)
238
239        @wraps(func)
240        def wrap(*args: t.Any, **kwargs: t.Any) -> t.Any:
241            return cache(int(time.time() / ttl), *args, **kwargs)
242
243        return wrap
244
245    return decorator
246
247
248class classproperty(property):
249    """Similar to a normal property but works for class methods"""
250
251    def __get__(self, obj: t.Any, owner: t.Any = None) -> t.Any:
252        return classmethod(self.fget).__get__(None, owner)()  # type: ignore
253
254
255@contextmanager
256def env_vars(environ: dict[str, str]) -> t.Iterator[None]:
257    """A context manager to temporarily modify environment variables."""
258    old_environ = os.environ.copy()
259    os.environ.update(environ)
260
261    try:
262        yield
263    finally:
264        os.environ.clear()
265        os.environ.update(old_environ)
266
267
268def merge_dicts(*args: t.Dict) -> t.Dict:
269    """
270    Merges dicts. Just does key collision replacement
271    """
272
273    def merge(a: t.Dict, b: t.Dict) -> t.Dict:
274        for b_key, b_value in b.items():
275            a_value = a.get(b_key)
276            if isinstance(a_value, dict) and isinstance(b_value, dict):
277                merge(a_value, b_value)
278            elif isinstance(b_value, dict):
279                a[b_key] = copy.deepcopy(b_value)
280            else:
281                a[b_key] = b_value
282        return a
283
284    return reduce(merge, args, {})
285
286
287def sqlglot_dialects() -> str:
288    return "'" + "', '".join(Dialects.__members__.values()) + "'"
289
290
291NON_ALNUM = re.compile(r"[^a-zA-Z0-9_]")
292
293NON_ALUM_INCLUDE_UNICODE = re.compile(r"\W", flags=re.UNICODE)
294
295
296def sanitize_name(name: str, *, include_unicode: bool = False) -> str:
297    if include_unicode:
298        s = unicodedata.normalize("NFC", name)
299        s = NON_ALUM_INCLUDE_UNICODE.sub("_", s)
300        return s
301    return NON_ALNUM.sub("_", name)
302
303
304def groupby(
305    items: t.Iterable[ITEM],
306    func: t.Callable[[ITEM], GROUP],
307) -> t.DefaultDict[GROUP, t.List[ITEM]]:
308    grouped = defaultdict(list)
309    for item in items:
310        grouped[func(item)].append(item)
311    return grouped
312
313
314def columns_to_types_to_struct(
315    columns_to_types: t.Union[t.Dict[str, exp.DataType], t.Dict[str, str]],
316) -> exp.DataType:
317    """
318    Converts a dict of column names to types to a struct.
319    """
320    return exp.DataType(
321        this=exp.DataType.Type.STRUCT,
322        expressions=[
323            exp.ColumnDef(this=exp.to_identifier(k), kind=v) for k, v in columns_to_types.items()
324        ],
325        nested=True,
326    )
327
328
329def type_is_known(d_type: t.Union[exp.DataType, exp.ColumnDef]) -> bool:
330    """Checks that a given column type is known and not NULL."""
331    if isinstance(d_type, exp.ColumnDef):
332        if not d_type.kind:
333            return False
334        d_type = d_type.kind
335    if isinstance(d_type, exp.DataTypeParam):
336        return True
337    if d_type.is_type(exp.DataType.Type.UNKNOWN, exp.DataType.Type.NULL):
338        return False
339    if d_type.expressions:
340        return all(type_is_known(expression) for expression in d_type.expressions)
341    return True
342
343
344def columns_to_types_all_known(columns_to_types: t.Dict[str, exp.DataType]) -> bool:
345    """Checks that all column types are known and not NULL."""
346    return all(type_is_known(expression) for expression in columns_to_types.values())
347
348
349class Verbosity(IntEnum):
350    """Verbosity levels for SQLMesh output."""
351
352    DEFAULT = 0
353    VERBOSE = 1
354    VERY_VERBOSE = 2
355
356    @property
357    def is_default(self) -> bool:
358        return self == Verbosity.DEFAULT
359
360    @property
361    def is_verbose(self) -> bool:
362        return self == Verbosity.VERBOSE
363
364    @property
365    def is_very_verbose(self) -> bool:
366        return self == Verbosity.VERY_VERBOSE
367
368
369class CompletionStatus(Enum):
370    SUCCESS = "success"
371    FAILURE = "failure"
372    NOTHING_TO_DO = "nothing_to_do"
373
374    @property
375    def is_success(self) -> bool:
376        return self == CompletionStatus.SUCCESS
377
378    @property
379    def is_failure(self) -> bool:
380        return self == CompletionStatus.FAILURE
381
382    @property
383    def is_nothing_to_do(self) -> bool:
384        return self == CompletionStatus.NOTHING_TO_DO
385
386
387def to_snake_case(name: str) -> str:
388    return "".join(
389        f"_{c.lower()}" if c.isupper() and idx != 0 else c.lower() for idx, c in enumerate(name)
390    )
391
392
393class JobType(Enum):
394    PLAN = "SQLMESH_PLAN"
395    RUN = "SQLMESH_RUN"
396
397
398@dataclass(frozen=True)
399class CorrelationId:
400    """ID that is added to each query in order to identify the job that created it."""
401
402    job_type: JobType
403    job_id: str
404
405    def __str__(self) -> str:
406        return f"{self.job_type.value}: {self.job_id}"
407
408    @classmethod
409    def from_plan_id(cls, plan_id: str) -> CorrelationId:
410        return CorrelationId(JobType.PLAN, plan_id)
411
412
413def get_source_columns_to_types(
414    columns_to_types: t.Dict[str, exp.DataType],
415    source_columns: t.Optional[t.List[str]],
416) -> t.Dict[str, exp.DataType]:
417    source_column_lookup = set(source_columns) if source_columns else None
418    return {
419        k: v
420        for k, v in columns_to_types.items()
421        if not source_column_lookup or k in source_column_lookup
422    }