sqlmesh.dbt.common

  1from __future__ import annotations
  2
  3import re
  4import typing as t
  5from dataclasses import dataclass
  6from pathlib import Path
  7
  8from ruamel.yaml.constructor import DuplicateKeyError
  9from sqlglot.helper import ensure_list
 10
 11from sqlmesh.dbt.util import DBT_VERSION
 12from sqlmesh.core.config.base import BaseConfig, UpdateStrategy
 13from sqlmesh.core.config.common import DBT_PROJECT_FILENAME
 14from sqlmesh.utils import AttributeDict
 15from sqlmesh.utils.conversions import ensure_bool, try_str_to_bool
 16from sqlmesh.utils.errors import ConfigError
 17from sqlmesh.utils.jinja import MacroReference
 18from sqlmesh.utils.pydantic import PydanticModel, field_validator
 19from sqlmesh.utils.yaml import load
 20
 21T = t.TypeVar("T", bound="GeneralConfig")
 22
 23PROJECT_FILENAME = DBT_PROJECT_FILENAME
 24RAW_CODE_KEY = "raw_code" if DBT_VERSION >= (1, 3, 0) else "raw_sql"  # type: ignore
 25
 26JINJA_ONLY = {
 27    "adapter",
 28    "api",
 29    "exceptions",
 30    "flags",
 31    "load_result",
 32    "modules",
 33    "run_query",
 34    "statement",
 35    "store_result",
 36    "target",
 37}
 38
 39
 40def load_yaml(source: str | Path) -> t.Dict:
 41    try:
 42        return load(
 43            source, render_jinja=False, allow_duplicate_keys=True, keep_last_duplicate_key=True
 44        )
 45    except DuplicateKeyError as ex:
 46        raise ConfigError(f"{source}: {ex}" if isinstance(source, Path) else f"{ex}")
 47
 48
 49def parse_meta(v: t.Optional[t.Dict[str, t.Any]]) -> t.Dict[str, t.Any]:
 50    if v is None:
 51        return {}
 52    for key, value in v.items():
 53        if isinstance(value, str):
 54            v[key] = try_str_to_bool(value)
 55
 56    return v
 57
 58
 59class SqlStr(str):
 60    pass
 61
 62
 63sql_str_validator = field_validator("sql", mode="before", check_fields=False)(
 64    lambda v: SqlStr(v) if isinstance(v, str) else v
 65)
 66
 67
 68class DbtConfig(BaseConfig, extra="allow", validate_assignment=True, frozen=False):
 69    pass
 70
 71
 72class GeneralConfig(DbtConfig):
 73    """
 74    General DBT configuration properties for models, sources, seeds, columns, etc.
 75
 76    Args:
 77        description: Description of element
 78        tests: Tests for the element
 79        enabled: When false, the element is ignored
 80        docs: Documentation specific configuration
 81        perist_docs: Persist resource descriptions as column and/or relation comments in the database
 82        tags: List of tags that can be used for element grouping
 83        meta: Dictionary of metadata for the element
 84    """
 85
 86    start: t.Optional[str] = None
 87    description: t.Optional[str] = None
 88    enabled: bool = True
 89    docs: t.Dict[str, t.Any] = {"show": True}
 90    persist_docs: t.Dict[str, t.Any] = {}
 91    tags: t.List[str] = []
 92    meta: t.Dict[str, t.Any] = {}
 93
 94    @field_validator("enabled", mode="before")
 95    @classmethod
 96    def _validate_bool(cls, v: str) -> bool:
 97        return ensure_bool(v)
 98
 99    @field_validator("docs", mode="before")
