Repository: tobymao/sqlglot

Issue:
Annotating types for dialect snowflake does not work
**TLDR**
Firstly, thanks for an amazing library! Annotating the types for dialect `snowflake` does not seem to work. Furthermore supplying a MappingSchema even seems to break the qualifying step. I hope I have not overlooked anything important here. We are using sqlglot for a variety of dialects and this problem only happens for `snowflake`:

**Example 1**

To Reproduce supply a simple schema and a simple query an qualify and annotate_types:

```py
import sqlglot
from sqlglot.optimizer.qualify import qualify
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.schema import ensure_schema

dialect = "snowflake"
schema = {
        "SCHEMA": {
            "TBL": {"COL": "INT", "col2": "VARCHAR"},
        }
    }
example_query = 'SELECT * FROM "SCHEMA"."TBL"'

# get qualified expression
expression = sqlglot.parse_one(example_query, dialect=dialect)
qual = qualify(expression, schema=schema, dialect=dialect)

# print the qualified statement
print(qual.sql(dialect=dialect))

# get the annotated types
expression = annotate_types(expression, schema=schema, dialect=dialect)
for col in expression.selects:
    print(col.output_name)
    print(col.type)
```

In this example you can see that the qualifier works (star select is expaneded) but the column types are set to UNKNOWN after annotating the types. The exact same code works for dialect `duckdb`.

**Example 2**

```py
import sqlglot
from sqlglot.optimizer.qualify import qualify
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.schema import ensure_schema

dialect = "snowflake"
schema = ensure_schema({
        "SCHEMA": {
            "TBL": {"COL": "INT", "col2": "VARCHAR"},
        }
    })
example_query = 'SELECT * FROM "SCHEMA"."TBL"'

# get qualified expression
expression = sqlglot.parse_one(example_query, dialect=dialect)
qual = qualify(expression, schema=schema, dialect=dialect)

# print the qualified statement
print(qual.sql(dialect=dialect))
```

When supplying a `MappingSchema` instead of a python dict, the star select expansion of the qualifying step already fails. For `duckdb` a `MappingSchema` also works.

Discussion:
@georgesittas: Thanks for the report @sh-rp, for whatever reason we weren't propagating the dialect properly. You need to pass it to `ensure_schema`, otherwise Snowflake's normalization logic doesn't kick in and there's a discrepancy between the schema's and query's identifiers.@sh-rp: @georgesittas thanks for the super fast fix and help :)

Solution patch:
diff --git a/sqlglot/optimizer/annotate_types.py b/sqlglot/optimizer/annotate_types.py
index 8c9c0aef0a..e779958aec 100644
--- a/sqlglot/optimizer/annotate_types.py
+++ b/sqlglot/optimizer/annotate_types.py
@@ -32,7 +32,7 @@ def annotate_types(
     schema: t.Optional[t.Dict | Schema] = None,
     annotators: t.Optional[AnnotatorsType] = None,
     coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
-    dialect: t.Optional[DialectType] = None,
+    dialect: DialectType = None,
 ) -> E:
     """
     Infers the types of an expression, annotating its AST accordingly.
@@ -55,9 +55,9 @@ def annotate_types(
         The expression annotated with types.
     """
 
-    schema = ensure_schema(schema)
+    schema = ensure_schema(schema, dialect=dialect)
 
-    return TypeAnnotator(schema, annotators, coerces_to, dialect=dialect).annotate(expression)
+    return TypeAnnotator(schema, annotators, coerces_to).annotate(expression)
 
 
 def _coerce_date_literal(l: exp.Expression, unit: t.Optional[exp.Expression]) -> exp.DataType.Type:
@@ -182,10 +182,9 @@ def __init__(
         annotators: t.Optional[AnnotatorsType] = None,
         coerces_to: t.Optional[t.Dict[exp.DataType.Type, t.Set[exp.DataType.Type]]] = None,
         binary_coercions: t.Optional[BinaryCoercions] = None,
-        dialect: t.Optional[DialectType] = None,
     ) -> None:
         self.schema = schema
-        self.annotators = annotators or Dialect.get_or_raise(dialect).ANNOTATORS
+        self.annotators = annotators or Dialect.get_or_raise(schema.dialect).ANNOTATORS
         self.coerces_to = coerces_to or self.COERCES_TO
         self.binary_coercions = binary_coercions or self.BINARY_COERCIONS
 
diff --git a/sqlglot/optimizer/isolate_table_selects.py b/sqlglot/optimizer/isolate_table_selects.py
index a6524b8f98..e0c1d9c1d2 100644
--- a/sqlglot/optimizer/isolate_table_selects.py
+++ b/sqlglot/optimizer/isolate_table_selects.py
@@ -1,17 +1,32 @@
+from __future__ import annotations
+
+import typing as t
+
 from sqlglot import alias, exp
 from sqlglot.errors import OptimizeError
 from sqlglot.optimizer.scope import traverse_scope
 from sqlglot.schema import ensure_schema
 
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+    from sqlglot.schema import Schema
+    from sqlglot.dialects.dialect import DialectType
 
-def isolate_table_selects(expression, schema=None):
-    schema = ensure_schema(schema)
+
+def isolate_table_selects(
+    expression: E,
+    schema: t.Optional[t.Dict | Schema] = None,
+    dialect: DialectType = None,
+) -> E:
+    schema = ensure_schema(schema, dialect=dialect)
 
     for scope in traverse_scope(expression):
         if len(scope.selected_sources) == 1:
             continue
 
         for _, source in scope.selected_sources.values():
