1# ruff: noqa: E402
  2"""
  3.. include:: ../README.md
  4"""
  5
  6from __future__ import annotations
  7
  8import glob
  9import logging
 10import os
 11import sys
 12import typing as t
 13from datetime import datetime
 14from enum import Enum
 15from pathlib import Path
 16
 17from sqlmesh.core.dialect import extend_sqlglot
 18
 19extend_sqlglot()
 20
 21from sqlmesh.core import constants as c
 22from sqlmesh.core.config import Config as Config
 23from sqlmesh.core.context import Context as Context, ExecutionContext as ExecutionContext
 24from sqlmesh.core.engine_adapter import EngineAdapter as EngineAdapter
 25from sqlmesh.core.macros import SQL as SQL, macro as macro
 26from sqlmesh.core.model import Model as Model, model as model
 27from sqlmesh.core.signal import signal as signal
 28from sqlmesh.core.snapshot import Snapshot as Snapshot
 29from sqlmesh.core.snapshot.evaluator import (
 30    CustomMaterialization as CustomMaterialization,
 31)
 32from sqlmesh.core.model.kind import CustomKind as CustomKind
 33from sqlmesh.utils import (
 34    debug_mode_enabled as debug_mode_enabled,
 35    enable_debug_mode as enable_debug_mode,
 36    str_to_bool,
 37)
 38from sqlmesh.utils.date import DatetimeRanges as DatetimeRanges
 39
 40try:
 41    from sqlmesh._version import __version__ as __version__, __version_tuple__ as __version_tuple__
 42except ImportError:
 43    pass
 44
 45
 46if t.TYPE_CHECKING:
 47    from sqlmesh.core.engine_adapter._typing import QueryOrDF  # noqa: F401
 48
 49
 50class RuntimeEnv(str, Enum):
 51    """Enum defining what environment SQLMesh is running in."""
 52
 53    TERMINAL = "terminal"
 54    DATABRICKS = "databricks"
 55    GOOGLE_COLAB = "google_colab"  # Not currently officially supported
 56    JUPYTER = "jupyter"
 57    DEBUGGER = "debugger"
 58    CI = "ci"  # CI or other envs that shouldn't use emojis
 59
 60    @classmethod
 61    def get(cls) -> RuntimeEnv:
 62        """Get the console class to use based on the environment that the code is running in
 63        Reference implementation: https://github.com/noklam/rich/blob/d3a1ae61a77d934844563514370084971bc3e143/rich/console.py#L511-L528
 64
 65        Unlike the rich implementation we try to split out by notebook type instead of treating it all as Jupyter.
 66        """
 67        runtime_env_var = os.getenv("SQLMESH_RUNTIME_ENVIRONMENT")
 68        if runtime_env_var:
 69            try:
 70                return RuntimeEnv(runtime_env_var)
 71            except ValueError:
 72                valid_values = [f'"{member.value}"' for member in RuntimeEnv]
 73                raise ValueError(
 74                    f"Invalid SQLMESH_RUNTIME_ENVIRONMENT value: {runtime_env_var}. Must be one of {', '.join(valid_values)}."
 75                )
 76
 77        try:
 78            shell = get_ipython()  # type: ignore
 79            if os.getenv("DATABRICKS_RUNTIME_VERSION"):
 80                return RuntimeEnv.DATABRICKS
 81            if "google.colab" in str(shell.__class__):  # type: ignore
 82                return RuntimeEnv.GOOGLE_COLAB
 83            if shell.__class__.__name__ == "ZMQInteractiveShell":  # type: ignore
 84                return RuntimeEnv.JUPYTER
 85        except NameError:
 86            pass
 87
 88        if debug_mode_enabled():
 89            return RuntimeEnv.DEBUGGER
 90
 91        if is_cicd_environment() or not is_interactive_environment():
 92            return RuntimeEnv.CI
 93
 94        return RuntimeEnv.TERMINAL
 95
 96    @property
 97    def is_terminal(self) -> bool:
 98        return self == RuntimeEnv.TERMINAL
 99
