sqlmesh.core.schema_diff

  1from __future__ import annotations
  2
  3import logging
  4import typing as t
  5from collections import defaultdict
  6from enum import Enum, auto
  7
  8from sqlglot import exp
  9from sqlglot.helper import ensure_list, seq_get
 10
 11from sqlmesh.utils import columns_to_types_to_struct
 12from sqlmesh.utils.pydantic import PydanticModel
 13
 14if t.TYPE_CHECKING:
 15    from sqlmesh.core._typing import TableName
 16
 17logger = logging.getLogger(__name__)
 18
 19
 20class TableAlterOperationType(Enum):
 21    ADD = auto()
 22    DROP = auto()
 23    ALTER_TYPE = auto()
 24
 25    @property
 26    def is_add(self) -> bool:
 27        return self == TableAlterOperationType.ADD
 28
 29    @property
 30    def is_drop(self) -> bool:
 31        return self == TableAlterOperationType.DROP
 32
 33    @property
 34    def is_alter_type(self) -> bool:
 35        return self == TableAlterOperationType.ALTER_TYPE
 36
 37
 38class TableAlterColumn(PydanticModel):
 39    name: str
 40    is_struct: bool
 41    is_array_of_struct: bool
 42    is_array_of_primitive: bool
 43    quoted: bool = False
 44
 45    @classmethod
 46    def primitive(self, name: str, quoted: bool = False) -> TableAlterColumn:
 47        return self(
 48            name=name,
 49            is_struct=False,
 50            is_array_of_struct=False,
 51            is_array_of_primitive=False,
 52            quoted=quoted,
 53        )
 54
 55    @classmethod
 56    def struct(self, name: str, quoted: bool = False) -> TableAlterColumn:
 57        return self(
 58            name=name,
 59            is_struct=True,
 60            is_array_of_struct=False,
 61            is_array_of_primitive=False,
 62            quoted=quoted,
 63        )
 64
 65    @classmethod
 66    def array_of_struct(self, name: str, quoted: bool = False) -> TableAlterColumn:
 67        return self(
 68            name=name,
 69            is_struct=False,
 70            is_array_of_struct=True,
 71            is_array_of_primitive=False,
 72            quoted=quoted,
 73        )
 74
 75    @classmethod
 76    def array_of_primitive(self, name: str, quoted: bool = False) -> TableAlterColumn:
 77        return self(
 78            name=name,
 79            is_struct=False,
 80            is_array_of_struct=False,
 81            is_array_of_primitive=True,
 82            quoted=quoted,
 83        )
 84
 85    @classmethod
 86    def from_struct_kwarg(self, struct: exp.ColumnDef) -> TableAlterColumn:
 87        name = struct.alias_or_name
 88        quoted = struct.this.quoted
 89        kwarg_type = struct.args["kind"]
 90
 91        if kwarg_type.is_type(exp.DataType.Type.STRUCT):
 92            return self.struct(name, quoted=quoted)
 93        elif kwarg_type.is_type(exp.DataType.Type.ARRAY):
 94            if kwarg_type.expressions[0].is_type(exp.DataType.Type.STRUCT):
 95                return self.array_of_struct(name, quoted=quoted)
 96            else:
 97                return self.array_of_primitive(name, quoted=quoted)
 98        else:
 99            return self.primitive(name, quoted=quoted)