+            assert source.parent
+
             if (
                 not isinstance(source, exp.Table)
                 or not schema.column_names(source)
diff --git a/sqlglot/optimizer/pushdown_projections.py b/sqlglot/optimizer/pushdown_projections.py
index 6a3bc4ef1a..a315751fec 100644
--- a/sqlglot/optimizer/pushdown_projections.py
+++ b/sqlglot/optimizer/pushdown_projections.py
@@ -1,3 +1,6 @@
+from __future__ import annotations
+
+import typing as t
 from collections import defaultdict
 
 from sqlglot import alias, exp
@@ -7,6 +10,11 @@
 from sqlglot.errors import OptimizeError
 from sqlglot.helper import seq_get
 
+if t.TYPE_CHECKING:
+    from sqlglot._typing import E
+    from sqlglot.schema import Schema
+    from sqlglot.dialects.dialect import DialectType
+
 # Sentinel value that means an outer query selecting ALL columns
 SELECT_ALL = object()
 
@@ -16,7 +24,12 @@ def default_selection(is_agg: bool) -> exp.Alias:
     return alias(exp.Max(this=exp.Literal.number(1)) if is_agg else "1", "_")
 
 
-def pushdown_projections(expression, schema=None, remove_unused_selections=True):
+def pushdown_projections(
+    expression: E,
+    schema: t.Optional[t.Dict | Schema] = None,
+    remove_unused_selections: bool = True,
+    dialect: DialectType = None,
+) -> E:
     """
     Rewrite sqlglot AST to remove unused columns projections.
 
@@ -34,9 +47,9 @@ def pushdown_projections(expression, schema=None, remove_unused_selections=True)
         sqlglot.Expression: optimized expression
     """
     # Map of Scope to all columns being selected by outer queries.
-    schema = ensure_schema(schema)
-    source_column_alias_count = {}
-    referenced_columns = defaultdict(set)
+    schema = ensure_schema(schema, dialect=dialect)
+    source_column_alias_count: t.Dict[exp.Expression | Scope, int] = {}
+    referenced_columns: t.DefaultDict[Scope, t.Set[str | object]] = defaultdict(set)
 
     # We build the scope tree (which is traversed in DFS postorder), then iterate
     # over the result in reverse order. This should ensure that the set of selected
@@ -69,12 +82,12 @@ def pushdown_projections(expression, schema=None, remove_unused_selections=True)
                     if scope.expression.args.get("by_name"):
                         referenced_columns[right] = referenced_columns[left]
                     else:
-                        referenced_columns[right] = [
+                        referenced_columns[right] = {
                             right.expression.selects[i].alias_or_name
                             for i, select in enumerate(left.expression.selects)
                             if SELECT_ALL in parent_selections
                             or select.alias_or_name in parent_selections
-                        ]
+                        }
 
         if isinstance(scope.expression, exp.Select):
             if remove_unused_selections:
diff --git a/sqlglot/optimizer/qualify_columns.py b/sqlglot/optimizer/qualify_columns.py
index 8875c6d88e..dedd973d8f 100644
--- a/sqlglot/optimizer/qualify_columns.py
+++ b/sqlglot/optimizer/qualify_columns.py
@@ -23,6 +23,7 @@ def qualify_columns(
     expand_stars: bool = True,
     infer_schema: t.Optional[bool] = None,
     allow_partial_qualification: bool = False,
+    dialect: DialectType = None,
 ) -> exp.Expression:
     """
     Rewrite sqlglot AST to have fully qualified columns.
@@ -50,7 +51,7 @@ def qualify_columns(
     Notes:
         - Currently only handles a single PIVOT or UNPIVOT operator
     """
-    schema = ensure_schema(schema)
+    schema = ensure_schema(schema, dialect=dialect)
     annotator = TypeAnnotator(schema)
     infer_schema = schema.empty if infer_schema is None else infer_schema
     dialect = Dialect.get_or_raise(schema.dialect)


Test patch:
diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py
index f7e7537ae2..4f9307e66f 100644
--- a/tests/test_optimizer.py
+++ b/tests/test_optimizer.py
@@ -1552,3 +1552,19 @@ def gen_expr(depth: int) -> exp.Expression:
         self.assertEqual(4, normalization_distance(gen_expr(2), max_=100))
         self.assertEqual(18, normalization_distance(gen_expr(3), max_=100))
         self.assertEqual(110, normalization_distance(gen_expr(10), max_=100))
+
+    def test_manually_annotate_snowflake(self):
+        dialect = "snowflake"
+        schema = {
+            "SCHEMA": {
+                "TBL": {"COL": "INT", "col2": "VARCHAR"},
+            }
+        }
+        example_query = 'SELECT * FROM "SCHEMA"."TBL"'
+
+        expression = parse_one(example_query, dialect=dialect)
+        qual = optimizer.qualify.qualify(expression, schema=schema, dialect=dialect)
+        annotated = optimizer.annotate_types.annotate_types(qual, schema=schema, dialect=dialect)
+
+        self.assertTrue(annotated.selects[0].is_type("INT"))
+        self.assertTrue(annotated.selects[1].is_type("VARCHAR"))