sqlglot.optimizer.optimizer

 1from __future__ import annotations
 2
 3import typing as t
 4
 5import sqlglot
 6from sqlglot import Schema, exp
 7from sqlglot.dialects.dialect import DialectType
 8from sqlglot.optimizer.annotate_types import annotate_types
 9from sqlglot.optimizer.canonicalize import canonicalize
10from sqlglot.optimizer.eliminate_ctes import eliminate_ctes
11from sqlglot.optimizer.eliminate_joins import eliminate_joins
12from sqlglot.optimizer.eliminate_subqueries import eliminate_subqueries
13from sqlglot.optimizer.merge_subqueries import merge_subqueries
14from sqlglot.optimizer.normalize import normalize
15from sqlglot.optimizer.optimize_joins import optimize_joins
16from sqlglot.optimizer.pushdown_predicates import pushdown_predicates
17from sqlglot.optimizer.pushdown_projections import pushdown_projections
18from sqlglot.optimizer.qualify import qualify
19from sqlglot.optimizer.qualify_columns import quote_identifiers
20from sqlglot.optimizer.simplify import simplify
21from sqlglot.optimizer.unnest_subqueries import unnest_subqueries
22from sqlglot.schema import ensure_schema
23
24RULES = (
25    qualify,
26    pushdown_projections,
27    normalize,
28    unnest_subqueries,
29    pushdown_predicates,
30    optimize_joins,
31    eliminate_subqueries,
32    merge_subqueries,
33    eliminate_joins,
34    eliminate_ctes,
35    quote_identifiers,
36    annotate_types,
37    canonicalize,
38    simplify,
39)
40
41
42def optimize(
43    expression: str | exp.Expression,
44    schema: t.Optional[dict | Schema] = None,
45    db: t.Optional[str] = None,
46    catalog: t.Optional[str] = None,
47    dialect: DialectType = None,
48    rules: t.Sequence[t.Callable] = RULES,
49    **kwargs,
50) -> exp.Expression:
51    """
52    Rewrite a sqlglot AST into an optimized form.
53
54    Args:
55        expression: expression to optimize
56        schema: database schema.
57            This can either be an instance of `sqlglot.optimizer.Schema` or a mapping in one of
58            the following forms:
59                1. {table: {col: type}}
60                2. {db: {table: {col: type}}}
61                3. {catalog: {db: {table: {col: type}}}}
62            If no schema is provided then the default schema defined at `sqlgot.schema` will be used
63        db: specify the default database, as might be set by a `USE DATABASE db` statement
64        catalog: specify the default catalog, as might be set by a `USE CATALOG c` statement
65        dialect: The dialect to parse the sql string.
66        rules: sequence of optimizer rules to use.
67            Many of the rules require tables and columns to be qualified.
68            Do not remove `qualify` from the sequence of rules unless you know what you're doing!
69        **kwargs: If a rule has a keyword argument with a same name in **kwargs, it will be passed in.
70
71    Returns:
72        The optimized expression.
73    """
74    schema = ensure_schema(schema or sqlglot.schema, dialect=dialect)
75    possible_kwargs = {
76        "db": db,
77        "catalog": catalog,
78        "schema": schema,
79        "dialect": dialect,
80        "isolate_tables": True,  # needed for other optimizations to perform well
81        "quote_identifiers": False,
82        **kwargs,
83    }
84
85    expression = exp.maybe_parse(expression, dialect=dialect, copy=True)
86    for rule in rules:
87        # Find any additional rule parameters, beyond `expression`
88        rule_params = rule.__code__.co_varnames
89        rule_kwargs = {
90            param: possible_kwargs[param] for param in rule_params if param in possible_kwargs
91        }
92        expression = rule(expression, **rule_kwargs)
93
94    return t.cast(exp.Expression, expression)