sqlglot.schema

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