From 90150543f9314be683d22a16339effd774192f6d Mon Sep 17 00:00:00 2001
From: Daniel Baumann <daniel.baumann@progress-linux.org>
Date: Thu, 22 Sep 2022 06:31:28 +0200
Subject: Merging upstream version 6.1.1.

Signed-off-by: Daniel Baumann <daniel.baumann@progress-linux.org>
---
 sqlglot/optimizer/optimizer.py | 39 +++++++++++++++++++++++++++------------
 1 file changed, 27 insertions(+), 12 deletions(-)

(limited to 'sqlglot/optimizer/optimizer.py')

diff --git a/sqlglot/optimizer/optimizer.py b/sqlglot/optimizer/optimizer.py
index c03fe3c..c8c2403 100644
--- a/sqlglot/optimizer/optimizer.py
+++ b/sqlglot/optimizer/optimizer.py
@@ -1,6 +1,7 @@
 from sqlglot.optimizer.eliminate_subqueries import eliminate_subqueries
 from sqlglot.optimizer.expand_multi_table_selects import expand_multi_table_selects
 from sqlglot.optimizer.isolate_table_selects import isolate_table_selects
+from sqlglot.optimizer.merge_derived_tables import merge_derived_tables
 from sqlglot.optimizer.normalize import normalize
 from sqlglot.optimizer.optimize_joins import optimize_joins
 from sqlglot.optimizer.pushdown_predicates import pushdown_predicates
@@ -10,8 +11,23 @@ from sqlglot.optimizer.qualify_tables import qualify_tables
 from sqlglot.optimizer.quote_identities import quote_identities
 from sqlglot.optimizer.unnest_subqueries import unnest_subqueries
 
+RULES = (
+    qualify_tables,
+    isolate_table_selects,
+    qualify_columns,
+    pushdown_projections,
+    normalize,
+    unnest_subqueries,
+    expand_multi_table_selects,
+    pushdown_predicates,
+    optimize_joins,
+    eliminate_subqueries,
+    merge_derived_tables,
+    quote_identities,
+)
 
-def optimize(expression, schema=None, db=None, catalog=None):
+
+def optimize(expression, schema=None, db=None, catalog=None, rules=RULES, **kwargs):
     """
     Rewrite a sqlglot AST into an optimized form.
 
@@ -25,19 +41,18 @@ def optimize(expression, schema=None, db=None, catalog=None):
                 3. {catalog: {db: {table: {col: type}}}}
         db (str): specify the default database, as might be set by a `USE DATABASE db` statement
         catalog (str): specify the default catalog, as might be set by a `USE CATALOG c` statement
+        rules (list): sequence of optimizer rules to use
+        **kwargs: If a rule has a keyword argument with a same name in **kwargs, it will be passed in.
     Returns:
         sqlglot.Expression: optimized expression
     """
+    possible_kwargs = {"db": db, "catalog": catalog, "schema": schema, **kwargs}
     expression = expression.copy()
-    expression = qualify_tables(expression, db=db, catalog=catalog)
-    expression = isolate_table_selects(expression)
-    expression = qualify_columns(expression, schema)
-    expression = pushdown_projections(expression)
-    expression = normalize(expression)
-    expression = unnest_subqueries(expression)
-    expression = expand_multi_table_selects(expression)
-    expression = pushdown_predicates(expression)
-    expression = optimize_joins(expression)
-    expression = eliminate_subqueries(expression)
-    expression = quote_identities(expression)
+    for rule in rules:
+
+        # Find any additional rule parameters, beyond `expression`
+        rule_params = rule.__code__.co_varnames
+        rule_kwargs = {param: possible_kwargs[param] for param in rule_params if param in possible_kwargs}
+
+        expression = rule(expression, **rule_kwargs)
     return expression
-- 
cgit v1.2.3