sqlmesh.core.model.decorator

  1from __future__ import annotations
  2
  3import typing as t
  4from pathlib import Path
  5import inspect
  6import re
  7
  8from sqlglot import exp
  9from sqlglot.dialects.dialect import DialectType
 10
 11from sqlmesh.core.config.common import VirtualEnvironmentMode
 12from sqlmesh.core.macros import MacroRegistry
 13from sqlmesh.core.signal import SignalRegistry
 14from sqlmesh.utils.jinja import JinjaMacroRegistry
 15from sqlmesh.core import constants as c
 16from sqlmesh.core.dialect import MacroFunc, parse_one
 17from sqlmesh.core.model.definition import (
 18    Model,
 19    create_python_model,
 20    create_sql_model,
 21    create_models_from_blueprints,
 22    get_model_name,
 23    parse_defaults_properties,
 24    render_meta_fields,
 25    render_model_defaults,
 26)
 27from sqlmesh.core.model.kind import ModelKindName, _ModelKind
 28from sqlmesh.utils import registry_decorator, DECORATOR_RETURN_TYPE
 29from sqlmesh.utils.errors import ConfigError, raise_config_error
 30from sqlmesh.utils.metaprogramming import build_env, serialize_env
 31
 32
 33if t.TYPE_CHECKING:
 34    from sqlmesh.core.audit import ModelAudit
 35
 36
 37class model(registry_decorator):
 38    """Specifies a function is a python based model."""
 39
 40    registry_name = "python_models"
 41    _dialect: DialectType = None
 42
 43    def __init__(self, name: t.Optional[str] = None, is_sql: bool = False, **kwargs: t.Any) -> None:
 44        if not is_sql and "columns" not in kwargs:
 45            raise ConfigError("Python model must define column schema.")
 46
 47        self.name_provided = bool(name)
 48        self.name = name or ""
 49        self.is_sql = is_sql
 50        self.kwargs = kwargs
 51
 52        # Make sure that argument values are expressions in order to pass validation in ModelMeta.
 53        for function_call_attribute in ("audits", "signals"):
 54            calls = self.kwargs.pop(function_call_attribute, [])
 55            self.kwargs[function_call_attribute] = [
 56                (
 57                    (call, {})
 58                    if isinstance(call, str)
 59                    else (
 60                        call[0],
 61                        {
 62                            arg_key: exp.convert(
 63                                tuple(arg_value) if isinstance(arg_value, list) else arg_value
 64                            )
 65                            for arg_key, arg_value in call[1].items()
 66                        },
 67                    )
 68                )
 69                for call in calls
 70            ]
 71
 72        if "default_catalog" in kwargs:
 73            raise ConfigError("`default_catalog` cannot be set on a per-model basis.")
 74
 75        self.columns = {
 76            column_name: (
 77                column_type
 78                if isinstance(column_type, exp.DataType)
 79                else exp.DataType.build(
 80                    str(column_type), dialect=self.kwargs.get("dialect", self._dialect)
 81                )
 82            )
 83            for column_name, column_type in self.kwargs.pop("columns", {}).items()
 84        }
 85
 86    def __call__(
 87        self, func: t.Callable[..., DECORATOR_RETURN_TYPE]
 88    ) -> t.Callable[..., DECORATOR_RETURN_TYPE]:
 89        if not self.name_provided:
 90            self.name = get_model_name(Path(inspect.getfile(func)))
 91        return super().__call__(func)
 92
 93    def models(
 94        self,
 95        get_variables: t.Callable[[t.Optional[str]], t.Dict[str, str]],
 96        path: Path,
 97        module_path: Path,
 98        dialect: t.Optional[str] = None,
 99        default_catalog_per_gateway: t.Optional[t.Dict[str, str]] = None,
100        **loader_kwargs: t.Any,
101    ) -> t.List[Model]:
102        blueprints = self.kwargs.pop("blueprints", None)
103
104        if isinstance(blueprints, str):
105            blueprints = parse_one(blueprints, dialect=dialect)
106
107        if isinstance(blueprints, MacroFunc):
108            from sqlmesh.core.model.definition import render_expression
109
110            blueprints = render_expression(
111                expression=blueprints,
112                module_path=module_path,
113                macros=loader_kwargs.get("macros"),
114                jinja_macros=loader_kwargs.get("jinja_macros"),
115                variables=get_variables(None),
116                path=path,
117                dialect=dialect,
118                default_catalog=loader_kwargs.get("default_catalog"),
119            )
120            if not blueprints:
121                raise_config_error("Failed to render blueprints property", path)
122
123            if len(blueprints) > 1:
124                blueprints = [exp.Tuple(expressions=blueprints)]
125
126            blueprints = blueprints[0]
127
128        return create_models_from_blueprints(
129            gateway=self.kwargs.get("gateway"),
130            blueprints=blueprints,
131            get_variables=get_variables,
132            loader=self.model,
133            path=path,
134            module_path=module_path,
135            dialect=dialect,
136            default_catalog_per_gateway=default_catalog_per_gateway,
137            **loader_kwargs,
138        )
139
140    def model(
141        self,
142        *,
143        module_path: Path,
144        path: Path,
145        defaults: t.Optional[t.Dict[str, t.Any]] = None,
146        macros: t.Optional[MacroRegistry] = None,
147        jinja_macros: t.Optional[JinjaMacroRegistry] = None,
148        signal_definitions: t.Optional[SignalRegistry] = None,
149        audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
150        dialect: t.Optional[str] = None,
151        time_column_format: str = c.DEFAULT_TIME_COLUMN_FORMAT,
152        physical_schema_mapping: t.Optional[t.Dict[re.Pattern, str]] = None,
153        project: str = "",
154        default_catalog: t.Optional[str] = None,
155        variables: t.Optional[t.Dict[str, t.Any]] = None,
156        infer_names: t.Optional[bool] = False,
157        blueprint_variables: t.Optional[t.Dict[str, t.Any]] = None,
158        virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
159    ) -> Model:
160        """Get the model registered by this function."""
161        env: t.Dict[str, t.Tuple[t.Any, t.Optional[bool]]] = {}
162        entrypoint = self.func.__name__
163
164        if not self.name_provided and not infer_names:
165            raise ConfigError("Python model must have a name.")
166
167        kind = self.kwargs.get("kind", None)
168        if kind is not None:
169            if isinstance(kind, _ModelKind):
170                from sqlmesh.core.console import get_console
171
172                get_console().log_warning(
173                    f"""Python model "{self.name}"'s `kind` argument was passed a SQLMesh `{type(kind).__name__}` object. This may result in unexpected behavior - provide a dictionary instead."""
174                )
175            elif isinstance(kind, dict):
176                if "name" not in kind or not isinstance(kind.get("name"), ModelKindName):
177                    raise ConfigError(
178                        f"""Python model "{self.name}"'s `kind` dictionary must contain a `name` key with a valid ModelKindName enum value."""
179                    )
180
181        build_env(self.func, env=env, name=entrypoint, path=module_path)
182
183        rendered_fields = render_meta_fields(
184            fields={"name": self.name, **self.kwargs},
185            module_path=module_path,
186            macros=macros,
187            jinja_macros=jinja_macros,
188            variables=variables,
189            path=path,
190            dialect=dialect,
191            default_catalog=default_catalog,
192            blueprint_variables=blueprint_variables,
193        )
194
195        rendered_name = rendered_fields["name"]
196        if isinstance(rendered_name, exp.Expr):
197            rendered_fields["name"] = rendered_name.sql(dialect=dialect)
198
199        rendered_defaults = (
200            render_model_defaults(
201                defaults=defaults,
202                module_path=module_path,
203                macros=macros,
204                jinja_macros=jinja_macros,
205                variables=variables,
206                path=path,
207                dialect=dialect,
208                default_catalog=default_catalog,
209            )
210            if defaults
211            else {}
212        )
213
214        rendered_defaults = parse_defaults_properties(rendered_defaults, dialect=dialect)
215
216        common_kwargs = {
217            "defaults": rendered_defaults,
218            "path": path,
219            "time_column_format": time_column_format,
220            "python_env": serialize_env(env, path=module_path),
221            "physical_schema_mapping": physical_schema_mapping,
222            "project": project,
223            "default_catalog": default_catalog,
224            "variables": variables,
225            "dialect": dialect,
226            "columns": self.columns if self.columns else None,
227            "module_path": module_path,
228            "macros": macros,
229            "jinja_macros": jinja_macros,
230            "audit_definitions": audit_definitions,
231            "signal_definitions": signal_definitions,
232            "blueprint_variables": blueprint_variables,
233            "virtual_environment_mode": virtual_environment_mode,
234            **rendered_fields,
235        }
236
237        for key in ("pre_statements", "post_statements", "on_virtual_update"):
238            statements = common_kwargs.get(key)
239            if statements:
240                common_kwargs[key] = [
241                    parse_one(s, dialect=common_kwargs.get("dialect")) if isinstance(s, str) else s
242                    for s in statements
243                ]
244
245        if self.is_sql:
246            query = MacroFunc(this=exp.Anonymous(this=entrypoint))
247            return create_sql_model(query=query, **common_kwargs)
248        return create_python_model(entrypoint=entrypoint, **common_kwargs)