Edit on GitHub

sqlglot.optimizer.canonicalize

  1from __future__ import annotations
  2
  3import itertools
  4
  5from sqlglot import exp
  6
  7
  8def canonicalize(expression: exp.Expression) -> exp.Expression:
  9    """Converts a sql expression into a standard form.
 10
 11    This method relies on annotate_types because many of the
 12    conversions rely on type inference.
 13
 14    Args:
 15        expression: The expression to canonicalize.
 16    """
 17    exp.replace_children(expression, canonicalize)
 18
 19    expression = add_text_to_concat(expression)
 20    expression = replace_date_funcs(expression)
 21    expression = coerce_type(expression)
 22    expression = remove_redundant_casts(expression)
 23    expression = ensure_bool_predicates(expression)
 24    expression = remove_ascending_order(expression)
 25
 26    return expression
 27
 28
 29def add_text_to_concat(node: exp.Expression) -> exp.Expression:
 30    if isinstance(node, exp.Add) and node.type and node.type.this in exp.DataType.TEXT_TYPES:
 31        node = exp.Concat(expressions=[node.left, node.right])
 32    return node
 33
 34
 35def replace_date_funcs(node: exp.Expression) -> exp.Expression:
 36    if isinstance(node, exp.Date) and not node.expressions and not node.args.get("zone"):
 37        return exp.cast(node.this, to=exp.DataType.Type.DATE)
 38    if isinstance(node, exp.Timestamp) and not node.expression:
 39        return exp.cast(node.this, to=exp.DataType.Type.TIMESTAMP)
 40    return node
 41
 42
 43def coerce_type(node: exp.Expression) -> exp.Expression:
 44    if isinstance(node, exp.Binary):
 45        _coerce_date(node.left, node.right)
 46    elif isinstance(node, exp.Between):
 47        _coerce_date(node.this, node.args["low"])
 48    elif isinstance(node, exp.Extract) and not node.expression.type.is_type(
 49        *exp.DataType.TEMPORAL_TYPES
 50    ):
 51        _replace_cast(node.expression, exp.DataType.Type.DATETIME)
 52
 53    return node
 54
 55
 56def remove_redundant_casts(expression: exp.Expression) -> exp.Expression:
 57    if (
 58        isinstance(expression, exp.Cast)
 59        and expression.to.type
 60        and expression.this.type
 61        and expression.to.type.this == expression.this.type.this
 62    ):
 63        return expression.this
 64    return expression
 65
 66
 67def ensure_bool_predicates(expression: exp.Expression) -> exp.Expression:
 68    if isinstance(expression, exp.Connector):
 69        _replace_int_predicate(expression.left)
 70        _replace_int_predicate(expression.right)
 71
 72    elif isinstance(expression, (exp.Where, exp.Having)) or (
 73        # We can't replace num in CASE x WHEN num ..., because it's not the full predicate
 74        isinstance(expression, exp.If)
 75        and not (isinstance(expression.parent, exp.Case) and expression.parent.this)
 76    ):
 77        _replace_int_predicate(expression.this)
 78
 79    return expression
 80
 81
 82def remove_ascending_order(expression: exp.Expression) -> exp.Expression:
 83    if isinstance(expression, exp.Ordered) and expression.args.get("desc") is False:
 84        # Convert ORDER BY a ASC to ORDER BY a
 85        expression.set("desc", None)
 86
 87    return expression
 88
 89
 90def _coerce_date(a: exp.Expression, b: exp.Expression) -> None:
 91    for a, b in itertools.permutations([a, b]):
 92        if (
 93            a.type
 94            and a.type.this == exp.DataType.Type.DATE
 95            and b.type
 96            and b.type.this not in (exp.DataType.Type.DATE, exp.DataType.Type.INTERVAL)
 97        ):
 98            _replace_cast(b, exp.DataType.Type.DATE)
 99
