Chore: Make the debug mode easier to use (#1501)

Author: izeigerman

sqlmesh/__init__.py

@@ -7,6 +7,7 @@
 import os
 import sys
 import typing as t
+from datetime import datetime
 from enum import Enum
 
 from sqlmesh.core.dialect import extend_sqlglot
@@ -88,6 +89,9 @@ def is_notebook(self) -> bool:
         pass
 
 
+LOG_FORMAT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
+
+
 # SO: https://stackoverflow.com/questions/384076/how-can-i-color-python-logging-output
 class CustomFormatter(logging.Formatter):
     """Custom logging formatter."""
@@ -97,14 +101,13 @@ class CustomFormatter(logging.Formatter):
     red = "\x1b[31;20m"
     bold_red = "\x1b[31;1m"
     reset = "\x1b[0m"
-    log_format = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
 
     FORMATS = {
-        logging.DEBUG: grey + log_format + reset,
-        logging.INFO: grey + log_format + reset,
-        logging.WARNING: yellow + log_format + reset,
-        logging.ERROR: red + log_format + reset,
-        logging.CRITICAL: bold_red + log_format + reset,
+        logging.DEBUG: grey + LOG_FORMAT + reset,
+        logging.INFO: grey + LOG_FORMAT + reset,
+        logging.WARNING: yellow + LOG_FORMAT + reset,
+        logging.ERROR: red + LOG_FORMAT + reset,
+        logging.CRITICAL: bold_red + LOG_FORMAT + reset,
     }
 
     def format(self, record: logging.LogRecord) -> str:
@@ -113,7 +116,7 @@ def format(self, record: logging.LogRecord) -> str:
         return formatter.format(record)
 
 
-def enable_logging(level: t.Optional[int] = None) -> None:
+def enable_logging(level: t.Optional[int] = None, write_to_file: bool = False) -> None:
     """Enable logging to send to stdout and color different levels"""
     level = level or (logging.DEBUG if debug_mode_enabled() else logging.INFO)
     logger = logging.getLogger()
@@ -123,3 +126,10 @@ def enable_logging(level: t.Optional[int] = None) -> None:
         handler.setLevel(level)
         handler.setFormatter(CustomFormatter())
         logger.addHandler(handler)
+
+        if write_to_file:
+            filename = f"sqlmesh_{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log"
+            file_handler = logging.FileHandler(filename, mode="w", encoding="utf-8")
+            file_handler.setLevel(level)
+            file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
+            logger.addHandler(file_handler)

sqlmesh/cli/__init__.py

@@ -1,3 +1,4 @@
+import logging
 import typing as t
 from functools import wraps
 
@@ -11,17 +12,35 @@
 DECORATOR_RETURN_TYPE = t.TypeVar("DECORATOR_RETURN_TYPE")
 
 
+logger = logging.getLogger(__name__)
+
+
 def error_handler(
     func: t.Callable[..., DECORATOR_RETURN_TYPE]
 ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
     @wraps(func)
     def wrapper(*args: t.Any, **kwargs: t.Any) -> DECORATOR_RETURN_TYPE:
-        try:
-            return func(*args, **kwargs)
-        except NodeExecutionFailedError as ex:
-            cause = ex.__cause__
-            raise click.ClickException(f"Failed processing {ex.node}. {cause}")
-        except (SQLMeshError, SqlglotError, ValueError) as ex:
-            raise click.ClickException(str(ex))
-
-    return wrapper if not debug_mode_enabled() else func
+        handler = _debug_exception_handler if debug_mode_enabled() else _default_exception_handler
+        return handler(lambda: func(*args, **kwargs))
+
+    return wrapper
+
+
+def _default_exception_handler(
+    func: t.Callable[[], DECORATOR_RETURN_TYPE]
+) -> DECORATOR_RETURN_TYPE:
+    try:
+        return func()
+    except NodeExecutionFailedError as ex:
+        cause = ex.__cause__
+        raise click.ClickException(f"Failed processing {ex.node}. {cause}")
+    except (SQLMeshError, SqlglotError, ValueError) as ex:
+        raise click.ClickException(str(ex))
+
+
+def _debug_exception_handler(func: t.Callable[[], DECORATOR_RETURN_TYPE]) -> DECORATOR_RETURN_TYPE:
+    try:
+        return func()
+    except Exception:
+        logger.exception("Unhandled exception")
+        raise

sqlmesh/cli/main.py

@@ -12,10 +12,12 @@
 from sqlmesh.cli import options as opt
 from sqlmesh.cli.example_project import ProjectTemplate, init_example_project
 from sqlmesh.core.context import Context
-from sqlmesh.utils import debug_mode_enabled
+from sqlmesh.utils import debug_mode_enabled, enable_debug_mode
 from sqlmesh.utils.date import TimeLike
 from sqlmesh.utils.errors import MissingDependencyError
 
+logger = logging.getLogger(__name__)
+
 
 def _sqlmesh_version() -> str:
     try:
@@ -40,6 +42,11 @@ def _sqlmesh_version() -> str:
     is_flag=True,
     help="Ignore warnings.",
 )
+@click.option(
+    "--debug",
+    is_flag=True,
+    help="Enable debug mode.",
+)
 @click.pass_context
 @error_handler
 def cli(
@@ -48,6 +55,7 @@ def cli(
     config: t.Optional[str] = None,
     gateway: t.Optional[str] = None,
     ignore_warnings: bool = False,
+    debug: bool = False,
 ) -> None:
     """SQLMesh command line tool."""
     if ctx.invoked_subcommand == "version":
@@ -67,25 +75,33 @@ def cli(
     if "--help" in sys.argv:
         return
 
-    if debug_mode_enabled():
+    debug = debug or debug_mode_enabled()
+    if debug:
         import faulthandler
         import signal
 
+        enable_debug_mode()
+
         # Enable threadumps.
         faulthandler.enable()
         # Windows doesn't support register so we check for it here
         if hasattr(faulthandler, "register"):
             faulthandler.register(signal.SIGUSR1.value)
-        enable_logging(level=logging.DEBUG)
+        enable_logging(level=logging.DEBUG, write_to_file=True)
     elif ignore_warnings:
         logging.getLogger().setLevel(logging.ERROR)
 
-    context = Context(
-        paths=paths,
-        config=config,
-        gateway=gateway,
-        load=load,
-    )
+    try:
+        context = Context(
+            paths=paths,
+            config=config,
+            gateway=gateway,
+            load=load,
+        )
+    except Exception:
+        if debug:
+            logger.exception("Failed to initialize SQLMesh context")
+        raise
 
     if load and not context.models:
         raise click.ClickException(

sqlmesh/utils/__init__.py

@@ -190,8 +190,16 @@ def str_to_bool(s: t.Optional[str]) -> bool:
     return s.lower() in ("true", "1", "t", "y", "yes", "on")
 
 
+_debug_mode_enabled: bool = False
+
+
+def enable_debug_mode() -> None:
+    global _debug_mode_enabled
+    _debug_mode_enabled = True
+
+
 def debug_mode_enabled() -> bool:
-    return str_to_bool(os.environ.get("SQLMESH_DEBUG"))
+    return _debug_mode_enabled or str_to_bool(os.environ.get("SQLMESH_DEBUG"))
 
 
 def ttl_cache(ttl: int = 60, maxsize: int = 128000) -> t.Callable: