sqlglot.schema

  1from __future__ import annotations
  2
  3import abc
  4import typing as t
  5
  6import sqlglot
  7from sqlglot import expressions as exp
  8from sqlglot._typing import T
  9from sqlglot.dialects.dialect import Dialect
 10from sqlglot.errors import ParseError, SchemaError
 11from sqlglot.helper import dict_depth
 12from sqlglot.trie import TrieResult, in_trie, new_trie
 13
 14if t.TYPE_CHECKING:
 15    from sqlglot.dataframe.sql.types import StructType
 16    from sqlglot.dialects.dialect import DialectType
 17
 18    ColumnMapping = t.Union[t.Dict, str, StructType, t.List]
 19
 20TABLE_ARGS = ("this", "db", "catalog")
 21
 22
 23class Schema(abc.ABC):
 24    """Abstract base class for database schemas"""
 25
 26    dialect: DialectType
 27
 28    @abc.abstractmethod
 29    def add_table(
 30        self,
 31        table: exp.Table | str,
 32        column_mapping: t.Optional[ColumnMapping] = None,
 33        dialect: DialectType = None,
 34        normalize: t.Optional[bool] = None,
 35        match_depth: bool = True,
 36    ) -> None:
 37        """
 38        Register or update a table. Some implementing classes may require column information to also be provided.
 39        The added table must have the necessary number of qualifiers in its path to match the schema's nesting level.
 40
 41        Args:
 42            table: the `Table` expression instance or string representing the table.
 43            column_mapping: a column mapping that describes the structure of the table.
 44            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 45            normalize: whether to normalize identifiers according to the dialect of interest.
 46            match_depth: whether to enforce that the table must match the schema's depth or not.
 47        """
 48
 49    @abc.abstractmethod
 50    def column_names(
 51        self,
 52        table: exp.Table | str,
 53        only_visible: bool = False,
 54        dialect: DialectType = None,
 55        normalize: t.Optional[bool] = None,
 56    ) -> t.List[str]:
 57        """
 58        Get the column names for a table.
 59
 60        Args:
 61            table: the `Table` expression instance.
 62            only_visible: whether to include invisible columns.
 63            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 64            normalize: whether to normalize identifiers according to the dialect of interest.
 65
 66        Returns:
 67            The list of column names.
 68        """
 69
 70    @abc.abstractmethod
 71    def get_column_type(
 72        self,
 73        table: exp.Table | str,
 74        column: exp.Column,
 75        dialect: DialectType = None,
 76        normalize: t.Optional[bool] = None,
 77    ) -> exp.DataType:
 78        """
 79        Get the `sqlglot.exp.DataType` type of a column in the schema.
 80
 81        Args:
 82            table: the source table.
 83            column: the target column.
 84            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 85            normalize: whether to normalize identifiers according to the dialect of interest.
 86
 87        Returns:
 88            The resulting column type.
 89        """
 90
 91    @property
 92    @abc.abstractmethod
 93    def supported_table_args(self) -> t.Tuple[str, ...]:
 94        """
 95        Table arguments this schema support, e.g. `("this", "db", "catalog")`
 96        """
 97
 98    @property
 99    def empty(self) -> bool:
100        """Returns whether or not the schema is empty."""
101        return True
102
103
104class AbstractMappingSchema(t.Generic[T]):
105    def __init__(
106        self,
107        mapping: t.Optional[t.Dict] = None,
108    ) -> None:
109        self.mapping = mapping or {}
110        self.mapping_trie = new_trie(
111            tuple(reversed(t)) for t in flatten_schema(self.mapping, depth=self.depth())
112        )
113        self._supported_table_args: t.Tuple[str, ...] = tuple()
114
115    @property
116    def empty(self) -> bool:
117        return not self.mapping
118
119    def depth(self) -> int:
120        return dict_depth(self.mapping)
121
122    @property
123    def supported_table_args(self) -> t.Tuple[str, ...]:
124        if not self._supported_table_args and self.mapping:
125            depth = self.depth()
126
127            if not depth:  # None
128                self._supported_table_args = tuple()
129            elif 1 <= depth <= 3:
130                self._supported_table_args = TABLE_ARGS[:depth]
131            else:
132                raise SchemaError(f"Invalid mapping shape. Depth: {depth}")
133
134        return self._supported_table_args
135
136    def table_parts(self, table: exp.Table) -> t.List[str]:
137        if isinstance(table.this, exp.ReadCSV):
138            return [table.this.name]
139        return [table.text(part) for part in TABLE_ARGS if table.text(part)]
140
141    def find(
142        self, table: exp.Table, trie: t.Optional[t.Dict] = None, raise_on_missing: bool = True
143    ) -> t.Optional[T]:
144        parts = self.table_parts(table)[0 : len(self.supported_table_args)]
145        value, trie = in_trie(self.mapping_trie if trie is None else trie, parts)
146
147        if value == TrieResult.FAILED:
148            return None
149
150        if value == TrieResult.PREFIX:
151            possibilities = flatten_schema(trie, depth=dict_depth(trie) - 1)
152
153            if len(possibilities) == 1:
154                parts.extend(possibilities[0])
155            else:
156                message = ", ".join(".".join(parts) for parts in possibilities)
157                if raise_on_missing:
158                    raise SchemaError(f"Ambiguous mapping for {table}: {message}.")
159                return None
160
161        return self.nested_get(parts, raise_on_missing=raise_on_missing)
162
163    def nested_get(
164        self, parts: t.Sequence[str], d: t.Optional[t.Dict] = None, raise_on_missing=True
165    ) -> t.Optional[t.Any]:
166        return nested_get(
167            d or self.mapping,
168            *zip(self.supported_table_args, reversed(parts)),
169            raise_on_missing=raise_on_missing,
170        )
171
172
173class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
174    """
175    Schema based on a nested mapping.
176
177    Args:
178        schema: Mapping in one of the following forms:
179            1. {table: {col: type}}
180            2. {db: {table: {col: type}}}
181            3. {catalog: {db: {table: {col: type}}}}
182            4. None - Tables will be added later
183        visible: Optional mapping of which columns in the schema are visible. If not provided, all columns
184            are assumed to be visible. The nesting should mirror that of the schema:
185            1. {table: set(*cols)}}
186            2. {db: {table: set(*cols)}}}
187            3. {catalog: {db: {table: set(*cols)}}}}
188        dialect: The dialect to be used for custom type mappings & parsing string arguments.
189        normalize: Whether to normalize identifier names according to the given dialect or not.
190    """
191
192    def __init__(
193        self,
194        schema: t.Optional[t.Dict] = None,
195        visible: t.Optional[t.Dict] = None,
196        dialect: DialectType = None,
197        normalize: bool = True,
198    ) -> None:
199        self.dialect = dialect
200        self.visible = visible or {}
201        self.normalize = normalize
202        self._type_mapping_cache: t.Dict[str, exp.DataType] = {}
203        self._depth = 0
204
205        super().__init__(self._normalize(schema or {}))
206
207    @classmethod
208    def from_mapping_schema(cls, mapping_schema: MappingSchema) -> MappingSchema:
209        return MappingSchema(
210            schema=mapping_schema.mapping,
211            visible=mapping_schema.visible,
212            dialect=mapping_schema.dialect,
213            normalize=mapping_schema.normalize,
214        )
215
216    def copy(self, **kwargs) -> MappingSchema:
217        return MappingSchema(
218            **{  # type: ignore
219                "schema": self.mapping.copy(),
220                "visible": self.visible.copy(),
221                "dialect": self.dialect,
222                "normalize": self.normalize,
223                **kwargs,
224            }
225        )
226
227    def add_table(
228        self,
229        table: exp.Table | str,
230        column_mapping: t.Optional[ColumnMapping] = None,
231        dialect: DialectType = None,
232        normalize: t.Optional[bool] = None,
233        match_depth: bool = True,
234    ) -> None:
235        """
236        Register or update a table. Updates are only performed if a new column mapping is provided.
237        The added table must have the necessary number of qualifiers in its path to match the schema's nesting level.
238
239        Args:
240            table: the `Table` expression instance or string representing the table.
241            column_mapping: a column mapping that describes the structure of the table.
242            dialect: the SQL dialect that will be used to parse `table` if it's a string.
243            normalize: whether to normalize identifiers according to the dialect of interest.
244            match_depth: whether to enforce that the table must match the schema's depth or not.
245        """
246        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
247
248        if match_depth and not self.empty and len(normalized_table.parts) != self.depth():
249            raise SchemaError(
250                f"Table {normalized_table.sql(dialect=self.dialect)} must match the "
251                f"schema's nesting level: {self.depth()}."
252            )
253
254        normalized_column_mapping = {
255            self._normalize_name(key, dialect=dialect, normalize=normalize): value
256            for key, value in ensure_column_mapping(column_mapping).items()
257        }
258
259        schema = self.find(normalized_table, raise_on_missing=False)
260        if schema and not normalized_column_mapping:
261            return
262
263        parts = self.table_parts(normalized_table)
264
265        nested_set(self.mapping, tuple(reversed(parts)), normalized_column_mapping)
266        new_trie([parts], self.mapping_trie)
267
268    def column_names(
269        self,
270        table: exp.Table | str,
271        only_visible: bool = False,
272        dialect: DialectType = None,
273        normalize: t.Optional[bool] = None,
274    ) -> t.List[str]:
275        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
276
277        schema = self.find(normalized_table)
278        if schema is None:
279            return []
280
281        if not only_visible or not self.visible:
282            return list(schema)
283
284        visible = self.nested_get(self.table_parts(normalized_table), self.visible) or []
285        return [col for col in schema if col in visible]
286
287    def get_column_type(
288        self,
289        table: exp.Table | str,
290        column: exp.Column,
291        dialect: DialectType = None,
292        normalize: t.Optional[bool] = None,
293    ) -> exp.DataType:
294        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
295
296        normalized_column_name = self._normalize_name(
297            column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
298        )
299
300        table_schema = self.find(normalized_table, raise_on_missing=False)
301        if table_schema:
302            column_type = table_schema.get(normalized_column_name)
303
304            if isinstance(column_type, exp.DataType):
305                return column_type
306            elif isinstance(column_type, str):
307                return self._to_data_type(column_type.upper(), dialect=dialect)
308
309        return exp.DataType.build("unknown")
310
311    def _normalize(self, schema: t.Dict) -> t.Dict:
312        """
313        Normalizes all identifiers in the schema.
314
315        Args:
316            schema: the schema to normalize.
317
318        Returns:
319            The normalized schema mapping.
320        """
321        normalized_mapping: t.Dict = {}
322        flattened_schema = flatten_schema(schema, depth=dict_depth(schema) - 1)
323
324        for keys in flattened_schema:
325            columns = nested_get(schema, *zip(keys, keys))
326
327            if not isinstance(columns, dict):
328                raise SchemaError(
329                    f"Table {'.'.join(keys[:-1])} must match the schema's nesting level: {len(flattened_schema[0])}."
330                )
331
332            normalized_keys = [
333                self._normalize_name(key, dialect=self.dialect, is_table=True) for key in keys
334            ]
335            for column_name, column_type in columns.items():
336                nested_set(
337                    normalized_mapping,
338                    normalized_keys + [self._normalize_name(column_name, dialect=self.dialect)],
339                    column_type,
340                )
341
342        return normalized_mapping
343
344    def _normalize_table(
345        self,
346        table: exp.Table | str,
347        dialect: DialectType = None,
348        normalize: t.Optional[bool] = None,
349    ) -> exp.Table:
350        normalized_table = exp.maybe_parse(
351            table, into=exp.Table, dialect=dialect or self.dialect, copy=True
352        )
353
354        for arg in TABLE_ARGS:
355            value = normalized_table.args.get(arg)
356            if isinstance(value, (str, exp.Identifier)):
357                normalized_table.set(
358                    arg,
359                    exp.to_identifier(
360                        self._normalize_name(
361                            value, dialect=dialect, is_table=True, normalize=normalize
362                        )
363                    ),
364                )
365
366        return normalized_table
367
368    def _normalize_name(
369        self,
370        name: str | exp.Identifier,
371        dialect: DialectType = None,
372        is_table: bool = False,
373        normalize: t.Optional[bool] = None,
374    ) -> str:
375        return normalize_name(
376            name,
377            dialect=dialect or self.dialect,
378            is_table=is_table,
379            normalize=self.normalize if normalize is None else normalize,
380        )
381
382    def depth(self) -> int:
383        if not self.empty and not self._depth:
384            # The columns themselves are a mapping, but we don't want to include those
385            self._depth = super().depth() - 1
386        return self._depth
387
388    def _to_data_type(self, schema_type: str, dialect: DialectType = None) -> exp.DataType:
389        """
390        Convert a type represented as a string to the corresponding `sqlglot.exp.DataType` object.
391
392        Args:
393            schema_type: the type we want to convert.
394            dialect: the SQL dialect that will be used to parse `schema_type`, if needed.
395
396        Returns:
397            The resulting expression type.
398        """
399        if schema_type not in self._type_mapping_cache:
400            dialect = dialect or self.dialect
401
402            try:
403                expression = exp.DataType.build(schema_type, dialect=dialect)
404                self._type_mapping_cache[schema_type] = expression
405            except AttributeError:
406                in_dialect = f" in dialect {dialect}" if dialect else ""
407                raise SchemaError(f"Failed to build type '{schema_type}'{in_dialect}.")
408
409        return self._type_mapping_cache[schema_type]
410
411
412def normalize_name(
413    name: str | exp.Identifier,
414    dialect: DialectType = None,
415    is_table: bool = False,
416    normalize: t.Optional[bool] = True,
417) -> str:
418    try:
419        identifier = sqlglot.maybe_parse(name, dialect=dialect, into=exp.Identifier)
420    except ParseError:
421        return name if isinstance(name, str) else name.name
422
423    name = identifier.name
424    if not normalize:
425        return name
426
427    # This can be useful for normalize_identifier
428    identifier.meta["is_table"] = is_table
429    return Dialect.get_or_raise(dialect).normalize_identifier(identifier).name
430
431
432def ensure_schema(schema: Schema | t.Optional[t.Dict], **kwargs: t.Any) -> Schema:
433    if isinstance(schema, Schema):
434        return schema
435
436    return MappingSchema(schema, **kwargs)
437
438
439def ensure_column_mapping(mapping: t.Optional[ColumnMapping]) -> t.Dict:
440    if mapping is None:
441        return {}
442    elif isinstance(mapping, dict):
443        return mapping
444    elif isinstance(mapping, str):
445        col_name_type_strs = [x.strip() for x in mapping.split(",")]
446        return {
447            name_type_str.split(":")[0].strip(): name_type_str.split(":")[1].strip()
448            for name_type_str in col_name_type_strs
449        }
450    # Check if mapping looks like a DataFrame StructType
451    elif hasattr(mapping, "simpleString"):
452        return {struct_field.name: struct_field.dataType.simpleString() for struct_field in mapping}
453    elif isinstance(mapping, list):
454        return {x.strip(): None for x in mapping}
455
456    raise ValueError(f"Invalid mapping provided: {type(mapping)}")
457
458
459def flatten_schema(
460    schema: t.Dict, depth: int, keys: t.Optional[t.List[str]] = None
461) -> t.List[t.List[str]]:
462    tables = []
463    keys = keys or []
464
465    for k, v in schema.items():
466        if depth >= 2:
467            tables.extend(flatten_schema(v, depth - 1, keys + [k]))
468        elif depth == 1:
469            tables.append(keys + [k])
470
471    return tables
472
473
474def nested_get(
475    d: t.Dict, *path: t.Tuple[str, str], raise_on_missing: bool = True
476) -> t.Optional[t.Any]:
477    """
478    Get a value for a nested dictionary.
479
480    Args:
481        d: the dictionary to search.
482        *path: tuples of (name, key), where:
483            `key` is the key in the dictionary to get.
484            `name` is a string to use in the error if `key` isn't found.
485
486    Returns:
487        The value or None if it doesn't exist.
488    """
489    for name, key in path:
490        d = d.get(key)  # type: ignore
491        if d is None:
492            if raise_on_missing:
493                name = "table" if name == "this" else name
494                raise ValueError(f"Unknown {name}: {key}")
495            return None
496
497    return d
498
499
500def nested_set(d: t.Dict, keys: t.Sequence[str], value: t.Any) -> t.Dict:
501    """
502    In-place set a value for a nested dictionary
503
504    Example:
505        >>> nested_set({}, ["top_key", "second_key"], "value")
506        {'top_key': {'second_key': 'value'}}
507
508        >>> nested_set({"top_key": {"third_key": "third_value"}}, ["top_key", "second_key"], "value")
509        {'top_key': {'third_key': 'third_value', 'second_key': 'value'}}
510
511    Args:
512        d: dictionary to update.
513        keys: the keys that makeup the path to `value`.
514        value: the value to set in the dictionary for the given key path.
515
516    Returns:
517        The (possibly) updated dictionary.
518    """
519    if not keys:
520        return d
521
522    if len(keys) == 1:
523        d[keys[0]] = value
524        return d
525
526    subd = d
527    for key in keys[:-1]:
528        if key not in subd:
529            subd = subd.setdefault(key, {})
530        else:
531            subd = subd[key]
532
533    subd[keys[-1]] = value
534    return d