100
101    @property
102    def is_array(self) -> bool:
103        return self.is_array_of_struct or self.is_array_of_primitive
104
105    @property
106    def is_primitive(self) -> bool:
107        return not self.is_struct and not self.is_array
108
109    @property
110    def is_nested(self) -> bool:
111        return not self.is_primitive
112
113    @property
114    def identifier(self) -> exp.Identifier:
115        return exp.to_identifier(self.name, quoted=self.quoted)
116
117
118class TableAlterColumnPosition(PydanticModel):
119    is_first: bool
120    is_last: bool
121    after: t.Optional[exp.Identifier] = None
122
123    @classmethod
124    def first(self) -> TableAlterColumnPosition:
125        return self(is_first=True, is_last=False, after=None)
126
127    @classmethod
128    def last(
129        self, after: t.Optional[t.Union[str, exp.Identifier]] = None
130    ) -> TableAlterColumnPosition:
131        return self(is_first=False, is_last=True, after=exp.to_identifier(after) if after else None)
132
133    @classmethod
134    def middle(self, after: t.Union[str, exp.Identifier]) -> TableAlterColumnPosition:
135        return self(is_first=False, is_last=False, after=exp.to_identifier(after))
136
137    @classmethod
138    def create(
139        self,
140        pos: int,
141        current_kwargs: t.List[exp.ColumnDef],
142        replacing_col: bool = False,
143    ) -> TableAlterColumnPosition:
144        is_first = pos == 0
145        is_last = pos == len(current_kwargs) - int(replacing_col)
146        after = None
147        if not is_first:
148            prior_kwarg = current_kwargs[pos - 1]
149            after, _ = _get_name_and_type(prior_kwarg)
150        return self(is_first=is_first, is_last=is_last, after=after)
151
152    @property
153    def column_position_node(self) -> t.Optional[exp.ColumnPosition]:
154        column = self.after if not self.is_last else None
155        position = None
156        if self.is_first:
157            position = "FIRST"
158        elif column and not self.is_last:
159            position = "AFTER"
160        return exp.ColumnPosition(this=column, position=position)
161
162
163class TableAlterOperation(PydanticModel):
164    op: TableAlterOperationType
165    columns: t.List[TableAlterColumn]
166    column_type: exp.DataType
167    expected_table_struct: exp.DataType
168    add_position: t.Optional[TableAlterColumnPosition] = None
169    current_type: t.Optional[exp.DataType] = None
170
171    @classmethod
172    def add(
173        self,
174        columns: t.Union[TableAlterColumn, t.List[TableAlterColumn]],
175        column_type: t.Union[str, exp.DataType],
176        expected_table_struct: t.Union[str, exp.DataType],
177        position: t.Optional[TableAlterColumnPosition] = None,
178    ) -> TableAlterOperation:
179        return self(
180            op=TableAlterOperationType.ADD,
181            columns=ensure_list(columns),
182            column_type=exp.DataType.build(column_type),
183            add_position=position,
184            expected_table_struct=exp.DataType.build(expected_table_struct),
185        )
186
187    @classmethod
188    def drop(
189        self,
190        columns: t.Union[TableAlterColumn, t.List[TableAlterColumn]],
191        expected_table_struct: t.Union[str, exp.DataType],
192        column_type: t.Optional[t.Union[str, exp.DataType]] = None,
193    ) -> TableAlterOperation:
194        column_type = exp.DataType.build(column_type) if column_type else exp.DataType.build("INT")
195        return self(
196            op=TableAlterOperationType.DROP,
197            columns=ensure_list(columns),
198            column_type=column_type,
199            expected_table_struct=exp.DataType.build(expected_table_struct),
200        )
201
202    @classmethod
203    def alter_type(
204        self,
205        columns: t.Union[TableAlterColumn, t.List[TableAlterColumn]],
206        column_type: t.Union[str, exp.DataType],
207        current_type: t.Union[str, exp.DataType],
208        expected_table_struct: t.Union[str, exp.DataType],
209        position: t.Optional[TableAlterColumnPosition] = None,
210    ) -> TableAlterOperation:
211        return self(
212            op=TableAlterOperationType.ALTER_TYPE,
213            columns=ensure_list(columns),
214            column_type=exp.DataType.build(column_type),
215            add_position=position,
216            current_type=exp.DataType.build(current_type),
217            expected_table_struct=exp.DataType.build(expected_table_struct),
218        )
219
220    @property
221    def is_add(self) -> bool:
222        return self.op.is_add
223
224    @property
225    def is_drop(self) -> bool:
226        return self.op.is_drop
227
228    @property
229    def is_alter_type(self) -> bool:
230        return self.op.is_alter_type
231
232    def column_identifiers(self, array_element_selector: str) -> t.List[exp.Identifier]:
233        results = []
234        for column in self.columns:
235            results.append(column.identifier)
236            if column.is_array_of_struct and len(self.columns) > 1 and array_element_selector:
237                results.append(exp.to_identifier(array_element_selector))
238        return results
239
240    def column(self, array_element_selector: str) -> t.Union[exp.Dot, exp.Identifier]:
241        columns = self.column_identifiers(array_element_selector)
242        if len(columns) == 1:
243            return columns[0]
244        return exp.Dot.build(columns)
245
246    def column_def(self, array_element_selector: str) -> exp.ColumnDef:
247        return exp.ColumnDef(
248            this=self.column(array_element_selector),
249            kind=self.column_type,
250        )
251
252    def expression(
253        self, table_name: t.Union[str, exp.Table], array_element_selector: str
254    ) -> exp.AlterTable:
255        if self.is_alter_type:
256            return exp.AlterTable(
257                this=exp.to_table(table_name),
258                actions=[
259                    exp.AlterColumn(
260                        this=self.column(array_element_selector),
261                        dtype=self.column_type,
262                    )
263                ],
264            )
265        elif self.is_add:
266            alter_table = exp.AlterTable(this=exp.to_table(table_name))
267            column = self.column_def(array_element_selector)
268            alter_table.set("actions", [column])
269            if self.add_position:
270                column.set("position", self.add_position.column_position_node)
271            return alter_table
272        elif self.is_drop:
273            alter_table = exp.AlterTable(this=exp.to_table(table_name))
274            drop_column = exp.Drop(this=self.column(array_element_selector), kind="COLUMN")
275            alter_table.set("actions", [drop_column])
276            return alter_table
277        else:
278            raise ValueError(f"Unknown operation {self.op}")
279
280
281class SchemaDiffer(PydanticModel):
282    """
283    Compares a source schema against a target schema and returns a list of alter statements to have the source
284    match the structure of target. Some engines have constraints on the types of operations that can be performed
285    therefore the final structure may not match the target exactly but it will be as close as possible. Two potential
286    differences that can happen:
287    1. Column order can be different if the engine doesn't support positional additions. Another reason for difference
288    is if a column is just moved since we don't currently support fixing moves.
289    2. Nested operations will be represented using a drop/add of the root column if the engine doesn't support nested
290    operations. As a result historical data is lost.
291    3. Column type changes will be reflected but it can be done through a drop/add if the change is not a compatible
292    change. As a result historical data is lost.
293
294    Potential future improvements:
295    1. Support precision changes on columns like VARCHAR and DECIMAL. Each engine has different rules on what is allowed
296    2. Support column moves. Databricks Delta supports moves and would allow exact matches.
297    """
298
299    support_positional_add: bool = False
300    support_nested_operations: bool = False
301    array_element_selector: str = ""
302    compatible_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = {}
303    support_coercing_compatible_types: bool = False
304
305    _coerceable_types: t.Dict[exp.DataType, t.Set[exp.DataType]] = {}
306
307    @property
308    def coerceable_types(self) -> t.Dict[exp.DataType, t.Set[exp.DataType]]:
309        if not self._coerceable_types:
310            if not self.support_coercing_compatible_types or not self.compatible_types:
311                return {}
312            coerceable_types = defaultdict(set)
313            for source_type, target_types in self.compatible_types.items():
314                for target_type in target_types:
315                    coerceable_types[target_type].add(source_type)
316            self._coerceable_types = coerceable_types
317        return self._coerceable_types
318
319    def _is_compatible_type(self, current_type: exp.DataType, new_type: exp.DataType) -> bool:
320        if current_type == new_type:
321            return True
322        if current_type in self.compatible_types:
323            return new_type in self.compatible_types[current_type]
324        return False
325
326    def _is_coerceable_type(self, current_type: exp.DataType, new_type: exp.DataType) -> bool:
327        if not self.support_coercing_compatible_types:
328            return False
329        if current_type in self.coerceable_types:
330            is_coerceable = new_type in self.coerceable_types[current_type]
331            if is_coerceable:
332                logger.warning(
333                    f"Coercing type {current_type} to {new_type} which means an alter will not be performed and therefore the resulting table structure will not match what is in the query.\nUpdate your model to cast the value to {current_type} type in order to remove this warning.",
334                )
335            return is_coerceable
336        return False
337
338    def _get_matching_kwarg(
339        self,
340        current_kwarg: t.Union[str, exp.ColumnDef],
341        new_struct: exp.DataType,
342        current_pos: int,
343    ) -> t.Tuple[t.Optional[int], t.Optional[exp.ColumnDef]]:
344        current_name = (
345            exp.to_identifier(current_kwarg)
346            if isinstance(current_kwarg, str)
347            else _get_name_and_type(current_kwarg)[0]
348        )
349        # First check if we have the same column in the same position to get O(1) complexity
350        new_kwarg = seq_get(new_struct.expressions, current_pos)
351        if new_kwarg:
352            new_name, new_type = _get_name_and_type(new_kwarg)
353            if current_name.this == new_name.this:
354                return current_pos, new_kwarg
355        # If not, check if we have the same column in all positions with O(n) complexity
356        for i, new_kwarg in enumerate(new_struct.expressions):
357            new_name, new_type = _get_name_and_type(new_kwarg)
358            if current_name.this == new_name.this:
359                return i, new_kwarg
360        return None, None
361
362    def _drop_operation(
363        self,
364        columns: t.Union[TableAlterColumn, t.List[TableAlterColumn]],
365        struct: exp.DataType,
366        pos: int,
367        root_struct: exp.DataType,
368    ) -> t.List[TableAlterOperation]:
369        columns = ensure_list(columns)
370        operations = []
371        column_pos, column_kwarg = self._get_matching_kwarg(columns[-1].name, struct, pos)
372        assert column_pos is not None
373        assert column_kwarg
374        struct.expressions.pop(column_pos)
375        operations.append(
376            TableAlterOperation.drop(columns, root_struct.copy(), column_kwarg.args["kind"])
377        )
378        return operations
379
380    def _resolve_drop_operation(
381        self,
382        parent_columns: t.List[TableAlterColumn],
383        current_struct: exp.DataType,
384        new_struct: exp.DataType,
385        root_struct: exp.DataType,
386    ) -> t.List[TableAlterOperation]:
387        operations = []
388        for current_pos, current_kwarg in enumerate(current_struct.expressions.copy()):
389            new_pos, _ = self._get_matching_kwarg(current_kwarg, new_struct, current_pos)
390            columns = parent_columns + [TableAlterColumn.from_struct_kwarg(current_kwarg)]
391            if new_pos is None:
392                operations.extend(
393                    self._drop_operation(columns, current_struct, current_pos, root_struct)
394                )
395        return operations
396
397    def _add_operation(
398        self,
399        columns: t.List[TableAlterColumn],
400        new_pos: int,
401        new_kwarg: exp.ColumnDef,
402        current_struct: exp.DataType,
403        root_struct: exp.DataType,
404    ) -> t.List[TableAlterOperation]:
405        if self.support_positional_add:
406            col_pos = TableAlterColumnPosition.create(new_pos, current_struct.expressions)
407            current_struct.expressions.insert(new_pos, new_kwarg)
408        else:
409            col_pos = None
410            current_struct.expressions.append(new_kwarg)
411        return [
412            TableAlterOperation.add(
413                columns,
414                new_kwarg.args["kind"],
415                root_struct.copy(),
416                col_pos,
417            )
418        ]
419
420    def _resolve_add_operations(
421        self,
422        parent_columns: t.List[TableAlterColumn],
423        current_struct: exp.DataType,
424        new_struct: exp.DataType,
425        root_struct: exp.DataType,
426    ) -> t.List[TableAlterOperation]:
427        operations = []
428        for new_pos, new_kwarg in enumerate(new_struct.expressions):
429            possible_current_pos, _ = self._get_matching_kwarg(new_kwarg, current_struct, new_pos)
430            if possible_current_pos is None:
431                columns = parent_columns + [TableAlterColumn.from_struct_kwarg(new_kwarg)]
432                operations.extend(
433                    self._add_operation(columns, new_pos, new_kwarg, current_struct, root_struct)
434                )
435        return operations
436
437    def _alter_operation(
438        self,
439        columns: t.List[TableAlterColumn],
440        pos: int,
441        struct: exp.DataType,
442        new_type: exp.DataType,
443        current_type: t.Union[str, exp.DataType],
444        root_struct: exp.DataType,
445        new_kwarg: exp.ColumnDef,
446    ) -> t.List[TableAlterOperation]:
447        # We don't copy on purpose here because current_type may need to be mutated inside
448        # _get_operations (struct.expressions.pop and struct.expressions.insert)
449        current_type = exp.DataType.build(current_type, copy=False)
450        if self.support_nested_operations:
451            if new_type.this == current_type.this == exp.DataType.Type.STRUCT:
452                return self._get_operations(
453                    columns,
454                    current_type,
455                    new_type,
456                    root_struct,
457                )
458            if new_type.this == current_type.this == exp.DataType.Type.ARRAY:
459                new_array_type = new_type.expressions[0]
460                current_array_type = current_type.expressions[0]
461                if new_array_type.this == current_array_type.this == exp.DataType.Type.STRUCT:
462                    return self._get_operations(
463                        columns,
464                        current_array_type,
465                        new_array_type,
466                        root_struct,
467                    )
468        if self._is_coerceable_type(current_type, new_type):
469            return []
470        elif self._is_compatible_type(current_type, new_type):
471            struct.expressions.pop(pos)
472            struct.expressions.insert(pos, new_kwarg)
473            col_pos = (
474                TableAlterColumnPosition.create(pos, struct.expressions, replacing_col=True)
475                if self.support_positional_add
476                else None
477            )
478            return [
479                TableAlterOperation.alter_type(
480                    columns,
481                    new_type,
482                    current_type,
483                    root_struct.copy(),
484                    col_pos,
485                )
486            ]
487        else:
488            return self._drop_operation(
489                columns, root_struct, pos, root_struct
490            ) + self._add_operation(columns, pos, new_kwarg, struct, root_struct)
491
492    def _resolve_alter_operations(
493        self,
494        parent_columns: t.List[TableAlterColumn],
495        current_struct: exp.DataType,
496        new_struct: exp.DataType,
497        root_struct: exp.DataType,
498    ) -> t.List[TableAlterOperation]:
499        operations = []
500        for current_pos, current_kwarg in enumerate(current_struct.expressions.copy()):
501            _, new_kwarg = self._get_matching_kwarg(current_kwarg, new_struct, current_pos)
502            assert new_kwarg
503            _, new_type = _get_name_and_type(new_kwarg)
504            _, current_type = _get_name_and_type(current_kwarg)
505            columns = parent_columns + [TableAlterColumn.from_struct_kwarg(current_kwarg)]
506            if new_type == current_type:
507                continue
508            operations.extend(
509                self._alter_operation(
510                    columns,
511                    current_pos,
512                    current_struct,
513                    new_type,
514                    current_type,
515                    root_struct,
516                    new_kwarg,
517                )
518            )
519        return operations
520
521    def _get_operations(
522        self,
523        parent_columns: t.List[TableAlterColumn],
524        current_struct: exp.DataType,
525        new_struct: exp.DataType,
526        root_struct: exp.DataType,
527    ) -> t.List[TableAlterOperation]:
528        root_struct = root_struct or current_struct
529        parent_columns = parent_columns or []
530        operations = []
531        operations.extend(
532            self._resolve_drop_operation(parent_columns, current_struct, new_struct, root_struct)
533        )
534        operations.extend(
535            self._resolve_add_operations(parent_columns, current_struct, new_struct, root_struct)
536        )
537        operations.extend(
538            self._resolve_alter_operations(parent_columns, current_struct, new_struct, root_struct)
539        )
540        return operations
541
542    def _from_structs(
543        self, current_struct: exp.DataType, new_struct: exp.DataType
544    ) -> t.List[TableAlterOperation]:
545        return self._get_operations([], current_struct, new_struct, current_struct)
546
547    def compare_structs(
548        self, table_name: t.Union[str, exp.Table], current: exp.DataType, new: exp.DataType
549    ) -> t.List[exp.AlterTable]:
550        """
551        Compares two schemas represented as structs.
552
553        Args:
554            current: The current schema.
555            new: The new schema.
556
557        Returns:
558            The list of table alter operations.
559        """
560        return [
561            op.expression(table_name, self.array_element_selector)
562            for op in self._from_structs(current, new)
563        ]
564
565    def compare_columns(
566        self,
567        table_name: TableName,
568        current: t.Dict[str, exp.DataType],
569        new: t.Dict[str, exp.DataType],
570    ) -> t.List[exp.AlterTable]:
571        """
572        Compares two schemas represented as dictionaries of column names and types.
573
574        Args:
575            current: The current schema.
576            new: The new schema.
577
578        Returns:
579            The list of schema deltas.
580        """
581        return self.compare_structs(
582            table_name, columns_to_types_to_struct(current), columns_to_types_to_struct(new)
583        )
584
585
586def _get_name_and_type(struct: exp.ColumnDef) -> t.Tuple[exp.Identifier, exp.DataType]:
587    return struct.this, struct.args["kind"]