sqlmesh.core.model.seed

  1from __future__ import annotations
  2
  3import logging
  4import typing as t
  5import zlib
  6from io import StringIO
  7from pathlib import Path
  8
  9from sqlglot import exp
 10from sqlglot.dialects.dialect import UNESCAPED_SEQUENCES
 11from sqlglot.helper import seq_get
 12from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 13
 14from sqlmesh.core.model.common import parse_bool
 15from sqlmesh.utils.pandas import columns_to_types_from_df
 16from sqlmesh.utils.pydantic import PydanticModel, field_validator
 17
 18if t.TYPE_CHECKING:
 19    import pandas as pd
 20
 21logger = logging.getLogger(__name__)
 22
 23NaHashables = t.List[t.Union[int, str, bool, t.Literal[None]]]
 24NaValues = t.Union[NaHashables, t.Dict[str, NaHashables]]
 25
 26
 27class CsvSettings(PydanticModel):
 28    """Settings for CSV seeds."""
 29
 30    delimiter: t.Optional[str] = None
 31    quotechar: t.Optional[str] = None
 32    doublequote: t.Optional[bool] = None
 33    escapechar: t.Optional[str] = None
 34    skipinitialspace: t.Optional[bool] = None
 35    lineterminator: t.Optional[str] = None
 36    encoding: t.Optional[str] = None
 37    na_values: t.Optional[NaValues] = None
 38    keep_default_na: t.Optional[bool] = None
 39
 40    @field_validator("doublequote", "skipinitialspace", "keep_default_na", mode="before")
 41    @classmethod
 42    def _bool_validator(cls, v: t.Any) -> t.Optional[bool]:
 43        if v is None:
 44            return v
 45        return parse_bool(v)
 46
 47    @field_validator(
 48        "delimiter", "quotechar", "escapechar", "lineterminator", "encoding", mode="before"
 49    )
 50    @classmethod
 51    def _str_validator(cls, v: t.Any) -> t.Optional[str]:
 52        if v is None or not isinstance(v, exp.Expr):
 53            return v
 54
 55        # SQLGlot parses escape sequences like \t as \\t for dialects that don't treat \ as
 56        # an escape character, so we map them back to the corresponding escaped sequence
 57        v = v.this
 58        return UNESCAPED_SEQUENCES.get(v, v)
 59
 60    @field_validator("na_values", mode="before")
 61    @classmethod
 62    def _na_values_validator(cls, v: t.Any) -> t.Optional[NaValues]:
 63        if v is None or not isinstance(v, exp.Expr):
 64            return v
 65
 66        try:
 67            if isinstance(v, exp.Paren) or not isinstance(v, (exp.Tuple, exp.Array)):
 68                v = exp.Tuple(expressions=[v.unnest()])
 69
 70            expressions = v.expressions
 71            if isinstance(seq_get(expressions, 0), (exp.PropertyEQ, exp.EQ)):
 72                return {
 73                    e.left.name: [
 74                        rhs_val.to_py()
 75                        for rhs_val in (
 76                            [e.right.unnest()]
 77                            if isinstance(e.right, exp.Paren)
 78                            else e.right.expressions
 79                        )
 80                    ]
 81                    for e in expressions
 82                }
 83
 84            return [e.to_py() for e in expressions]
 85        except ValueError as e:
 86            logger.warning(f"Failed to coerce na_values '{v}', proceeding with defaults. {str(e)}")
 87
 88        return None
 89
 90
 91class CsvSeedReader:
 92    def __init__(self, content: str, dialect: str, settings: CsvSettings):
 93        self.content = content
 94        self.dialect = dialect
 95        self.settings = settings
 96        self._df: t.Optional[pd.DataFrame] = None
 97
 98    @property
 99    def columns_to_types(self) -> t.Dict[str, exp.DataType]:
100        return columns_to_types_from_df(self._get_df())
101
102    @property
103    def column_hashes(self) -> t.Dict[str, str]:
104        df = self._get_df()
105        return {
106            column_name: str(zlib.crc32(df[column_name].to_json().encode("utf-8")))
107            for column_name in df.columns
108        }
109
110    def read(self, batch_size: t.Optional[int] = None) -> t.Generator[pd.DataFrame, None, None]:
111        df = self._get_df()
112
113        batch_size = batch_size or df.size
114        batch_start = 0
115        while batch_start < df.shape[0]:
116            yield df.iloc[batch_start : batch_start + batch_size, :]
117            batch_start += batch_size
118
119    def _get_df(self) -> pd.DataFrame:
120        import pandas as pd
121
122        if self._df is None:
123            self._df = pd.read_csv(
124                StringIO(self.content),
125                index_col=False,
126                on_bad_lines="error",
127                low_memory=False,
128                **{k: v for k, v in self.settings.dict().items() if v is not None},
129            )
130            self._df = self._df.rename(
131                columns={
132                    col: normalize_identifiers(col, dialect=self.dialect).name
133                    for col in self._df.columns
134                },
135            )
136
137        return self._df
138
139
140class Seed(PydanticModel):
141    """Represents content of a seed.
142
143    Presently only CSV format is supported.
144    """
145
146    content: str
147
148    def reader(self, dialect: str = "", settings: t.Optional[CsvSettings] = None) -> CsvSeedReader:
149        return CsvSeedReader(self.content, dialect, settings or CsvSettings())
150
151
152def create_seed(path: str | Path) -> Seed:
153    with open(Path(path), "r", encoding="utf-8") as fd:
154        return Seed(content=fd.read())