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