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 _resolve_import_module(obj: t.Any, name: str) -> str:
448    """Resolve the most appropriate module path for importing an object.
449
450    When a callable's ``__module__`` points to a submodule of a known public
451    module (e.g. ``sqlglot.expressions.builders`` is a submodule of
452    ``sqlglot.expressions``), and the object is re-exported from that public
453    parent module, prefer the public parent so that generated import statements
454    remain stable across internal restructurings of third-party packages.
455
456    Args:
457        obj: The callable to resolve.
458        name: The name under which the object will be imported.
459
460    Returns:
461        The module path to use in the ``from <module> import <name>`` statement.
462    """
463    module_name = getattr(obj, "__module__", None) or ""
464    parts = module_name.split(".")
465
466    # Walk from the shallowest ancestor (excluding the top-level package) up to
467    # the immediate parent, returning the shallowest one that re-exports the object.
468    # We skip the top-level package to avoid over-normalizing (e.g. ``sqlglot``
469    # re-exports everything, but callers expect ``sqlglot.expressions``).
470    for i in range(2, len(parts)):
471        parent = ".".join(parts[:i])
472        try:
473            parent_module = sys.modules.get(parent) or importlib.import_module(parent)
474            if getattr(parent_module, name, None) is obj:
475                return parent
476        except Exception:
477            continue
478
479    return module_name
480
481
482def serialize_env(env: t.Dict[str, t.Any], path: Path) -> t.Dict[str, Executable]:
483    """Serializes a python function into a self contained dictionary.
484
485    Recursively walks a function's globals to store all other references inside of env.
486
487    Args:
488        env: Dictionary to store the env.
489        path: The root path to seralize. Other modules will not be walked and treated as imports.
490    """
491    serialized = {}
492
493    for k, (v, is_metadata) in env.items():
494        # We don't store `False` for `is_metadata` to reduce the pydantic model's payload size
495        is_metadata = is_metadata or None
496
497        if isinstance(v, LITERALS) or v is None:
498            serialized[k] = Executable.value(v, is_metadata=is_metadata)
499        elif inspect.ismodule(v):
500            name = v.__name__
501            if hasattr(v, "__file__") and _is_relative_to(v.__file__, path):
502                raise SQLMeshError(
503                    f"Cannot serialize 'import {name}'. Use 'from {name} import ...' instead."
504                )
505            postfix = "" if name == k else f" as {k}"
506            serialized[k] = Executable(
507                payload=f"import {name}{postfix}",
508                kind=ExecutableKind.IMPORT,
509                is_metadata=is_metadata,
510            )
511        elif callable(v):
512            name = v.__name__
513            name = k if name == "<lambda>" else name
514
515            # getfile raises a `TypeError` for built-in modules, classes, or functions
516            # https://docs.python.org/3/library/inspect.html#inspect.getfile
517            try:
518                file_path = Path(inspect.getfile(v))
519                relative_obj_file_path = _is_relative_to(file_path, path)
520
521                # A callable can be a "wrapper" that is defined in a third-party library [1], in which case the file
522                # containing its definition won't be relative to the project's path. This can lead to serializing
523                # it as a "relative import", such as `from models.some_python_model import foo`, because the `wraps`
524                # decorator preserves the wrapped function's module [2]. Payloads like this are invalid, as they
525                # can result in `ModuleNotFoundError`s when hydrating python environments, e.g. if a project's files
526                # are not available during a scheduled cadence run.
527                #
528                # [1]: https://github.com/jd/tenacity/blob/0d40e76f7d06d631fb127e1ec58c8bd776e70d49/tenacity/__init__.py#L322-L346
529                # [2]: https://github.com/python/cpython/blob/f502c8f6a6db4be27c97a0e5466383d117859b7f/Lib/functools.py#L33-L57
530                if not relative_obj_file_path and (wrapped := getattr(v, "__wrapped__", None)):
531                    v = wrapped
532                    file_path = Path(inspect.getfile(wrapped))
533                    relative_obj_file_path = _is_relative_to(file_path, path)
534            except TypeError:
535                file_path = None
536                relative_obj_file_path = False
537
538            if relative_obj_file_path:
539                serialized[k] = Executable(
540                    name=name,
541                    payload=normalize_source(v),
542                    kind=ExecutableKind.DEFINITION,
543                    # Do `as_posix` to serialize windows path back to POSIX
544                    path=t.cast(Path, file_path).relative_to(path.absolute()).as_posix(),
545                    alias=k if name != k else None,
546                    is_metadata=is_metadata,
547                )
548            else:
549                serialized[k] = Executable(
550                    payload=f"from {_resolve_import_module(v, name)} import {name}",
551                    kind=ExecutableKind.IMPORT,
552                    is_metadata=is_metadata,
553                )
554        else:
555            raise SQLMeshError(
556                f"Object '{v}' cannot be serialized. If it's defined in a library, import the corresponding "
557                "module and reference the object using its fully-qualified name. For example, the datetime "
558                "module's 'UTC' object should be accessed as 'datetime.UTC'."
559            )
560
561    return serialized
562
563
564def prepare_env(
565    python_env: t.Dict[str, Executable],
566    env: t.Optional[t.Dict[str, t.Any]] = None,
567) -> t.Dict[str, t.Any]:
568    """Prepare a python env by hydrating and executing functions.
569
570    The Python ENV is stored in a json serializable format.
571    Functions and imports are stored as a special data class.
572
573    Args:
574        python_env: The dictionary containing the serialized python environment.
575        env: The dictionary to execute code in.
576
577    Returns:
578        The prepared environment with hydrated functions.
579    """
580    env = {} if env is None else env
581
582    for name, executable in sorted(
583        python_env.items(), key=lambda item: 0 if item[1].is_import else 1
584    ):
585        if executable.is_value:
586            env[name] = eval(executable.payload)
587        else:
588            exec(executable.payload, env)
589            if executable.alias and executable.name:
590                env[executable.alias] = env[executable.name]
591
592    return env
593
594
595def format_evaluated_code_exception(
596    exception: Exception,
597    python_env: t.Dict[str, Executable],
598) -> str:
599    """Formats exceptions that occur from evaled code.
600
601    Stack traces generated by evaled code lose code context and are difficult to debug.
602    This intercepts the default stack trace and tries to make it debuggable.
603
604    Args:
605        exception: The exception to print the stack trace for.
606        python_env: The environment containing stringified python code.
607    """
608    tb: t.List[str] = []
609    indent = ""
610
611    skip_patterns = re.compile(
612        r"Traceback \(most recent call last\):|"
613        r'File ".*?core/model/definition\.py|'
614        r'File ".*?core/snapshot/definition\.py|'
615        r'File ".*?core/macros\.py|'
616        r'File ".*?inspect\.py'
617    )
618
619    for error_line in format_exception(exception):
620        if skip_patterns.search(error_line):
621            continue
622
623        error_match = re.search("^.*?Error: ", error_line)
624        if error_match:
625            tb.append(f"{indent * 2}  {error_line}")
626            continue
627
628        eval_code_match = re.search('File "<string>", line (.*), in (.*)', error_line)
629        if not eval_code_match:
630            tb.append(f"{indent}{error_line}")
631            continue
632
633        line_num = int(eval_code_match.group(1))
634        func = eval_code_match.group(2)
635
636        if func not in python_env:
637            tb.append(error_line)
638            continue
639
640        executable = python_env[func]
641        indent = error_line[: eval_code_match.start()]
642
643        error_line = (
644            f"{indent}File '{executable.path}' (or imported file), line {line_num}, in {func}"
645        )
646
647        code = executable.payload
648        formatted = []
649
650        for i, code_line in enumerate(code.splitlines()):
651            if i < line_num:
652                pad = len(code_line) - len(code_line.lstrip())
653                if i + 1 == line_num:
654                    formatted.append(f"{code_line[:pad]}{code_line[pad:]}")
655                else:
656                    formatted.append(code_line)
657
658        tb.extend(
659            (
660                error_line,
661                textwrap.indent(
662                    os.linesep.join(formatted),
663                    indent + "  ",
664                ),
665            )
666        )
667
668    return os.linesep.join(tb)
669
670
671def print_exception(
672    exception: Exception,
673    python_env: t.Dict[str, Executable],
674    out: t.TextIO = sys.stderr,
675) -> None:
676    """Prints exceptions that occur from evaled code.
677
678    Stack traces generated by evaled code lose code context and are difficult to debug.
679    This intercepts the default stack trace and tries to make it debuggable.
680
681    Args:
682        exception: The exception to print the stack trace for.
683        python_env: The environment containing stringified python code.
684        out: The output stream to write to.
685    """
686    tb = format_evaluated_code_exception(exception, python_env)
687    out.write(tb)
688
689
690def _dict_sort(obj: t.Any) -> str:
691    try:
692        if isinstance(obj, dict):
693            obj = dict(sorted(obj.items(), key=lambda x: str(x[0])))
694    except Exception:
695        logger.warning("Failed to sort non-recursive dict", exc_info=True)
696    return repr(obj)
697
698
699def import_python_file(path: Path, relative_base: Path = Path()) -> types.ModuleType:
700    relative_path = path.absolute().relative_to(relative_base.absolute())
701    module_name = str(relative_path.with_suffix("")).replace(os.path.sep, ".")
702
703    # remove the entire module hierarchy in case they were already loaded
704    parts = module_name.split(".")
705    for i in range(len(parts)):
706        sys.modules.pop(".".join(parts[0 : i + 1]), None)
707
708    return importlib.import_module(module_name)