From ef2db38de92f2329c1c366318bddfc7e3dee8415 Mon Sep 17 00:00:00 2001
From: Daniel Baumann <daniel.baumann@progress-linux.org>
Date: Sun, 19 Feb 2023 14:44:59 +0100
Subject: Adding upstream version 11.1.3.

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

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

diff --git a/sqlglot/optimizer/optimizer.py b/sqlglot/optimizer/optimizer.py
index 96fd56b..d9d04be 100644
--- a/sqlglot/optimizer/optimizer.py
+++ b/sqlglot/optimizer/optimizer.py
@@ -1,4 +1,10 @@
+from __future__ import annotations
+
+import typing as t
+
 import sqlglot
+from sqlglot import Schema, exp
+from sqlglot.dialects.dialect import DialectType
 from sqlglot.optimizer.annotate_types import annotate_types
 from sqlglot.optimizer.canonicalize import canonicalize
 from sqlglot.optimizer.eliminate_ctes import eliminate_ctes
@@ -24,8 +30,8 @@ RULES = (
     isolate_table_selects,
     qualify_columns,
     expand_laterals,
-    validate_qualify_columns,
     pushdown_projections,
+    validate_qualify_columns,
     normalize,
     unnest_subqueries,
     expand_multi_table_selects,
@@ -40,22 +46,31 @@ RULES = (
 )
 
 
-def optimize(expression, schema=None, db=None, catalog=None, rules=RULES, **kwargs):
+def optimize(
+    expression: str | exp.Expression,
+    schema: t.Optional[dict | Schema] = None,
+    db: t.Optional[str] = None,
+    catalog: t.Optional[str] = None,
+    dialect: DialectType = None,
+    rules: t.Sequence[t.Callable] = RULES,
+    **kwargs,
+):
     """
     Rewrite a sqlglot AST into an optimized form.
 
     Args:
-        expression (sqlglot.Expression): expression to optimize
-        schema (dict|sqlglot.optimizer.Schema): database schema.
+        expression: expression to optimize
+        schema: database schema.
             This can either be an instance of `sqlglot.optimizer.Schema` or a mapping in one of
             the following forms:
                 1. {table: {col: type}}
                 2. {db: {table: {col: type}}}
                 3. {catalog: {db: {table: {col: type}}}}
             If no schema is provided then the default schema defined at `sqlgot.schema` will be used
-        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 (sequence): sequence of optimizer rules to use.
+        db: specify the default database, as might be set by a `USE DATABASE db` statement
+        catalog: specify the default catalog, as might be set by a `USE CATALOG c` statement
+        dialect: The dialect to parse the sql string.
+        rules: sequence of optimizer rules to use.
             Many of the rules require tables and columns to be qualified.
             Do not remove qualify_tables or qualify_columns from the sequence of rules unless you know
             what you're doing!
@@ -65,7 +80,7 @@ def optimize(expression, schema=None, db=None, catalog=None, rules=RULES, **kwar
     """
     schema = ensure_schema(schema or sqlglot.schema)
     possible_kwargs = {"db": db, "catalog": catalog, "schema": schema, **kwargs}
-    expression = expression.copy()
+    expression = exp.maybe_parse(expression, dialect=dialect, copy=True)
     for rule in rules:
         # Find any additional rule parameters, beyond `expression`
         rule_params = rule.__code__.co_varnames
-- 
cgit v1.2.3