100
101def _replace_cast(node: exp.Expression, to: exp.DataType.Type) -> None:
102    node.replace(exp.cast(node.copy(), to=to))
103
104
105def _replace_int_predicate(expression: exp.Expression) -> None:
106    if isinstance(expression, exp.Coalesce):
107        for _, child in expression.iter_expressions():
108            _replace_int_predicate(child)
109    elif expression.type and expression.type.this in exp.DataType.INTEGER_TYPES:
110        expression.replace(exp.NEQ(this=expression.copy(), expression=exp.Literal.number(0)))
def canonicalize( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
 9def canonicalize(expression: exp.Expression) -> exp.Expression:
10    """Converts a sql expression into a standard form.
11
12    This method relies on annotate_types because many of the
13    conversions rely on type inference.
14
15    Args:
16        expression: The expression to canonicalize.
17    """
18    exp.replace_children(expression, canonicalize)
19
20    expression = add_text_to_concat(expression)
21    expression = replace_date_funcs(expression)
22    expression = coerce_type(expression)
23    expression = remove_redundant_casts(expression)
24    expression = ensure_bool_predicates(expression)
25    expression = remove_ascending_order(expression)
26
27    return expression

Converts a sql expression into a standard form.

This method relies on annotate_types because many of the conversions rely on type inference.

Arguments:
  • expression: The expression to canonicalize.
def add_text_to_concat(node: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
30def add_text_to_concat(node: exp.Expression) -> exp.Expression:
31    if isinstance(node, exp.Add) and node.type and node.type.this in exp.DataType.TEXT_TYPES:
32        node = exp.Concat(expressions=[node.left, node.right])
33    return node
def replace_date_funcs(node: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
36def replace_date_funcs(node: exp.Expression) -> exp.Expression:
37    if isinstance(node, exp.Date) and not node.expressions and not node.args.get("zone"):
38        return exp.cast(node.this, to=exp.DataType.Type.DATE)
39    if isinstance(node, exp.Timestamp) and not node.expression:
40        return exp.cast(node.this, to=exp.DataType.Type.TIMESTAMP)
41    return node
def coerce_type(node: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
44def coerce_type(node: exp.Expression) -> exp.Expression:
45    if isinstance(node, exp.Binary):
46        _coerce_date(node.left, node.right)
47    elif isinstance(node, exp.Between):
48        _coerce_date(node.this, node.args["low"])
49    elif isinstance(node, exp.Extract) and not node.expression.type.is_type(
50        *exp.DataType.TEMPORAL_TYPES
51    ):
52        _replace_cast(node.expression, exp.DataType.Type.DATETIME)
53
54    return node
def remove_redundant_casts( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
57def remove_redundant_casts(expression: exp.Expression) -> exp.Expression:
58    if (
59        isinstance(expression, exp.Cast)
60        and expression.to.type
61        and expression.this.type
62        and expression.to.type.this == expression.this.type.this
63    ):
64        return expression.this
65    return expression
def ensure_bool_predicates( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
68def ensure_bool_predicates(expression: exp.Expression) -> exp.Expression:
69    if isinstance(expression, exp.Connector):
70        _replace_int_predicate(expression.left)
71        _replace_int_predicate(expression.right)
72
73    elif isinstance(expression, (exp.Where, exp.Having)) or (
74        # We can't replace num in CASE x WHEN num ..., because it's not the full predicate
75        isinstance(expression, exp.If)
76        and not (isinstance(expression.parent, exp.Case) and expression.parent.this)
77    ):
78        _replace_int_predicate(expression.this)
79
80    return expression
def remove_ascending_order( expression: sqlglot.expressions.Expression) -> sqlglot.expressions.Expression:
83def remove_ascending_order(expression: exp.Expression) -> exp.Expression:
84    if isinstance(expression, exp.Ordered) and expression.args.get("desc") is False:
85        # Convert ORDER BY a ASC to ORDER BY a
86        expression.set("desc", None)
87
88    return expression