sqlglot.optimizer.optimizer

 1import sqlglot
 2from sqlglot.optimizer.annotate_types import annotate_types
 3from sqlglot.optimizer.canonicalize import canonicalize
 4from sqlglot.optimizer.eliminate_ctes import eliminate_ctes
 5from sqlglot.optimizer.eliminate_joins import eliminate_joins
 6from sqlglot.optimizer.eliminate_subqueries import eliminate_subqueries
 7from sqlglot.optimizer.expand_laterals import expand_laterals
 8from sqlglot.optimizer.expand_multi_table_selects import expand_multi_table_selects
 9from sqlglot.optimizer.isolate_table_selects import isolate_table_selects
10from sqlglot.optimizer.lower_identities import lower_identities
11from sqlglot.optimizer.merge_subqueries import merge_subqueries
12from sqlglot.optimizer.normalize import normalize
13from sqlglot.optimizer.optimize_joins import optimize_joins
14from sqlglot.optimizer.pushdown_predicates import pushdown_predicates
15from sqlglot.optimizer.pushdown_projections import pushdown_projections
16from sqlglot.optimizer.qualify_columns import qualify_columns, validate_qualify_columns
17from sqlglot.optimizer.qualify_tables import qualify_tables
18from sqlglot.optimizer.unnest_subqueries import unnest_subqueries
19from sqlglot.schema import ensure_schema
20
21RULES = (
22    lower_identities,
23    qualify_tables,
24    isolate_table_selects,
25    qualify_columns,
26    expand_laterals,
27    validate_qualify_columns,
28    pushdown_projections,
29    normalize,
30    unnest_subqueries,
31    expand_multi_table_selects,
32    pushdown_predicates,
33    optimize_joins,
34    eliminate_subqueries,
35    merge_subqueries,
36    eliminate_joins,
37    eliminate_ctes,
38    annotate_types,
39    canonicalize,
40)
41
42
43def optimize(expression, schema=None, db=None, catalog=None, rules=RULES, **kwargs):
44    """
45    Rewrite a sqlglot AST into an optimized form.
46
47    Args:
48        expression (sqlglot.Expression): expression to optimize
49        schema (dict|sqlglot.optimizer.Schema): database schema.
50            This can either be an instance of `sqlglot.optimizer.Schema` or a mapping in one of
51            the following forms:
52                1. {table: {col: type}}
53                2. {db: {table: {col: type}}}
54                3. {catalog: {db: {table: {col: type}}}}
55            If no schema is provided then the default schema defined at `sqlgot.schema` will be used
56        db (str): specify the default database, as might be set by a `USE DATABASE db` statement
57        catalog (str): specify the default catalog, as might be set by a `USE CATALOG c` statement
58        rules (sequence): sequence of optimizer rules to use.
59            Many of the rules require tables and columns to be qualified.
60            Do not remove qualify_tables or qualify_columns from the sequence of rules unless you know
61            what you're doing!
62        **kwargs: If a rule has a keyword argument with a same name in **kwargs, it will be passed in.
63    Returns:
64        sqlglot.Expression: optimized expression
65    """
66    schema = ensure_schema(schema or sqlglot.schema)
67    possible_kwargs = {"db": db, "catalog": catalog, "schema": schema, **kwargs}
68    expression = expression.copy()
69    for rule in rules:
70        # Find any additional rule parameters, beyond `expression`
71        rule_params = rule.__code__.co_varnames
72        rule_kwargs = {
73            param: possible_kwargs[param] for param in rule_params if param in possible_kwargs
74        }
75        expression = rule(expression, **rule_kwargs)
76    return expression