sqlmesh.utils.metaprogramming

  1from __future__ import annotations
  2
  3import ast
  4import dis
  5import importlib
  6import inspect
  7import linecache
  8import logging
  9import os
 10import re
 11import sys
 12import textwrap
 13import types
 14import typing as t
 15from dataclasses import dataclass
 16from enum import Enum
 17from numbers import Number
 18from pathlib import Path
 19
 20from astor import to_source
 21
 22from sqlmesh.core import constants as c
 23from sqlmesh.utils import format_exception, unique
 24from sqlmesh.utils.errors import SQLMeshError
 25from sqlmesh.utils.pydantic import PydanticModel
 26
 27logger = logging.getLogger(__name__)
 28
 29
 30IGNORE_DECORATORS = {"macro", "model", "signal"}
 31SERIALIZABLE_CALLABLES = (type, types.FunctionType)
 32LITERALS = (Number, str, bytes, tuple, list, dict, set, bool)
 33
 34
 35def _is_relative_to(path: t.Optional[Path | str], other: t.Optional[Path | str]) -> bool:
 36    if path is None or other is None:
 37        return False
 38
 39    if isinstance(path, str):
 40        path = Path(path)
 41    if isinstance(other, str):
 42        other = Path(other)
 43
 44    if "site-packages" in str(path) or not path.exists() or not other.exists():
 45        return False
 46
 47    try:
 48        path.absolute().relative_to(other.absolute())
 49        return True
 50    except ValueError:
 51        return False
 52
 53
 54def _code_globals(code: types.CodeType) -> t.Dict[str, None]:
 55    variables = {
 56        instruction.argval: None
 57        for instruction in dis.get_instructions(code)
 58        if instruction.opname == "LOAD_GLOBAL"
 59    }
 60
 61    for const in code.co_consts:
 62        if isinstance(const, types.CodeType):
 63            variables.update(_code_globals(const))
 64
 65    return variables
 66
 67
 68def _globals_match(obj1: t.Any, obj2: t.Any) -> bool:
 69    return type(obj1) == type(obj2) and (
 70        obj1 == obj2
 71        or (
 72            getattr(obj1, "__module__", None) == getattr(obj2, "__module__", None)
 73            and getattr(obj1, "__name__", None) == getattr(obj2, "__name__", None)
 74        )
 75    )
 76
 77
 78def func_globals(func: t.Callable) -> t.Dict[str, t.Any]:
 79    """Finds all global references and closures in a function and nested functions.
 80
 81    This function treats closures as global variables, which could cause problems in the future.
 82
 83    Args:
 84        func: The function to introspect
 85
 86    Returns:
 87        A dictionary of all global references.
 88    """
 89    variables = {}
 90
 91    if hasattr(func, "__code__"):
 92        root_node = parse_source(func)
 93
 94        func_args = next(node for node in ast.walk(root_node) if isinstance(node, ast.arguments))
 95        arg_defaults = (d for d in func_args.defaults + func_args.kw_defaults if d is not None)
 96
 97        # ast.Name corresponds to variable references, such as foo or x.foo. The former is
 98        # represented as Name(id=foo), and the latter as Attribute(value=Name(id=x) attr=foo)
 99        arg_globals = [
100            n.id for default in arg_defaults for n in ast.walk(default) if isinstance(n, ast.Name)
101        ]
102
103        code = func.__code__
104        for var in (
105            arg_globals + list(_code_globals(code)) + decorator_vars(func, root_node=root_node)
106        ):
107            if var in func.__globals__:
108                variables[var] = func.__globals__[var]
109
110        if func.__closure__:
111            for var, value in zip(code.co_freevars, func.__closure__):
112                variables[var] = value.cell_contents
113
114    return variables
115
116
117class ClassFoundException(Exception):
118    pass
119
120
121class _ClassFinder(ast.NodeVisitor):
122    def __init__(self, qualname: str) -> None:
123        self.stack: t.List[str] = []
124        self.qualname = qualname
125
126    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
127        self.stack.append(node.name)
128        self.stack.append("<locals>")
129        self.generic_visit(node)
130        self.stack.pop()
131        self.stack.pop()
132
133    visit_AsyncFunctionDef = visit_FunctionDef  # type: ignore
134
135    def visit_ClassDef(self, node: ast.ClassDef) -> None:
136        self.stack.append(node.name)
137        if self.qualname == ".".join(self.stack):
138            # Return the decorator for the class if present
139            if node.decorator_list:
140                line_number = node.decorator_list[0].lineno
141            else:
142                line_number = node.lineno
143
144            # decrement by one since lines starts with indexing by zero
145            line_number -= 1
146            raise ClassFoundException(line_number)
147        self.generic_visit(node)
148        self.stack.pop()
149
150
151class _DecoratorDependencyFinder(ast.NodeVisitor):
152    def __init__(self) -> None:
153        self.dependencies: t.List[str] = []
154
155    def _extract_dependencies(self, node: ast.ClassDef | ast.FunctionDef) -> None:
156        for decorator in node.decorator_list:
157            dependencies: t.List[str] = []
158            for n in ast.walk(decorator):
159                if isinstance(n, ast.Attribute):
160                    dep = n.attr
161                elif isinstance(n, ast.Name):
162                    dep = n.id
163                else:
164                    continue
165
166                if dep in IGNORE_DECORATORS:
167                    dependencies = []
168                    break
169
170                dependencies.append(dep)
171
172            self.dependencies.extend(dependencies)
173
174    def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
175        self._extract_dependencies(node)
176
177    def visit_ClassDef(self, node: ast.ClassDef) -> None:
178        self._extract_dependencies(node)
179
180    visit_AsyncFunctionDef = visit_FunctionDef  # type: ignore
181
182
183def getsource(obj: t.Any) -> str:
184    """Get the source of a function or class.
185
186    inspect.getsource doesn't find decorators in python < 3.9
187    https://github.com/python/cpython/commit/696136b993e11b37c4f34d729a0375e5ad544ade
188    """
189    path = inspect.getsourcefile(obj)
190    if path:
191        module = inspect.getmodule(obj, path)
192
193        if module:
194            lines = linecache.getlines(path, module.__dict__)
195        else:
196            lines = linecache.getlines(path)
197
198        def join_source(lnum: int) -> str:
199            return "".join(inspect.getblock(lines[lnum:]))
200
201        if inspect.isclass(obj):
202            qualname = obj.__qualname__
203            source = "".join(lines)
204            tree = ast.parse(source)
205            class_finder = _ClassFinder(qualname)
206            try:
207                class_finder.visit(tree)
208            except ClassFoundException as e:
209                return join_source(e.args[0])
210        elif inspect.isfunction(obj):
211            obj = obj.__code__
212            if hasattr(obj, "co_firstlineno"):
213                lnum = obj.co_firstlineno - 1
214                pat = re.compile(r"^(\s*def\s)|(\s*async\s+def\s)|(.*(?<!\w)lambda(:|\s))|^(\s*@)")
215                while lnum > 0:
216                    try:
217                        line = lines[lnum]
218                    except IndexError:
219                        raise OSError("lineno is out of bounds")
220                    if pat.match(line):
221                        break
222                    lnum = lnum - 1
223                return join_source(lnum)
224    raise SQLMeshError(f"Cannot find source for {obj}")
225
226
227def parse_source(func: t.Callable) -> ast.Module:
228    """Parse a function and returns an ast node."""
229    return ast.parse(textwrap.dedent(getsource(func)))
230
231
232def _decorator_name(decorator: ast.expr) -> str:
233    node = decorator
234    if isinstance(decorator, ast.Call):
235        node = decorator.func
236    return node.id if isinstance(node, ast.Name) else ""
237
238
239def decorator_vars(func: t.Callable, root_node: t.Optional[ast.Module] = None) -> t.List[str]:
240    """
241    Returns a list of all the decorators of a callable, as well as names of objects that
242    are referenced in their argument list. These objects may be transitive dependencies
243    that we need to include in the serialized python environments.
244    """
245    root_node = root_node or parse_source(func)
246    finder = _DecoratorDependencyFinder()
247    finder.visit(root_node)
248    return unique(finder.dependencies)
249
250
251def normalize_source(obj: t.Any) -> str:
252    """Rewrites an object's source with formatting and doc strings removed by using Python ast.
253
254    Args:
255        obj: The object to fetch source from and convert to a string.
256
257    Returns:
258        A string representation of the normalized function.
259    """
260    root_node = parse_source(obj)
261
262    for node in ast.walk(root_node):
263        if isinstance(node, (ast.FunctionDef, ast.ClassDef)):
264            for decorator in node.decorator_list:
265                if _decorator_name(decorator) in IGNORE_DECORATORS:
266                    node.decorator_list.remove(decorator)
267
268            # remove docstrings
269            body = node.body
270            if body and isinstance(body[0], ast.Expr) and isinstance(body[0].value, ast.Str):
271                node.body = body[1:]
272
273            # remove function return type annotation
274            if isinstance(node, ast.FunctionDef):
275                node.returns = None
276
277    return to_source(root_node).strip()
278
279
280def build_env(
281    obj: t.Any,
282    *,
283    env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]],
284    name: str,
285    path: Path,
286    is_metadata_obj: bool = False,
287) -> None:
288    """Fills in env dictionary with all globals needed to execute the object.
289
290    Recursively traverse classes and functions.
291
292    Args:
293        obj: Any python object.
294        env: Dictionary to store the env.
295        name: Name of the object in the env.
296        path: The module path to serialize. Other modules will not be walked and treated as imports.
297        is_metadata_obj: An optional flag that determines whether the input object is metadata-only.
298    """
299    # We don't rely on `env` to keep track of visited objects, because it's populated in post-order
300    visited: t.Set[str] = set()
301
302    def walk(obj: t.Any, name: str, is_metadata: bool = False) -> None:
303        obj_module = inspect.getmodule(obj)
304        if obj_module and obj_module.__name__ == "builtins":
305            return
306
307        if name in visited:
308            if name not in env or _globals_match(env[name][0], obj):
309                return
310
311            raise SQLMeshError(
312                f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
313            )
314
315        visited.add(name)
316        name_missing_from_env = name not in env
317
318        if name_missing_from_env or (not is_metadata and env[name] == (obj, True)):
319            if not name_missing_from_env:
320                # The existing object in the env is "metadata only" but we're walking it again as a
321                # non-"metadata only" dependency, so we update this flag to ensure all transitive
322                # dependencies are also not marked as "metadata only"
323                is_metadata = False
324
325            if hasattr(obj, c.SQLMESH_MACRO):
326                # We only need to add the undecorated code of @macro() functions in env, which
327                # is accessible through the `__wrapped__` attribute added by functools.wraps
328                obj = obj.__wrapped__
329            elif callable(obj) and not isinstance(obj, SERIALIZABLE_CALLABLES):
330                obj = getattr(obj, "__wrapped__", None)
331                name = getattr(obj, "__name__", "")
332
333                # Callable class instances shouldn't be serialized (e.g. tenacity.Retrying).
334                # We still want to walk the callables they decorate, though
335                if not isinstance(obj, SERIALIZABLE_CALLABLES) or name in env:
336                    return
337
338            if (
339                not obj_module
340                or not hasattr(obj_module, "__file__")
341                or not _is_relative_to(obj_module.__file__, path)
342            ):
343                env[name] = (obj, is_metadata)
344                return
345
346            if inspect.isclass(obj):
347                for var in decorator_vars(obj):
348                    if obj_module and var in obj_module.__dict__:
349                        walk(obj_module.__dict__[var], var, is_metadata)
350
351                for base in obj.__bases__:
352                    walk(base, base.__qualname__, is_metadata)
353
354                for k, v in obj.__dict__.items():
355                    # skip dunder methods bar __init__ as it might contain user defined logic with cross class references
356                    if k.startswith("__") and k != "__init__":
357                        continue
358
359                    # Traverse methods in a class to find global references
360                    if isinstance(v, (classmethod, staticmethod)):
361                        v = v.__func__
362
363                    if callable(v):
364                        # Walk the method if it's part of the object, else it's a global function and we just store it
365                        if v.__qualname__.startswith(obj.__qualname__):
366                            try:
367                                for k, v in func_globals(v).items():
368                                    walk(v, k, is_metadata)
369                            except (OSError, TypeError):
370                                # __init__ may come from built-ins or wrapped callables
371                                pass
372                    else:
373                        walk(v, k, is_metadata)
374            elif callable(obj):
375                for k, v in func_globals(obj).items():
376                    walk(v, k, is_metadata)
377
378            # We store the object in the environment after its dependencies, because otherwise we
379            # could crash at environment hydration time, since dicts are ordered and the top-level
380            # objects would be loaded before their dependencies.
381            env[name] = (obj, is_metadata)
382        elif not _globals_match(env[name][0], obj):
383            raise SQLMeshError(
384                f"Cannot store {obj} in environment, duplicate definitions found for '{name}'"
385            )
386
387    # The "metadata only" annotation of the object is transitive
388    walk(obj, name, is_metadata_obj or getattr(obj, c.SQLMESH_METADATA, False))
389
390
391@dataclass
392class SqlValue:
393    """A SQL string representing a generated SQLGlot AST."""
394
395    sql: str
396
397
398class ExecutableKind(str, Enum):
399    """The kind of of executable. The order of the members is used when serializing the python model to text."""
400
401    IMPORT = "import"
402    VALUE = "value"
403    DEFINITION = "definition"
404
405    def __lt__(self, other: t.Any) -> bool:
406        if not isinstance(other, ExecutableKind):
407            return NotImplemented
408        values = list(ExecutableKind.__dict__.values())
409        return values.index(self) < values.index(other)
410
411    def __str__(self) -> str:
412        return self.value
413
414
415class Executable(PydanticModel):
416    payload: str
417    kind: ExecutableKind = ExecutableKind.DEFINITION
418    name: t.Optional[str] = None
419    path: t.Optional[str] = None
420    alias: t.Optional[str] = None
421    is_metadata: t.Optional[bool] = None
422
423    @property
424    def is_definition(self) -> bool:
425        return self.kind == ExecutableKind.DEFINITION
426
427    @property
428    def is_import(self) -> bool:
429        return self.kind == ExecutableKind.IMPORT
430
431    @property
432    def is_value(self) -> bool:
433        return self.kind == ExecutableKind.VALUE
434
435    @classmethod
436    def value(
437        cls, v: t.Any, is_metadata: t.Optional[bool] = None, sort_root_dict: bool = False
438    ) -> Executable:
439        payload = _dict_sort(v) if sort_root_dict else repr(v)
440        return Executable(
441            payload=payload,
442            kind=ExecutableKind.VALUE,
443            is_metadata=is_metadata or None,
444        )
445
446
447def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable]:
448    """Serializes a python function into a self contained dictionary.
449
450    Recursively walks a function's globals to store all other references inside of env.
451
452    Args:
453        env: Dictionary to store the env.
454        path: The root path to seralize. Other modules will not be walked and treated as imports.
455    """
456    serialized = {}
457
458    for k, (v, is_metadata) in env.items():
459        # We don't store `False` for `is_metadata` to reduce the pydantic model's payload size
460        is_metadata = is_metadata or None
461
462        if isinstance(v, LITERALS) or v is None:
463            serialized[k] = Executable.value(v, is_metadata=is_metadata)
464        elif inspect.ismodule(v):
465            name = v.__name__
466            if hasattr(v, "__file__") and _is_relative_to(v.__file__, path):
467                raise SQLMeshError(
468                    f"Cannot serialize 'import {name}'. Use 'from {name} import ...' instead."
469                )
470            postfix = "" if name == k else f" as {k}"
471            serialized[k] = Executable(
472                payload=f"import {name}{postfix}",
473                kind=ExecutableKind.IMPORT,
474                is_metadata=is_metadata,
475            )
476        elif callable(v):
477            name = v.__name__
478            name = k if name == "<lambda>" else name
479
480            # getfile raises a `TypeError` for built-in modules, classes, or functions
481            # https://docs.python.org/3/library/inspect.html#inspect.getfile
482            try:
483                file_path = Path(inspect.getfile(v))
484                relative_obj_file_path = _is_relative_to(file_path, path)
485
486                # A callable can be a "wrapper" that is defined in a third-party library [1], in which case the file
487                # containing its definition won't be relative to the project's path. This can lead to serializing
488                # it as a "relative import", such as `from models.some_python_model import foo`, because the `wraps`
489                # decorator preserves the wrapped function's module [2]. Payloads like this are invalid, as they
490                # can result in `ModuleNotFoundError`s when hydrating python environments, e.g. if a project's files
491                # are not available during a scheduled cadence run.
492                #
493                # [1]: https://github.com/jd/tenacity/blob/0d40e76f7d06d631fb127e1ec58c8bd776e70d49/tenacity/__init__.py#L322-L346
494                # [2]: https://github.com/python/cpython/blob/f502c8f6a6db4be27c97a0e5466383d117859b7f/Lib/functools.py#L33-L57
495                if not relative_obj_file_path and (wrapped := getattr(v, "__wrapped__", None)):
496                    v = wrapped
497                    file_path = Path(inspect.getfile(wrapped))
498                    relative_obj_file_path = _is_relative_to(file_path, path)
499            except TypeError:
500                file_path = None
501                relative_obj_file_path = False
502
503            if relative_obj_file_path:
504                serialized[k] = Executable(
505                    name=name,
506                    payload=normalize_source(v),
507                    kind=ExecutableKind.DEFINITION,
508                    # Do `as_posix` to serialize windows path back to POSIX
509                    path=t.cast(Path, file_path).relative_to(path.absolute()).as_posix(),
510                    alias=k if name != k else None,
511                    is_metadata=is_metadata,
512                )
513            else:
514                serialized[k] = Executable(
515                    payload=f"from {v.__module__} import {name}",
516                    kind=ExecutableKind.IMPORT,
517                    is_metadata=is_metadata,
518                )
519        else:
520            raise SQLMeshError(
521                f"Object '{v}' cannot be serialized. If it's defined in a library, import the corresponding "
522                "module and reference the object using its fully-qualified name. For example, the datetime "
523                "module's 'UTC' object should be accessed as 'datetime.UTC'."
524            )
525
526    return serialized
527
528
529def prepare_env(
530    python_env: t.Dict[str, Executable],
531    env: t.Optional[t.Dict[str, t.Any]] = None,
532) -> t.Dict[str, t.Any]:
533    """Prepare a python env by hydrating and executing functions.
534
535    The Python ENV is stored in a json serializable format.
536    Functions and imports are stored as a special data class.
537
538    Args:
539        python_env: The dictionary containing the serialized python environment.
540        env: The dictionary to execute code in.
541
542    Returns:
543        The prepared environment with hydrated functions.
544    """
545    env = {} if env is None else env
546
547    for name, executable in sorted(
548        python_env.items(), key=lambda item: 0 if item[1].is_import else 1
549    ):
550        if executable.is_value:
551            env[name] = eval(executable.payload)
552        else:
553            exec(executable.payload, env)
554            if executable.alias and executable.name:
555                env[executable.alias] = env[executable.name]
556
557    return env
558
559
560def format_evaluated_code_exception(
561    exception: Exception,
562    python_env: t.Dict[str, Executable],
563) -> str:
564    """Formats exceptions that occur from evaled code.
565
566    Stack traces generated by evaled code lose code context and are difficult to debug.
567    This intercepts the default stack trace and tries to make it debuggable.
568
569    Args:
570        exception: The exception to print the stack trace for.
571        python_env: The environment containing stringified python code.
572    """
573    tb: t.List[str] = []
574    indent = ""
575
576    skip_patterns = re.compile(
577        r"Traceback \(most recent call last\):|"
578        r'File ".*?core/model/definition\.py|'
579        r'File ".*?core/snapshot/definition\.py|'
580        r'File ".*?core/macros\.py|'
581        r'File ".*?inspect\.py'
582    )
583
584    for error_line in format_exception(exception):
585        if skip_patterns.search(error_line):
586            continue
587
588        error_match = re.search("^.*?Error: ", error_line)
589        if error_match:
590            tb.append(f"{indent * 2}  {error_line}")
591            continue
592
593        eval_code_match = re.search('File "<string>", line (.*), in (.*)', error_line)
594        if not eval_code_match:
595            tb.append(f"{indent}{error_line}")
596            continue
597
598        line_num = int(eval_code_match.group(1))
599        func = eval_code_match.group(2)
600
601        if func not in python_env:
602            tb.append(error_line)
603            continue
604
605        executable = python_env[func]
606        indent = error_line[: eval_code_match.start()]
607
608        error_line = (
609            f"{indent}File '{executable.path}' (or imported file), line {line_num}, in {func}"
610        )
611
612        code = executable.payload
613        formatted = []
614
615        for i, code_line in enumerate(code.splitlines()):
616            if i < line_num:
617                pad = len(code_line) - len(code_line.lstrip())
618                if i + 1 == line_num:
619                    formatted.append(f"{code_line[:pad]}{code_line[pad:]}")
620                else:
621                    formatted.append(code_line)
622
623        tb.extend(
624            (
625                error_line,
626                textwrap.indent(
627                    os.linesep.join(formatted),
628                    indent + "  ",
629                ),
630            )
631        )
632
633    return os.linesep.join(tb)
634
635
636def print_exception(
637    exception: Exception,
638    python_env: t.Dict[str, Executable],
639    out: t.TextIO = sys.stderr,
640) -> None:
641    """Prints exceptions that occur from evaled code.
642
643    Stack traces generated by evaled code lose code context and are difficult to debug.
644    This intercepts the default stack trace and tries to make it debuggable.
645
646    Args:
647        exception: The exception to print the stack trace for.
648        python_env: The environment containing stringified python code.
649        out: The output stream to write to.
650    """
651    tb = format_evaluated_code_exception(exception, python_env)
652    out.write(tb)
653
654
655def _dict_sort(obj: t.Any) -> str:
656    try:
657        if isinstance(obj, dict):
658            obj = dict(sorted(obj.items(), key=lambda x: str(x[0])))
659    except Exception:
660        logger.warning("Failed to sort non-recursive dict", exc_info=True)
661    return repr(obj)
662
663
664def import_python_file(path: Path, relative_base: Path = Path()) -> types.ModuleType:
665    relative_path = path.absolute().relative_to(relative_base.absolute())
666    module_name = str(relative_path.with_suffix("")).replace(os.path.sep, ".")
667
668    # remove the entire module hierarchy in case they were already loaded
669    parts = module_name.split(".")
670    for i in range(len(parts)):
671        sys.modules.pop(".".join(parts[0 : i + 1]), None)
672
673    return importlib.import_module(module_name)