sqlmesh.dbt.column

 1from __future__ import annotations
 2
 3import typing as t
 4import logging
 5
 6from sqlglot import exp, parse_one
 7from sqlglot.helper import ensure_list
 8
 9from sqlmesh.dbt.common import GeneralConfig
10from sqlmesh.utils.conversions import ensure_bool
11from sqlmesh.utils.pydantic import field_validator
12
13logger = logging.getLogger(__name__)
14
15
16def yaml_to_columns(
17    yaml: t.Dict[str, ColumnConfig] | t.List[t.Dict[str, ColumnConfig]],
18) -> t.Dict[str, ColumnConfig]:
19    columns = {}
20    mappings: t.List[t.Dict[str, ColumnConfig]] = ensure_list(yaml)
21    for column in mappings:
22        config = ColumnConfig(**column)
23        columns[config.name] = config
24
25    return columns
26
27
28def column_types_to_sqlmesh(
29    columns: t.Dict[str, ColumnConfig], dialect: t.Optional[str] = None
30) -> t.Dict[str, exp.DataType]:
31    """
32    Get the sqlmesh column types
33
34    Returns:
35        A dict of column name to exp.DataType
36    """
37    col_types_to_sqlmesh: t.Dict[str, exp.DataType] = {}
38    for name, column in columns.items():
39        if column.enabled and column.data_type:
40            column_def = parse_one(
41                f"{name} {column.data_type}", into=exp.ColumnDef, dialect=dialect or ""
42            )
43            if column_def.args.get("constraints"):
44                logger.warning(
45                    f"Ignoring unsupported constraints for column '{name}' with definition '{column.data_type}'. Please refer to github.com/SQLMesh/sqlmesh/issues/4717 for more information."
46                )
47            kind = column_def.kind
48            if kind:
49                col_types_to_sqlmesh[name] = kind
50    return col_types_to_sqlmesh
51
52
53def column_descriptions_to_sqlmesh(columns: t.Dict[str, ColumnConfig]) -> t.Dict[str, str]:
54    """
55    Get the sqlmesh column types
56
57    Returns:
58        A dict of column name to description
59    """
60    return {
61        name: column.description
62        for name, column in columns.items()
63        if column.enabled and column.description
64    }
65
66
67class ColumnConfig(GeneralConfig):
68    """
69    Column configuration for a DBT project
70
71    Args:
72        name: Name of the column
73        data_type: The column's data type
74        quote: Boolean flag to use quoting for the column name
75    """
76
77    name: str
78    data_type: t.Optional[str] = None
79    quote: t.Optional[bool] = False
80
81    @field_validator("quote", mode="before")
82    @classmethod
83    def _validate_bool(cls, v: str) -> bool:
84        return ensure_bool(v)