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
412        # A new table can change how partially-qualified names resolve (or make them ambiguous),
413        # so any cached lookup may be stale, not just the one for this table's full path
414        self._find_cache.clear()
415
416    def column_names(
417        self,
418        table: exp.Table | str,
419        only_visible: bool = False,
420        dialect: DialectType = None,
421        normalize: bool | None = None,
422    ) -> list[str]:
423        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
424
425        schema: dict[str, object] | None = self.find(normalized_table)
426        if schema is None:
427            return []
428
429        if not only_visible or not self.visible:
430            return list(schema)
431
432        visible = self.nested_get(self.table_parts(normalized_table), self.visible) or []
433        return [col for col in schema if col in visible]
434
435    def get_column_type(
436        self,
437        table: exp.Table | str,
438        column: exp.Column | str,
439        dialect: DialectType = None,
440        normalize: bool | None = None,
441    ) -> exp.DataType:
442        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
443
444        normalized_column_name = self._normalize_name(
445            column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
446        )
447
448        table_schema: dict[str, object] | None = self.find(normalized_table, raise_on_missing=False)
449        if table_schema:
450            column_type = table_schema.get(normalized_column_name)
451
452            if isinstance(column_type, exp.DataType):
453                return column_type
454            elif isinstance(column_type, str):
455                return self._to_data_type(column_type, dialect=dialect)
456
457        return exp.DType.UNKNOWN.into_expr()
458
459    def get_udf_type(
460        self,
461        udf: exp.Anonymous | str,
462        dialect: DialectType = None,
463        normalize: bool | None = None,
464    ) -> exp.DataType:
465        """
466        Get the return type of a UDF.
467
468        Args:
469            udf: the UDF expression or string (e.g., "db.my_func()").
470            dialect: the SQL dialect for parsing string arguments.
471            normalize: whether to normalize identifiers.
472
473        Returns:
474            The return type as a DataType, or UNKNOWN if not found.
475        """
476        parts = self._normalize_udf(udf, dialect=dialect, normalize=normalize)
477        resolved_parts = self._find_in_trie(parts, self.udf_trie, raise_on_missing=False)
478
479        if resolved_parts is None:
480            return exp.DType.UNKNOWN.into_expr()
481
482        udf_type = nested_get(
483            self.udf_mapping,
484            *zip(resolved_parts, reversed(resolved_parts)),
485            raise_on_missing=False,
486        )
487
488        if isinstance(udf_type, exp.DataType):
489            return udf_type
490        elif isinstance(udf_type, str):
491            return self._to_data_type(udf_type, dialect=dialect)
492
493        return exp.DType.UNKNOWN.into_expr()
494
495    def has_column(
496        self,
497        table: exp.Table | str,
498        column: exp.Column | str,
499        dialect: DialectType = None,
500        normalize: bool | None = None,
501    ) -> bool:
502        normalized_table = self._normalize_table(table, dialect=dialect, normalize=normalize)
503
504        normalized_column_name = self._normalize_name(
505            column if isinstance(column, str) else column.this, dialect=dialect, normalize=normalize
506        )
507
508        table_schema: dict[str, object] | None = self.find(normalized_table, raise_on_missing=False)
509        return normalized_column_name in table_schema if table_schema else False
510
511    def _normalize(self, schema: dict[str, object]) -> dict[str, object]:
512        """
513        Normalizes all identifiers in the schema.
514
515        Args:
516            schema: the schema to normalize.
517
518        Returns:
519            The normalized schema mapping.
520        """
521        normalized_mapping: dict[str, object] = {}
522        flattened_schema = flatten_schema(schema)
523        error_msg = "Table {} must match the schema's nesting level: {}."
524
525        for keys in flattened_schema:
526            columns = nested_get(schema, *zip(keys, keys))
527
528            if not isinstance(columns, dict):
529                raise SchemaError(error_msg.format(".".join(keys[:-1]), len(flattened_schema[0])))
530            if not columns:
531                raise SchemaError(f"Table {'.'.join(keys[:-1])} must have at least one column")
532            if isinstance(first(columns.values()), dict):
533                raise SchemaError(
534                    error_msg.format(
535                        ".".join(keys + flatten_schema(columns)[0]), len(flattened_schema[0])
536                    ),
537                )
538
539            normalized_keys = [self._normalize_name(key, is_table=True) for key in keys]
540            for column_name, column_type in columns.items():
541                nested_set(
542                    normalized_mapping,
543                    normalized_keys + [self._normalize_name(column_name)],
544                    column_type,
545                )
546
547        return normalized_mapping
548
549    def _normalize_udfs(self, udfs: dict[str, object]) -> dict[str, object]:
550        """
551        Normalizes all identifiers in the UDF mapping.
552
553        Args:
554            udfs: the UDF mapping to normalize.
555
556        Returns:
557            The normalized UDF mapping.
558        """
559        normalized_mapping: dict[str, object] = {}
560
561        for keys in flatten_schema(udfs, depth=dict_depth(udfs)):
562            udf_type = nested_get(udfs, *zip(keys, keys))
563            normalized_keys = [self._normalize_name(key, is_table=True) for key in keys]
564            nested_set(normalized_mapping, normalized_keys, udf_type)
565
566        return normalized_mapping
567
568    def _normalize_udf(
569        self,
570        udf: exp.Anonymous | str,
571        dialect: DialectType = None,
572        normalize: bool | None = None,
573    ) -> list[str]:
574        """
575        Extract and normalize UDF parts for lookup.
576
577        Args:
578            udf: the UDF expression or qualified string (e.g., "db.my_func()").
579            dialect: the SQL dialect for parsing.
580            normalize: whether to normalize identifiers.
581
582        Returns:
583            A list of normalized UDF parts (reversed for trie lookup).
584        """
585        dialect = dialect or self.dialect
586        normalize = self.normalize if normalize is None else normalize
587
588        if isinstance(udf, str):
589            parsed: exp.Expr = exp.maybe_parse(udf, dialect=dialect)
590
591            if isinstance(parsed, exp.Anonymous):
592                udf = parsed
593            elif isinstance(parsed, exp.Dot) and isinstance(parsed.expression, exp.Anonymous):
594                udf = parsed.expression
595            else:
596                raise SchemaError(f"Unable to parse UDF from: {udf!r}")
597        parts = self.udf_parts(udf)
598
599        if normalize:
600            parts = [self._normalize_name(part, dialect=dialect, is_table=True) for part in parts]
601
602        return parts
603
604    def _normalize_table(
605        self,
606        table: exp.Table | str,
607        dialect: DialectType = None,
608        normalize: bool | None = None,
609    ) -> exp.Table:
610        dialect = dialect or self.dialect
611        normalize = self.normalize if normalize is None else normalize
612
613        # Cache normalized tables by object id for exp.Table inputs
614        # This is effective when the same Table object is looked up multiple times
615        if isinstance(table, exp.Table) and (
616            cached := self._normalized_table_cache.get((table, dialect, normalize))
617        ):
618            return cached
619
620        normalized_table = exp.maybe_parse(table, into=exp.Table, dialect=dialect, copy=normalize)
621
622        if normalize:
623            for part in normalized_table.parts:
624                if isinstance(part, exp.Identifier):
625                    part.replace(
626                        normalize_name(part, dialect=dialect, is_table=True, normalize=normalize)
627                    )
628
629        self._normalized_table_cache[(normalized_table, dialect, normalize)] = normalized_table
630        return normalized_table
631
632    def _normalize_name(
633        self,
634        name: str | exp.Identifier,
635        dialect: DialectType = None,
636        is_table: bool = False,
637        normalize: bool | None = None,
638    ) -> str:
639        normalize = self.normalize if normalize is None else normalize
640
641        dialect = dialect or self.dialect
642        name_str = name if isinstance(name, str) else name.name
643        cache_key = (name_str, dialect, is_table, normalize)
644
645        if cached := self._normalized_name_cache.get(cache_key):
646            return cached
647
648        result = normalize_name(
649            name,
650            dialect=dialect,
651            is_table=is_table,
652            normalize=normalize,
653        ).name
654
655        self._normalized_name_cache[cache_key] = result
656        return result
657
658    def depth(self) -> int:
659        if not self.empty and not self._depth:
660            # The columns themselves are a mapping, but we don't want to include those
661            self._depth = super().depth() - 1
662        return self._depth
663
664    def _to_data_type(self, schema_type: str, dialect: DialectType = None) -> exp.DataType:
665        """
666        Convert a type represented as a string to the corresponding `sqlglot.exp.DataType` object.
667
668        Args:
669            schema_type: the type we want to convert.
670            dialect: the SQL dialect that will be used to parse `schema_type`, if needed.
671
672        Returns:
673            The resulting expression type.
674        """
675        if schema_type not in self._type_mapping_cache:
676            dialect = Dialect.get_or_raise(dialect) if dialect else self.dialect
677            udt = dialect.SUPPORTS_USER_DEFINED_TYPES
678
679            try:
680                expression = exp.DataType.from_str(schema_type, dialect=dialect, udt=udt)
681                expression.transform(dialect.normalize_identifier, copy=False)
682                self._type_mapping_cache[schema_type] = expression
683            except AttributeError:
684                in_dialect = f" in dialect {dialect}" if dialect else ""
685                raise SchemaError(f"Failed to build type '{schema_type}'{in_dialect}.")
686
687        return self._type_mapping_cache[schema_type]
688
689
690def normalize_name(
691    identifier: str | exp.Identifier,
692    dialect: DialectType = None,
693    is_table: bool = False,
694    normalize: bool | None = True,
695) -> exp.Identifier:
696    if isinstance(identifier, str):
697        identifier = exp.parse_identifier(identifier, dialect=dialect)
698    elif normalize:
699        identifier = identifier.copy()
700
701    if not normalize:
702        return identifier
703
704    # this is used for normalize_identifier, bigquery has special rules pertaining tables
705    identifier.meta["is_table"] = is_table
706    return Dialect.get_or_raise(dialect).normalize_identifier(identifier)
707
708
709def ensure_schema(
710    schema: Schema | dict[str, object] | None, **kwargs: Unpack[SchemaArgs]
711) -> Schema:
712    if isinstance(schema, Schema):
713        return schema
714
715    return MappingSchema(schema, **kwargs)
716
717
718def ensure_column_mapping(mapping: ColumnMapping | None) -> dict[str, t.Any]:
719    if mapping is None:
720        return {}
721    elif isinstance(mapping, dict):
722        return mapping
723    elif isinstance(mapping, str):
724        col_name_type_strs = [x.strip() for x in mapping.split(",")]
725        return {
726            name_type_str.split(":")[0].strip(): name_type_str.split(":")[1].strip()
727            for name_type_str in col_name_type_strs
728        }
729    elif isinstance(mapping, list):
730        return {x.strip(): None for x in mapping}
731
732    raise ValueError(f"Invalid mapping provided: {type(mapping)}")
733
734
735def flatten_schema(
736    schema: dict[str, object], depth: int | None = None, keys: list[str] | None = None
737) -> list[list[str]]:
738    tables: list[list[str]] = []
739    keys = keys or []
740    depth = dict_depth(schema) - 1 if depth is None else depth
741
742    for k, v in schema.items():
743        if depth == 1 or not isinstance(v, dict):
744            tables.append(keys + [k])
745        elif depth >= 2:
746            tables.extend(flatten_schema(v, depth - 1, keys + [k]))
747
748    return tables
749
750
751def nested_get(
752    d: dict[str, object], *path: tuple[str, str], raise_on_missing: bool = True
753) -> t.Any | None:
754    """
755    Get a value for a nested dictionary.
756
757    Args:
758        d: the dictionary to search.
759        *path: tuples of (name, key), where:
760            `key` is the key in the dictionary to get.
761            `name` is a string to use in the error if `key` isn't found.
762
763    Returns:
764        The value or None if it doesn't exist.
765    """
766    result: t.Any = d
767    for name, key in path:
768        result = result.get(key)
769        if result is None:
770            if raise_on_missing:
771                name = "table" if name == "this" else name
772                raise ValueError(f"Unknown {name}: {key}")
773            return None
774
775    return result
776
777
778def nested_set(d: dict[str, t.Any], keys: Sequence[str], value: t.Any) -> dict[str, t.Any]:
779    """
780    In-place set a value for a nested dictionary
781
782    Example:
783        >>> nested_set({}, ["top_key", "second_key"], "value")
784        {'top_key': {'second_key': 'value'}}
785
786        >>> nested_set({"top_key": {"third_key": "third_value"}}, ["top_key", "second_key"], "value")
787        {'top_key': {'third_key': 'third_value', 'second_key': 'value'}}
788
789    Args:
790        d: dictionary to update.
791        keys: the keys that makeup the path to `value`.
792        value: the value to set in the dictionary for the given key path.
793
794    Returns:
795        The (possibly) updated dictionary.
796    """
797    if not keys:
798        return d
799
800    if len(keys) == 1:
801        d[keys[0]] = value
802        return d
803
804    subd = d
805    for key in keys[:-1]:
806        if key not in subd:
807            subd = subd.setdefault(key, {})
808        else:
809            subd = subd[key]
810
811    subd[keys[-1]] = value
812    return d