100    @property
101    def is_databricks(self) -> bool:
102        return self == RuntimeEnv.DATABRICKS
103
104    @property
105    def is_jupyter(self) -> bool:
106        return self == RuntimeEnv.JUPYTER
107
108    @property
109    def is_google_colab(self) -> bool:
110        return self == RuntimeEnv.GOOGLE_COLAB
111
112    @property
113    def is_ci(self) -> bool:
114        return self == RuntimeEnv.CI
115
116    @property
117    def is_notebook(self) -> bool:
118        return not self.is_terminal and not self.is_ci
119
120
121def is_cicd_environment() -> bool:
122    for key in ("CI", "GITHUB_ACTIONS", "TRAVIS", "CIRCLECI", "GITLAB_CI", "BUILDKITE"):
123        if str_to_bool(os.environ.get(key, "false")):
124            return True
125    return False
126
127
128def is_interactive_environment() -> bool:
129    if sys.stdin is None or sys.stdout is None:
130        return False
131    return sys.stdin.isatty() and sys.stdout.isatty()
132
133
134if RuntimeEnv.get().is_notebook:
135    try:
136        from sqlmesh.magics import register_magics
137
138        register_magics()
139    except ImportError:
140        pass
141
142
143LOG_FORMAT = "%(asctime)s - %(threadName)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)"
144LOG_FILENAME_PREFIX = "sqlmesh_"
145
146
147# SO: https://stackoverflow.com/questions/384076/how-can-i-color-python-logging-output
148class CustomFormatter(logging.Formatter):
149    """Custom logging formatter."""
150
151    grey = "\x1b[38;20m"
152    yellow = "\x1b[33;20m"
153    red = "\x1b[31;20m"
154    bold_red = "\x1b[31;1m"
155    reset = "\x1b[0m"
156
157    FORMATS = {
158        logging.DEBUG: grey + LOG_FORMAT + reset,
159        logging.INFO: grey + LOG_FORMAT + reset,
160        logging.WARNING: yellow + LOG_FORMAT + reset,
161        logging.ERROR: red + LOG_FORMAT + reset,
162        logging.CRITICAL: bold_red + LOG_FORMAT + reset,
163    }
164
165    def format(self, record: logging.LogRecord) -> str:
166        log_fmt = self.FORMATS.get(record.levelno)
167        formatter = logging.Formatter(log_fmt)
168        return formatter.format(record)
169
170
171def remove_excess_logs(
172    log_file_dir: t.Optional[t.Union[str, Path]] = None,
173    log_limit: int = c.DEFAULT_LOG_LIMIT,
174) -> None:
175    if log_limit <= 0:
176        return
177
178    log_file_dir = log_file_dir or c.DEFAULT_LOG_FILE_DIR
179    log_path_prefix = Path(log_file_dir) / LOG_FILENAME_PREFIX
180
181    for path in list(sorted(glob.glob(f"{log_path_prefix}*.log"), reverse=True))[log_limit:]:
182        os.remove(path)
183
184
185def configure_logging(
186    force_debug: bool = False,
187    write_to_stdout: bool = False,
188    write_to_file: bool = True,
189    log_file_dir: t.Optional[t.Union[str, Path]] = None,
190    ignore_warnings: bool = False,
191    log_level: t.Optional[t.Union[str, int]] = None,
192) -> None:
193    # Remove noisy grpc logs that are not useful for users
194    os.environ["GRPC_VERBOSITY"] = os.environ.get("GRPC_VERBOSITY", "NONE")
195
196    logger = logging.getLogger()
197    debug = force_debug or debug_mode_enabled()
198
199    if log_level is not None:
200        if isinstance(log_level, str):
201            level = logging._nameToLevel.get(log_level.upper()) or logging.INFO
202        else:
203            level = log_level
204    else:
205        # base logger needs to be the lowest level that we plan to log
206        level = logging.DEBUG if debug else logging.INFO
207
208    logger.setLevel(level)
209
210    if debug:
211        # Remove noisy snowflake connector logs that are not useful for users
212        logging.getLogger("snowflake.connector").setLevel(logging.INFO)
213
214    if write_to_stdout:
215        stdout_handler = logging.StreamHandler(sys.stdout)
216        stdout_handler.setFormatter(CustomFormatter())
217        stdout_handler.setLevel(logging.ERROR if ignore_warnings else level)
218        logger.addHandler(stdout_handler)
219
220    log_file_dir = log_file_dir or c.DEFAULT_LOG_FILE_DIR
221    log_path_prefix = Path(log_file_dir) / LOG_FILENAME_PREFIX
222
223    if write_to_file:
224        os.makedirs(str(log_file_dir), exist_ok=True)
225        filename = f"{log_path_prefix}{datetime.now().strftime('%Y_%m_%d_%H_%M_%S')}.log"
226        file_handler = logging.FileHandler(filename, mode="w", encoding="utf-8")
227
228        # the log files should always log at least info so that users will always have
229        # minimal info for debugging even if they specify "ignore_warnings"
230        file_handler.setLevel(level)
231        file_handler.setFormatter(logging.Formatter(LOG_FORMAT))
232        logger.addHandler(file_handler)
233
234    if debug:
235        import faulthandler
236
237        enable_debug_mode()
238
239        # Enable threadumps.
240        faulthandler.enable()
241
242        # Windows doesn't support register so we check for it here
243        if hasattr(faulthandler, "register"):
244            from signal import SIGUSR1
245
246            faulthandler.register(SIGUSR1.value)