sqlmesh.cli

 1import logging
 2import typing as t
 3from functools import wraps
 4
 5import click
 6from sqlglot.errors import SqlglotError
 7from sqlmesh.core.context import Context
 8from sqlmesh.utils import debug_mode_enabled
 9from sqlmesh.utils.errors import SQLMeshError
10
11DECORATOR_RETURN_TYPE = t.TypeVar("DECORATOR_RETURN_TYPE")
12
13
14logger = logging.getLogger(__name__)
15
16
17def error_handler(
18    func: t.Callable[..., DECORATOR_RETURN_TYPE],
19) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
20    @wraps(func)
21    def wrapper(*args: t.List[t.Any], **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
22        context_or_obj = args[0]
23        sqlmesh_context = (
24            context_or_obj.obj if isinstance(context_or_obj, click.Context) else context_or_obj
25        )
26        if not isinstance(sqlmesh_context, Context):
27            sqlmesh_context = None
28        handler = _debug_exception_handler if debug_mode_enabled() else _default_exception_handler
29        return handler(sqlmesh_context, lambda: func(*args, **kwargs))
30
31    return wrapper
32
33
34def _default_exception_handler(
35    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
36) -> DECORATOR_RETURN_TYPE:
37    try:
38        return func()
39    except (SQLMeshError, SqlglotError, ValueError) as ex:
40        click.echo(click.style("Error: " + str(ex), fg="red"))
41        exit(1)
42    finally:
43        if context:
44            context.close()
45
46
47def _debug_exception_handler(
48    context: t.Optional[Context], func: t.Callable[[], DECORATOR_RETURN_TYPE]
49) -> DECORATOR_RETURN_TYPE:
50    try:
51        return func()
52    except Exception:
53        logger.exception("Unhandled exception")
54        raise
55    finally:
56        if context:
57            context.close()