sqlmesh.dbt.seed

  1from __future__ import annotations
  2
  3import typing as t
  4
  5import agate
  6
  7from sqlmesh.dbt.util import DBT_VERSION
  8
  9if DBT_VERSION >= (1, 8, 0):
 10    from dbt_common.clients import agate_helper  # type: ignore
 11
 12    SUPPORTS_DELIMITER = True
 13else:
 14    from dbt.clients import agate_helper  # type: ignore
 15
 16    SUPPORTS_DELIMITER = False
 17from sqlglot import exp
 18
 19from sqlmesh.core.config.common import VirtualEnvironmentMode
 20from sqlmesh.core.model import Model, SeedKind, create_seed_model
 21from sqlmesh.core.model.seed import CsvSettings
 22from sqlmesh.dbt.basemodel import BaseModelConfig
 23from sqlmesh.dbt.column import ColumnConfig
 24
 25if t.TYPE_CHECKING:
 26    from sqlmesh.core.audit.definition import ModelAudit
 27    from sqlmesh.dbt.context import DbtContext
 28
 29
 30class SeedConfig(BaseModelConfig):
 31    """
 32    seedConfig contains all config parameters available to DBT seeds
 33
 34    See https://docs.getdbt.com/reference/configs-and-properties for
 35    a more detailed description of each config parameter under the
 36    General propreties, General configs, and For seeds sections.
 37    """
 38
 39    delimiter: str = ","
 40    column_types: t.Optional[t.Dict[str, str]] = None
 41    quote_columns: t.Optional[bool] = False
 42
 43    def to_sqlmesh(
 44        self,
 45        context: DbtContext,
 46        audit_definitions: t.Optional[t.Dict[str, ModelAudit]] = None,
 47        virtual_environment_mode: VirtualEnvironmentMode = VirtualEnvironmentMode.default,
 48    ) -> Model:
 49        """Converts the dbt seed into a SQLMesh model."""
 50        seed_path = self.path.absolute().as_posix()
 51
 52        column_types_override = {
 53            name: ColumnConfig(name=name, data_type=data_type, quote=self.quote_columns)
 54            for name, data_type in (self.column_types or {}).items()
 55        }
 56        kwargs = self.sqlmesh_model_kwargs(context, column_types_override)
 57
 58        columns = kwargs.get("columns") or {}
 59
 60        agate_table = (
 61            agate_helper.from_csv(seed_path, [], delimiter=self.delimiter)
 62            if SUPPORTS_DELIMITER
 63            else agate_helper.from_csv(seed_path, [])
 64        )
 65        inferred_types = {
 66            name: AGATE_TYPE_MAPPING[tpe.__class__]
 67            for name, tpe in zip(agate_table.column_names, agate_table.column_types)
 68        }
 69
 70        # The columns list built from the mixture of supplied and inferred types needs to
 71        # be in the same order as the data for assumptions elsewhere in the codebase to hold true
 72        new_columns = {}
 73        for column_name in agate_table.column_names:
 74            if column_name not in columns:
 75                new_columns[column_name] = inferred_types[column_name]
 76            else:
 77                new_columns[column_name] = columns[column_name]
 78
 79        kwargs["columns"] = new_columns
 80
 81        # dbt treats single whitespace as a null value
 82        csv_settings = CsvSettings(
 83            delimiter=self.delimiter,
 84            na_values=[" "],
 85            keep_default_na=True,
 86        )
 87
 88        return create_seed_model(
 89            self.canonical_name(context),
 90            SeedKind(path=seed_path, csv_settings=csv_settings),
 91            dialect=self.dialect(context),
 92            audit_definitions=audit_definitions,
 93            virtual_environment_mode=virtual_environment_mode,
 94            start=self.start or context.sqlmesh_config.model_defaults.start,
 95            dbt_node_info=self.node_info,
 96            **kwargs,
 97        )
 98
 99
100AGATE_TYPE_MAPPING = {
101    agate_helper.Number: exp.DataType.build("double"),
102    agate_helper.ISODateTime: exp.DataType.build("datetime"),
103    agate.Date: exp.DataType.build("date"),
104    agate.DateTime: exp.DataType.build("datetime"),
105    agate.Boolean: exp.DataType.build("boolean"),
106    agate.Text: exp.DataType.build("text"),
107}
108
109
110if DBT_VERSION >= (1, 7, 0):
111
112    class Integer(agate_helper.Integer):
113        def cast(self, d: t.Any) -> t.Optional[int]:
114            if isinstance(d, str):
115                # The dbt's implementation doesn't support coercion of strings to integers.
116                if d.strip().lower() in self.null_values:
117                    return None
118                try:
119                    return int(d)
120                except ValueError:
121                    raise agate.exceptions.CastError('Can not parse value "%s" as Integer.' % d)
122            return super().cast(d)
123
124        def jsonify(self, d: t.Any) -> str:
125            return d
126
127    agate_helper.Integer = Integer  # type: ignore
128
129    AGATE_TYPE_MAPPING[agate_helper.Integer] = exp.DataType.build("int")