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, first
 10from sqlglot.trie import TrieResult, in_trie, new_trie
 11
 12
 13if t.TYPE_CHECKING:
 14    from sqlglot._typing import SchemaArgs
 15    from sqlglot.dialects.dialect import DialectType
 16    from collections.abc import Sequence
 17    from typing_extensions import Unpack
 18
 19    ColumnMapping = t.Union[dict[str, t.Any], str, list[str]]
 20
 21
 22class Schema(abc.ABC):
 23    """Abstract base class for database schemas"""
 24
 25    @property
 26    def dialect(self) -> Dialect | None:
 27        """
 28        Returns None by default. Subclasses that require dialect-specific
 29        behavior should override this property.
 30        """
 31        return None
 32
 33    @abc.abstractmethod
 34    def add_table(
 35        self,
 36        table: exp.Table | str,
 37        column_mapping: ColumnMapping | None = None,
 38        dialect: DialectType = None,
 39        normalize: bool | None = None,
 40        match_depth: bool = True,
 41    ) -> None:
 42        """
 43        Register or update a table. Some implementing classes may require column information to also be provided.
 44        The added table must have the necessary number of qualifiers in its path to match the schema's nesting level.
 45
 46        Args:
 47            table: the `Table` expression instance or string representing the table.
 48            column_mapping: a column mapping that describes the structure of the table.
 49            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 50            normalize: whether to normalize identifiers according to the dialect of interest.
 51            match_depth: whether to enforce that the table must match the schema's depth or not.
 52        """
 53
 54    @abc.abstractmethod
 55    def column_names(
 56        self,
 57        table: exp.Table | str,
 58        only_visible: bool = False,
 59        dialect: DialectType = None,
 60        normalize: bool | None = None,
 61    ) -> Sequence[str]:
 62        """
 63        Get the column names for a table.
 64
 65        Args:
 66            table: the `Table` expression instance.
 67            only_visible: whether to include invisible columns.
 68            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 69            normalize: whether to normalize identifiers according to the dialect of interest.
 70
 71        Returns:
 72            The sequence of column names.
 73        """
 74
 75    @abc.abstractmethod
 76    def get_column_type(
 77        self,
 78        table: exp.Table | str,
 79        column: exp.Column | str,
 80        dialect: DialectType = None,
 81        normalize: bool | None = None,
 82    ) -> exp.DataType:
 83        """
 84        Get the `sqlglot.exp.DataType` type of a column in the schema.
 85
 86        Args:
 87            table: the source table.
 88            column: the target column.
 89            dialect: the SQL dialect that will be used to parse `table` if it's a string.
 90            normalize: whether to normalize identifiers according to the dialect of interest.
 91
 92        Returns:
 93            The resulting column type.
 94        """
 95
 96    def has_column(
 97        self,
 98        table: exp.Table | str,
 99        column: exp.Column | str,
100        dialect: DialectType = None,
101        normalize: bool | None = None,
102    ) -> bool:
103        """
104        Returns whether `column` appears in `table`'s schema.
105
106        Args:
107            table: the source table.
108            column: the target column.
109            dialect: the SQL dialect that will be used to parse `table` if it's a string.
110            normalize: whether to normalize identifiers according to the dialect of interest.
111
112        Returns:
113            True if the column appears in the schema, False otherwise.
114        """
115        name = column if isinstance(column, str) else column.name
116        return name in self.column_names(table, dialect=dialect, normalize=normalize)
117
118    def get_udf_type(
119        self,
120        udf: exp.Anonymous | str,
121        dialect: DialectType = None,
122        normalize: bool | None = None,
123    ) -> exp.DataType:
124        """
125        Get the return type of a UDF.
126
127        Args:
128            udf: the UDF expression or string.
129            dialect: the SQL dialect for parsing string arguments.
130            normalize: whether to normalize identifiers.
131
132        Returns:
133            The return type as a DataType, or UNKNOWN if not found.
134        """
135        return exp.DType.UNKNOWN.into_expr()
136
137    @property
138    @abc.abstractmethod
139    def supported_table_args(self) -> tuple[str, ...]:
140        """
141        Table arguments this schema support, e.g. `("this", "db", "catalog")`
142        """
143
144    @property
145    def empty(self) -> bool:
146        """Returns whether the schema is empty."""
147        return True
148
149
150class AbstractMappingSchema:
151    def __init__(
152        self,
153        mapping: dict[str, object] | None = None,
154        udf_mapping: dict[str, object] | None = None,
155    ) -> None:
156        self.mapping: dict[str, object] = mapping or {}
157        self.mapping_trie: dict[str, object] = new_trie(
158            tuple(reversed(t)) for t in flatten_schema(self.mapping, depth=self.depth())
159        )
160        self.udf_mapping: dict[str, object] = udf_mapping or {}
161        self.udf_trie: dict[str, object] = new_trie(
162            tuple(reversed(t)) for t in flatten_schema(self.udf_mapping, depth=self.udf_depth())
163        )
164
165        self._supported_table_args: tuple[str, ...] = tuple()
166
167    @property
168    def empty(self) -> bool:
169        return not self.mapping
170
171    def depth(self) -> int:
172        return dict_depth(self.mapping)
173
174    def udf_depth(self) -> int:
175        return dict_depth(self.udf_mapping)
176
177    @property
178    def supported_table_args(self) -> tuple[str, ...]:
179        if not self._supported_table_args and self.mapping:
180            depth = self.depth()
181
182            if not depth:  # None
183                self._supported_table_args = tuple()
184            elif 1 <= depth <= 3:
185                self._supported_table_args = exp.TABLE_PARTS[:depth]
186            else:
187                raise SchemaError(f"Invalid mapping shape. Depth: {depth}")
188
189        return self._supported_table_args
190
191    def table_parts(self, table: exp.Table) -> list[str]:
192        return [p.name for p in reversed(table.parts)]
193
194    def udf_parts(self, udf: exp.Anonymous) -> list[str]:
195        # a.b.c(...) is represented as Dot(Dot(a, b), Anonymous(c, ...))
196        parent = udf.parent
197        parts = [p.name for p in parent.flatten()] if isinstance(parent, exp.Dot) else [udf.name]
198        return list(reversed(parts))[0 : self.udf_depth()]
199
200    def _find_in_trie(
201        self,
202        parts: list[str],
203        trie: dict[str, object],
204        raise_on_missing: bool,
205    ) -> list[str] | None:
206        value, trie = in_trie(trie, parts)
207
208        if value == TrieResult.FAILED:
209            return None
210
211        if value == TrieResult.PREFIX:
212            possibilities = flatten_schema(trie)
213
214            if len(possibilities) == 1:
215                parts.extend(possibilities[0])
216            else:
217                if raise_on_missing:
218                    joined_parts = ".".join(parts)
219                    message = ", ".join(".".join(p) for p in possibilities)
220                    raise SchemaError(f"Ambiguous mapping for {joined_parts}: {message}.")
221
222                return None
223
224        return parts
225
226    def find(
227        self, table: exp.Table, raise_on_missing: bool = True, ensure_data_types: bool = False
228    ) -> t.Any | None:
229        """
230        Returns the schema of a given table.
231
232        Args:
233            table: the target table.
234            raise_on_missing: whether to raise in case the schema is not found.
235            ensure_data_types: whether to convert `str` types to their `DataType` equivalents.
236
237        Returns:
238            The schema of the target table.
239        """
240        parts = self.table_parts(table)[0 : len(self.supported_table_args)]
241        resolved_parts = self._find_in_trie(parts, self.mapping_trie, raise_on_missing)
242
243        if resolved_parts is None:
244            return None
245
246        return self.nested_get(resolved_parts, raise_on_missing=raise_on_missing)
247
248    def find_udf(self, udf: exp.Anonymous, raise_on_missing: bool = False) -> t.Any | None:
249        """
250        Returns the return type of a given UDF.
251
252        Args:
253            udf: the target UDF expression.
254            raise_on_missing: whether to raise if the UDF is not found.
255
256        Returns:
257            The return type of the UDF, or None if not found.
258        """
259        parts = self.udf_parts(udf)
260        resolved_parts = self._find_in_trie(parts, self.udf_trie, raise_on_missing)
261
262        if resolved_parts is None:
263            return None
264
265        return nested_get(
266            self.udf_mapping,
267            *zip(resolved_parts, reversed(resolved_parts)),
268            raise_on_missing=raise_on_missing,
269        )
270
271    def nested_get(
272        self,
273        parts: Sequence[str],
274        d: dict[str, object] | None = None,
275        raise_on_missing: bool = True,
276    ) -> t.Any | None:
277        return nested_get(
278            d or self.mapping,
279            *zip(self.supported_table_args, reversed(parts)),
280            raise_on_missing=raise_on_missing,
281        )
282
283
284class MappingSchema(AbstractMappingSchema, Schema):
285    """
286    Schema based on a nested mapping.
287
288    Args:
289        schema: Mapping in one of the following forms:
290            1. {table: {col: type}}
291            2. {db: {table: {col: type}}}
292            3. {catalog: {db: {table: {col: type}}}}
293            4. None - Tables will be added later
294        visible: Optional mapping of which columns in the schema are visible. If not provided, all columns
295            are assumed to be visible. The nesting should mirror that of the schema:
296            1. {table: set(*cols)}}
297            2. {db: {table: set(*cols)}}}
298            3. {catalog: {db: {table: set(*cols)}}}}
299        dialect: The dialect to be used for custom type mappings & parsing string arguments.
300        normalize: Whether to normalize identifier names according to the given dialect or not.
301    """
302
303    def __init__(
304        self,
305        schema: dict[str, object] | None = None,
306        visible: dict[str, object] | None = None,
307        dialect: DialectType = None,
308        normalize: bool = True,
309        udf_mapping: dict[str, object] | None = None,
310    ) -> None:
311        self.visible: dict[str, object] = {} if visible is None else visible
312        self.normalize: bool = normalize
313        self._dialect: Dialect = Dialect.get_or_raise(dialect)
314        self._type_mapping_cache: dict[str, exp.DataType] = {}
315        self._normalized_table_cache: dict[tuple[exp.Table, DialectType, bool], exp.Table] = {}
316        self._normalized_name_cache: dict[tuple[str, DialectType, bool, bool], str] = {}
317        self._find_cache: dict[tuple[exp.Table, bool], dict[str, object] | None] = {}
318        self._depth: int = 0
319        schema = {} if schema is None else schema
320        udf_mapping = {} if udf_mapping is None else udf_mapping
321
322        super().__init__(
323            self._normalize(schema) if self.normalize else schema,
324            self._normalize_udfs(udf_mapping) if self.normalize else udf_mapping,
325        )
326
327    @property
328    def dialect(self) -> Dialect:
329        """Returns the dialect for this mapping schema."""
330        return self._dialect
331
332    @classmethod
333    def from_mapping_schema(cls, mapping_schema: MappingSchema) -> MappingSchema:
334        return MappingSchema(
335            schema=mapping_schema.mapping,
336            visible=mapping_schema.visible,
337            dialect=mapping_schema.dialect,
338            normalize=mapping_schema.normalize,
339            udf_mapping=mapping_schema.udf_mapping,
340        )
341
342    def find(
343        self, table: exp.Table, raise_on_missing: bool = True, ensure_data_types: bool = False
344    ) -> t.Any | None:
345        cache_key = (table, ensure_data_types)
346        schema = self._find_cache.get(cache_key)
347
348        if schema is None:
349            schema = super().find(table, raise_on_missing=raise_on_missing)
350            if ensure_data_types and isinstance(schema, dict):
351                schema = {
352                    col: self._to_data_type(dtype) if isinstance(dtype, str) else dtype
353                    for col, dtype in schema.items()
354                }
355            self._find_cache[cache_key] = schema
356
357        return schema
358
359    def copy(
360        self, schema: dict[str, object] | None = None, **kwargs: Unpack[SchemaArgs]
361    ) -> MappingSchema:
362        mapping_kwargs: SchemaArgs = {
363            "visible": self.visible.copy(),
364            "dialect": self.dialect,
365            "normalize": self.normalize,
366            "udf_mapping": self.udf_mapping.copy(),
367            **kwargs,
368        }
369        return MappingSchema(self.mapping.copy() if schema is None else schema, **mapping_kwargs)
370
371    def add_table(
372        self,
373        table: exp.Table | str,
374        column_mapping: ColumnMapping | None = None,
375        dialect: DialectType = None,
376        normalize: bool | None = None,
377        match_depth: bool = True,
378    ) -> None:
379        """
380        Register or update a table. Updates are only performed if a new column mapping is provided.
381        The added table must have the necessary number of qualifiers in its path to match the schema's nesting level.
382
383        Args:
384            table: the `Table` expression instance or string representing the table.
385            column_mapping: a column mapping that describes the structure of the table.
386            dialect: the SQL dialect that will be used to parse `table` if it's a string.
387            normalize: whether to normalize identifiers according to the dialect of interest.
388            match_depth: whether to enforce that the table must match the schema's depth or not.
389        """
390        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
391
392        if match_depth and not self.empty and len(normalized_table.parts) != self.depth():
393            raise SchemaError(
394                f"Table {normalized_table.sql(dialect=self.dialect)} must match the "
395                f"schema's nesting level: {self.depth()}."
396            )
397
398        normalized_column_mapping = {
399            self._normalize_name(key, dialect=dialect, normalize=normalize): value
400            for key, value in ensure_column_mapping(column_mapping).items()
401        }
402
403        schema = self.find(normalized_table, raise_on_missing=False)
404        if schema and not normalized_column_mapping:
405            return
406
407        parts = self.table_parts(normalized_table)
408
409        nested_set(self.mapping, tuple(reversed(parts)), normalized_column_mapping)
410        new_trie([parts], self.mapping_trie)
411        self._find_cache.pop((normalized_table, True), None)
412        self._find_cache.pop((normalized_table, False), None)
413
414    def column_names(
415        self,
416        table: exp.Table | str,
417        only_visible: bool = False,
418        dialect: DialectType = None,
419        normalize: bool | None = None,
420    ) -> list[str]:
421        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
422
423        schema: dict[str, object] | None = self.find(normalized_table)
424        if schema is None:
425            return []
426
427        if not only_visible or not self.visible:
428            return list(schema)
429
430        visible = self.nested_get(self.table_parts(normalized_table), self.visible) or []
431        return [col for col in schema if col in visible]
432
433    def get_column_type(
434        self,
435        table: exp.Table | str,
436        column: exp.Column | str,
437        dialect: DialectType = None,
438        normalize: bool | None = None,
439    ) -> exp.DataType:
440        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
441
442        normalized_column_name = self._normalize_name(
443            column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
444        )
445
446        table_schema: dict[str, object] | None = self.find(normalized_table, raise_on_missing=False)
447        if table_schema:
448            column_type = table_schema.get(normalized_column_name)
449
450            if isinstance(column_type, exp.DataType):
451                return column_type
452            elif isinstance(column_type, str):
453                return self._to_data_type(column_type, dialect=dialect)
454
455        return exp.DType.UNKNOWN.into_expr()
456
457    def get_udf_type(
458        self,
459        udf: exp.Anonymous | str,
460        dialect: DialectType = None,
461        normalize: bool | None = None,
462    ) -> exp.DataType:
463        """
464        Get the return type of a UDF.
465
466        Args:
467            udf: the UDF expression or string (e.g., "db.my_func()").
468            dialect: the SQL dialect for parsing string arguments.
469            normalize: whether to normalize identifiers.
470
471        Returns:
472            The return type as a DataType, or UNKNOWN if not found.
473        """
474        parts = self._normalize_udf(udf, dialect=dialect, normalize=normalize)
475        resolved_parts = self._find_in_trie(parts, self.udf_trie, raise_on_missing=False)
476
477        if resolved_parts is None:
478            return exp.DType.UNKNOWN.into_expr()
479
480        udf_type = nested_get(
481            self.udf_mapping,
482            *zip(resolved_parts, reversed(resolved_parts)),
483            raise_on_missing=False,
484        )
485
486        if isinstance(udf_type, exp.DataType):
487            return udf_type
488        elif isinstance(udf_type, str):
489            return self._to_data_type(udf_type, dialect=dialect)
490
491        return exp.DType.UNKNOWN.into_expr()
492
493    def has_column(
494        self,
495        table: exp.Table | str,
496        column: exp.Column | str,
497        dialect: DialectType = None,
498        normalize: bool | None = None,
499    ) -> bool:
500        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
501
502        normalized_column_name = self._normalize_name(
503            column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
504        )
505
506        table_schema: dict[str, object] | None = self.find(normalized_table, raise_on_missing=False)
507        return normalized_column_name in table_schema if table_schema else False
508
509    def _normalize(self, schema: dict[str, object]) -> dict[str, object]:
510        """
511        Normalizes all identifiers in the schema.
512
513        Args:
514            schema: the schema to normalize.
515
516        Returns:
517            The normalized schema mapping.
518        """
519        normalized_mapping: dict[str, object] = {}
520        flattened_schema = flatten_schema(schema)
521        error_msg = "Table {} must match the schema's nesting level: {}."
522
523        for keys in flattened_schema:
524            columns = nested_get(schema, *zip(keys, keys))
525
526            if not isinstance(columns, dict):
527                raise SchemaError(error_msg.format(".".join(keys[:-1]), len(flattened_schema[0])))
528            if not columns:
529                raise SchemaError(f"Table {'.'.join(keys[:-1])} must have at least one column")
530            if isinstance(first(columns.values()), dict):
531                raise SchemaError(
532                    error_msg.format(
533                        ".".join(keys + flatten_schema(columns)[0]), len(flattened_schema[0])
534                    ),
535                )
536
537            normalized_keys = [self._normalize_name(key, is_table=True) for key in keys]
538            for column_name, column_type in columns.items():
539                nested_set(
540                    normalized_mapping,
541                    normalized_keys + [self._normalize_name(column_name)],
542                    column_type,
543                )
544
545        return normalized_mapping
546
547    def _normalize_udfs(self, udfs: dict[str, object]) -> dict[str, object]:
548        """
549        Normalizes all identifiers in the UDF mapping.
550
551        Args:
552            udfs: the UDF mapping to normalize.
553
554        Returns:
555            The normalized UDF mapping.
556        """
557        normalized_mapping: dict[str, object] = {}
558
559        for keys in flatten_schema(udfs, depth=dict_depth(udfs)):
560            udf_type = nested_get(udfs, *zip(keys, keys))
561            normalized_keys = [self._normalize_name(key, is_table=True) for key in keys]
562            nested_set(normalized_mapping, normalized_keys, udf_type)
563
564        return normalized_mapping
565
566    def _normalize_udf(
567        self,
568        udf: exp.Anonymous | str,
569        dialect: DialectType = None,
570        normalize: bool | None = None,
571    ) -> list[str]:
572        """
573        Extract and normalize UDF parts for lookup.
574
575        Args:
576            udf: the UDF expression or qualified string (e.g., "db.my_func()").
577            dialect: the SQL dialect for parsing.
578            normalize: whether to normalize identifiers.
579
580        Returns:
581            A list of normalized UDF parts (reversed for trie lookup).
582        """
583        dialect = dialect or self.dialect
584        normalize = self.normalize if normalize is None else normalize
585
586        if isinstance(udf, str):
587            parsed: exp.Expr = exp.maybe_parse(udf, dialect=dialect)
588
589            if isinstance(parsed, exp.Anonymous):
590                udf = parsed
591            elif isinstance(parsed, exp.Dot) and isinstance(parsed.expression, exp.Anonymous):
592                udf = parsed.expression
593            else:
594                raise SchemaError(f"Unable to parse UDF from: {udf!r}")
595        parts = self.udf_parts(udf)
596
597        if normalize:
598            parts = [self._normalize_name(part, dialect=dialect, is_table=True) for part in parts]
599
600        return parts
601
602    def _normalize_table(
603        self,
604        table: exp.Table | str,
605        dialect: DialectType = None,
606        normalize: bool | None = None,
607    ) -> exp.Table:
608        dialect = dialect or self.dialect
609        normalize = self.normalize if normalize is None else normalize
610
611        # Cache normalized tables by object id for exp.Table inputs
612        # This is effective when the same Table object is looked up multiple times
613        if isinstance(table, exp.Table) and (
614            cached := self._normalized_table_cache.get((table, dialect, normalize))
615        ):
616            return cached
617
618        normalized_table = exp.maybe_parse(table, into=exp.Table, dialect=dialect, copy=normalize)
619
620        if normalize:
621            for part in normalized_table.parts:
622                if isinstance(part, exp.Identifier):
623                    part.replace(
624                        normalize_name(part, dialect=dialect, is_table=True, normalize=normalize)
625                    )
626
627        self._normalized_table_cache[(normalized_table, dialect, normalize)] = normalized_table
628        return normalized_table
629
630    def _normalize_name(
631        self,
632        name: str | exp.Identifier,
633        dialect: DialectType = None,
634        is_table: bool = False,
635        normalize: bool | None = None,
636    ) -> str:
637        normalize = self.normalize if normalize is None else normalize
638
639        dialect = dialect or self.dialect
640        name_str = name if isinstance(name, str) else name.name
641        cache_key = (name_str, dialect, is_table, normalize)
642
643        if cached := self._normalized_name_cache.get(cache_key):
644            return cached
645
646        result = normalize_name(
647            name,
648            dialect=dialect,
649            is_table=is_table,
650            normalize=normalize,
651        ).name
652
653        self._normalized_name_cache[cache_key] = result
654        return result
655
656    def depth(self) -> int:
657        if not self.empty and not self._depth:
658            # The columns themselves are a mapping, but we don't want to include those
659            self._depth = super().depth() - 1
660        return self._depth
661
662    def _to_data_type(self, schema_type: str, dialect: DialectType = None) -> exp.DataType:
663        """
664        Convert a type represented as a string to the corresponding `sqlglot.exp.DataType` object.
665
666        Args:
667            schema_type: the type we want to convert.
668            dialect: the SQL dialect that will be used to parse `schema_type`, if needed.
669
670        Returns:
671            The resulting expression type.
672        """
673        if schema_type not in self._type_mapping_cache:
674            dialect = Dialect.get_or_raise(dialect) if dialect else self.dialect
675            udt = dialect.SUPPORTS_USER_DEFINED_TYPES
676
677            try:
678                expression = exp.DataType.from_str(schema_type, dialect=dialect, udt=udt)
679                expression.transform(dialect.normalize_identifier, copy=False)
680                self._type_mapping_cache[schema_type] = expression
681            except AttributeError:
682                in_dialect = f" in dialect {dialect}" if dialect else ""
683                raise SchemaError(f"Failed to build type '{schema_type}'{in_dialect}.")
684
685        return self._type_mapping_cache[schema_type]
686
687
688def normalize_name(
689    identifier: str | exp.Identifier,
690    dialect: DialectType = None,
691    is_table: bool = False,
692    normalize: bool | None = True,
693) -> exp.Identifier:
694    if isinstance(identifier, str):
695        identifier = exp.parse_identifier(identifier, dialect=dialect)
696    elif normalize:
697        identifier = identifier.copy()
698
699    if not normalize:
700        return identifier
701
702    # this is used for normalize_identifier, bigquery has special rules pertaining tables
703    identifier.meta["is_table"] = is_table
704    return Dialect.get_or_raise(dialect).normalize_identifier(identifier)
705
706
707def ensure_schema(
708    schema: Schema | dict[str, object] | None, **kwargs: Unpack[SchemaArgs]
709) -> Schema:
710    if isinstance(schema, Schema):
711        return schema
712
713    return MappingSchema(schema, **kwargs)
714
715
716def ensure_column_mapping(mapping: ColumnMapping | None) -> dict[str, t.Any]:
717    if mapping is None:
718        return {}
719    elif isinstance(mapping, dict):
720        return mapping
721    elif isinstance(mapping, str):
722        col_name_type_strs = [x.strip() for x in mapping.split(",")]
723        return {
724            name_type_str.split(":")[0].strip(): name_type_str.split(":")[1].strip()
725            for name_type_str in col_name_type_strs
726        }
727    elif isinstance(mapping, list):
728        return {x.strip(): None for x in mapping}
729
730    raise ValueError(f"Invalid mapping provided: {type(mapping)}")
731
732
733def flatten_schema(
734    schema: dict[str, object], depth: int | None = None, keys: list[str] | None = None
735) -> list[list[str]]:
736    tables: list[list[str]] = []
737    keys = keys or []
738    depth = dict_depth(schema) - 1 if depth is None else depth
739
740    for k, v in schema.items():
741        if depth == 1 or not isinstance(v, dict):
742            tables.append(keys + [k])
743        elif depth >= 2:
744            tables.extend(flatten_schema(v, depth - 1, keys + [k]))
745
746    return tables
747
748
749def nested_get(
750    d: dict[str, object], *path: tuple[str, str], raise_on_missing: bool = True
751) -> t.Any | None:
752    """
753    Get a value for a nested dictionary.
754
755    Args:
756        d: the dictionary to search.
757        *path: tuples of (name, key), where:
758            `key` is the key in the dictionary to get.
759            `name` is a string to use in the error if `key` isn't found.
760
761    Returns:
762        The value or None if it doesn't exist.
763    """
764    result: t.Any = d
765    for name, key in path:
766        result = result.get(key)
767        if result is None:
768            if raise_on_missing:
769                name = "table" if name == "this" else name
770                raise ValueError(f"Unknown {name}: {key}")
771            return None
772
773    return result
774
775
776def nested_set(d: dict[str, t.Any], keys: Sequence[str], value: t.Any) -> dict[str, t.Any]:
777    """
778    In-place set a value for a nested dictionary
779
780    Example:
781        >>> nested_set({}, ["top_key", "second_key"], "value")
782        {'top_key': {'second_key': 'value'}}
783
784        >>> nested_set({"top_key": {"third_key": "third_value"}}, ["top_key", "second_key"], "value")
785        {'top_key': {'third_key': 'third_value', 'second_key': 'value'}}
786
787    Args:
788        d: dictionary to update.
789        keys: the keys that makeup the path to `value`.
790        value: the value to set in the dictionary for the given key path.
791
792    Returns:
793        The (possibly) updated dictionary.
794    """
795    if not keys:
796        return d
797
798    if len(keys) == 1:
799        d[keys[0]] = value
800        return d
801
802    subd = d
803    for key in keys[:-1]:
804        if key not in subd:
805            subd = subd.setdefault(key, {})
806        else:
807            subd = subd[key]
808
809    subd[keys[-1]] = value
810    return d