sqlglot.schema

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