sqlmesh.cli

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