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