100    @classmethod
101    def _validate_dict(cls, v: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
102        for key, value in v.items():
103            if isinstance(value, str):
104                v[key] = try_str_to_bool(value)
105
106        return v
107
108    @field_validator("persist_docs", mode="before")
109    @classmethod
110    def _validate_persist_docs(cls, v: t.Dict[str, str]) -> t.Dict[str, bool]:
111        return {key: bool(value) for key, value in v.items()}
112
113    @field_validator("tags", mode="before")
114    @classmethod
115    def _validate_list(cls, v: t.Union[str, t.List[str]]) -> t.List[str]:
116        return ensure_list(v)
117
118    @field_validator("meta", mode="before")
119    @classmethod
120    def _validate_meta(cls, v: t.Optional[t.Dict[str, t.Union[str, t.Any]]]) -> t.Dict[str, t.Any]:
121        return parse_meta(v)
122
123    _FIELD_UPDATE_STRATEGY: t.ClassVar[t.Dict[str, UpdateStrategy]] = {
124        **BaseConfig._FIELD_UPDATE_STRATEGY,
125        **{
126            "tests": UpdateStrategy.KEY_UPDATE,
127            "docs": UpdateStrategy.KEY_UPDATE,
128            "persist_docs": UpdateStrategy.KEY_UPDATE,
129            "tags": UpdateStrategy.EXTEND,
130            "meta": UpdateStrategy.KEY_UPDATE,
131        },
132    }
133
134    _SQL_FIELDS: t.ClassVar[t.List[str]] = []
135
136    @property
137    def config_attribute_dict(self) -> AttributeDict[str, t.Any]:
138        return AttributeDict(self.dict(exclude=EXCLUDED_CONFIG_ATTRIBUTE_KEYS))
139
140    def _get_field_value(self, field: str) -> t.Optional[t.Any]:
141        field_val = getattr(self, field, None)
142        return field_val if field_val is not None else self.meta.get(field, None)
143
144    def replace(self, other: T) -> None:
145        """
146        Replace the contents of this instance with the passed in instance.
147
148        Args:
149            other: The instance to apply to this instance
150        """
151        for field in other.fields_set:
152            setattr(self, field, getattr(other, field))
153
154    @property
155    def sqlmesh_config_kwargs(self) -> t.Dict[str, t.Any]:
156        """
157        Create a kwargs dict of all sqlmesh config fields.
158
159        Returns:
160            Kwargs dict of sqlmesh config fields.
161        """
162        kwargs = {}
163        for field in self.sqlmesh_config_fields:
164            field_val = self._get_field_value(field)
165            if field_val is not None:
166                kwargs[field] = field_val
167        return kwargs
168
169    @property
170    def sqlmesh_config_fields(self) -> t.Set[str]:
171        """
172        SQLMesh config fields that can be set in dbt projects.
173
174        Returns:
175            A set of SQLMesh config fields that can be set in dbt projects.
176        """
177        return set()
178
179
180@dataclass
181class ModelAttrs:
182    attrs: t.Set[str]
183    all_attrs: bool = False
184
185
186class Dependencies(PydanticModel):
187    """
188    DBT dependencies for a model, macro, etc.
189
190    Args:
191        macros: The references to macros
192        sources: The "source_name.table_name" for source tables used
193        refs: The table_name for models used
194    """
195
196    macros: t.List[MacroReference] = []
197    sources: t.Set[str] = set()
198    refs: t.Set[str] = set()
199    variables: t.Set[str] = set()
200    model_attrs: ModelAttrs = ModelAttrs(attrs=set())
201
202    has_dynamic_var_names: bool = False
203
204    def union(self, other: Dependencies) -> Dependencies:
205        return Dependencies(
206            macros=list(set(self.macros) | set(other.macros)),
207            sources=self.sources | other.sources,
208            refs=self.refs | other.refs,
209            variables=self.variables | other.variables,
210            model_attrs=ModelAttrs(
211                attrs=self.model_attrs.attrs | other.model_attrs.attrs,
212                all_attrs=self.model_attrs.all_attrs or other.model_attrs.all_attrs,
213            ),
214            has_dynamic_var_names=self.has_dynamic_var_names or other.has_dynamic_var_names,
215        )
216
217    @field_validator("macros", mode="after")
218    @classmethod
219    def _sort_macros(cls, v: t.List[MacroReference]) -> t.List[MacroReference]:
220        return sorted(v, key=lambda x: (x.package or "", x.name))
221
222    def dict(self, *args: t.Any, **kwargs: t.Any) -> t.Dict[str, t.Any]:
223        # See https://github.com/pydantic/pydantic/issues/1090
224        exclude = kwargs.pop("exclude", None) or set()
225
226        out = super().dict(*args, **kwargs, exclude={*exclude, "macros"})
227        if "macros" not in exclude:
228            out["macros"] = [macro.dict() for macro in self.macros]
229
230        return out
231
232
233def extract_jinja_config(input: str) -> t.Tuple[str, str]:
234    def jinja_end(sql: str, start: int) -> int:
235        cursor = start
236        quote = None
237        while cursor < len(sql):
238            if sql[cursor] in ('"', "'"):
239                if quote is None:
240                    quote = sql[cursor]
241                elif quote == sql[cursor]:
242                    quote = None
243            if sql[cursor : cursor + 2] == "}}" and quote is None:
244                return cursor + 2
245            cursor += 1
246        return cursor
247
248    no_config = input
249    only_config = ""
250    matches = re.findall(r"{{\s*config\s*\(", no_config)
251    for match in matches:
252        start = no_config.find(match)
253        if start == -1:
254            continue
255        extracted = no_config[start : jinja_end(no_config, start)]
256        only_config = SqlStr("\n".join([only_config, extracted]) if only_config else extracted)
257        no_config = SqlStr(no_config.replace(extracted, "").strip())
258
259    return (no_config, only_config)
260
261
262EXCLUDED_CONFIG_ATTRIBUTE_KEYS = {
263    "config",
264    "config_call_dict",
265    "checksum",
266    "created_at",
267    "contract",
268    "depends_on",
269    "dependencies",
270    "docs",
271    "metrics",
272    "original_file_path",
273    "packages",
274    "patch_path",
275    "path",
276    "persist_docs",
277    "post_hook",
278    "pre_hook",
279    "raw_code",
280    "refs",
281    "resource_type",
282    "sources",
283    "sql",
284    "tests",
285    "unrendered_config",
286}