fix: expand util - align normalization behaviour with lazy and non-lazy source providers. (#4874)

* Fix sources usage

* Format

Author: omerhadari

sqlglot/expressions.py

@@ -11,6 +11,7 @@
 """
 
 from __future__ import annotations
+
 import datetime
 import math
 import numbers
@@ -37,6 +38,7 @@
 
 if t.TYPE_CHECKING:
     from typing_extensions import Self
+
     from sqlglot._typing import E, Lit
     from sqlglot.dialects.dialect import DialectType
 
@@ -8494,7 +8496,7 @@ def _replace_placeholders(node: Expression, args, **kwargs) -> Expression:
 
 def expand(
     expression: Expression,
-    sources: t.Dict[str, Query] | t.Callable[[str], t.Optional[Query]],
+    sources: t.Dict[str, Query | t.Callable[[], Query]],
     dialect: DialectType = None,
     copy: bool = True,
 ) -> Expression:
@@ -8517,27 +8519,17 @@ def expand(
     Returns:
         The transformed expression.
     """
+    normalized_sources = {normalize_table_name(k, dialect=dialect): v for k, v in sources.items()}
     # Create a query provider based on the sources parameter
-    if callable(sources):
-        get_source = sources
-    else:
-        # Pre-normalize table names in sources dictionary for consistent lookups
-        normalized_sources = {
-            normalize_table_name(k, dialect=dialect): v for k, v in sources.items()
-        }
-
-        def _get_source(name: str) -> t.Optional[Query]:
-            return normalized_sources.get(name)
-
-        get_source = _get_source
 
     def _expand(node: Expression):
         if isinstance(node, Table):
             name = normalize_table_name(node, dialect=dialect)
-            source = get_source(name)
+            source = normalized_sources.get(name)
             if source:
                 # Create a subquery with the same alias (or table name if no alias)
-                subquery = source.subquery(node.alias or name)
+                parsed_source = source() if callable(source) else source
+                subquery = parsed_source.subquery(node.alias or name)
                 subquery.comments = [f"source: {name}"]
                 # Continue expanding within the subquery
                 return subquery.transform(_expand, copy=False)

tests/test_expressions.py

@@ -1,11 +1,9 @@
-import sys
 import datetime
 import math
-import typing as t
+import sys
 import unittest
 
 from sqlglot import ParseError, alias, exp, parse_one
-from sqlglot.expressions import Query, normalize_table_name
 
 
 class TestExpressions(unittest.TestCase):
@@ -280,22 +278,10 @@ def test_expand(self):
         )
 
     def test_expand_with_lazy_source_provider(self):
-        class DynamicSourceProvider:
-            def __init__(self, dialect: str = "spark"):
-                self._sources = {normalize_table_name("`a-b`.`c`", dialect): "select 1"}
-
-            def get(self, name: str) -> t.Optional[Query]:
-                query_sql = self._sources.get(name)
-
-                if query_sql:
-                    return parse_one(query_sql)
-
-        dynamic_source_provider = DynamicSourceProvider()
-
         self.assertEqual(
             exp.expand(
                 parse_one('select * from "a-b"."C" AS a'),
-                lambda name: dynamic_source_provider.get(name),
+                {"`a-b`.c": lambda: parse_one("select 1", dialect="spark")},
                 dialect="spark",
             ).sql(),
             "SELECT * FROM (SELECT 1) AS a /* source: a-b.c */",
@@ -862,6 +848,7 @@ def test_properties_from_dict(self):
 
     def test_convert(self):
         from collections import namedtuple
+
         import pytz
 
         PointTuple = namedtuple("Point", ["x", "y"])