sqlglot.schema

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