Add a query provider to the expand function, enabling lazy loading of queries during AST transformations. (#4872)

* add a query provider to the expand function to load queries lazily

* change laziness implementation to use callables as source providers to simplify the usage

* fix linting issues

Author: DannyMor

sqlglot/expressions.py

@@ -8494,7 +8494,7 @@ def _replace_placeholders(node: Expression, args, **kwargs) -> Expression:
 
 def expand(
     expression: Expression,
-    sources: t.Dict[str, Query],
+    sources: t.Dict[str, Query] | t.Callable[[str], t.Optional[Query]],
     dialect: DialectType = None,
     copy: bool = True,
 ) -> Expression:
@@ -8510,22 +8510,36 @@ def expand(
 
     Args:
         expression: The expression to expand.
-        sources: A dictionary of name to Queries.
-        dialect: The dialect of the sources dict.
+        sources: A dict of name to query or a callable that provides a query on demand.
+        dialect: The dialect of the sources dict or the callable.
         copy: Whether to copy the expression during transformation. Defaults to True.
 
     Returns:
         The transformed expression.
     """
-    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 = sources.get(name)
+            source = get_source(name)
             if source:
+                # Create a subquery with the same alias (or table name if no alias)
                 subquery = source.subquery(node.alias or name)
                 subquery.comments = [f"source: {name}"]
+                # Continue expanding within the subquery
                 return subquery.transform(_expand, copy=False)
         return node
 

tests/test_expressions.py

@@ -1,9 +1,11 @@
 import sys
 import datetime
 import math
+import typing as t
 import unittest
 
 from sqlglot import ParseError, alias, exp, parse_one
+from sqlglot.expressions import Query, normalize_table_name
 
 
 class TestExpressions(unittest.TestCase):
@@ -277,6 +279,28 @@ def test_expand(self):
             "SELECT * FROM (SELECT 1) AS a /* source: a-b.c */",
         )
 
+    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),
+                dialect="spark",
+            ).sql(),
+            "SELECT * FROM (SELECT 1) AS a /* source: a-b.c */",
+        )
+
     def test_replace_placeholders(self):
         self.assertEqual(
             exp.replace_placeholders(