sqlglot.optimizer.optimizer

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