sql-glider 0.1.18__tar.gz → 0.1.20__tar.gz

{sql_glider-0.1.18 → sql_glider-0.1.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-glider
-Version: 0.1.18
+Version: 0.1.20
 Summary: SQL Utility Toolkit for better understanding, use, and governance of your queries in a native environment.
 Project-URL: Homepage, https://github.com/rycowhi/sql-glider/
 Project-URL: Repository, https://github.com/rycowhi/sql-glider/

{sql_glider-0.1.18 → sql_glider-0.1.20}/src/sqlglider/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.18'
-__version_tuple__ = version_tuple = (0, 1, 18)
+__version__ = version = '0.1.20'
+__version_tuple__ = version_tuple = (0, 1, 20)
 
 __commit_id__ = commit_id = None

{sql_glider-0.1.18 → sql_glider-0.1.20}/src/sqlglider/lineage/analyzer.py

@@ -1480,7 +1480,7 @@ class LineageAnalyzer:
 
     def _extract_schema_from_statement(self, expr: exp.Expression) -> None:
         """
-        Extract column definitions from CREATE VIEW/TABLE AS SELECT statements.
+        Extract column definitions from CREATE VIEW/TABLE AS SELECT and CACHE TABLE statements.
 
         This method builds up file-scoped schema context as statements are processed,
         enabling SQLGlot to correctly expand SELECT * and trace cross-statement references.
@@ -1488,29 +1488,39 @@ class LineageAnalyzer:
         Args:
             expr: The SQL expression to extract schema from
         """
-        # Only handle CREATE VIEW or CREATE TABLE (AS SELECT)
-        if not isinstance(expr, exp.Create):
-            return
-        if expr.kind not in ("VIEW", "TABLE"):
-            return
+        target_name: str | None = None
+        select_node: exp.Expression | None = None
 
-        # Get target table/view name
-        target = expr.this
-        if isinstance(target, exp.Schema):
-            target = target.this
-        if not isinstance(target, exp.Table):
-            return
+        if isinstance(expr, exp.Create):
+            if expr.kind not in ("VIEW", "TABLE"):
+                return
+
+            target = expr.this
+            if isinstance(target, exp.Schema):
+                target = target.this
+            if not isinstance(target, exp.Table):
+                return
+
+            target_name = self._get_qualified_table_name(target)
+            select_node = expr.expression
+
+            # Handle Subquery wrapper (e.g., CREATE VIEW AS (SELECT ...))
+            if isinstance(select_node, exp.Subquery):
+                select_node = select_node.this
+
+        elif isinstance(expr, exp.Cache):
+            target = expr.this
+            if not isinstance(target, exp.Table):
+                return
 
-        target_name = self._get_qualified_table_name(target)
+            target_name = self._get_qualified_table_name(target)
+            select_node = expr.expression
 
-        # Get the SELECT node from the CREATE statement
-        select_node = expr.expression
-        if select_node is None:
+        else:
             return
 
-        # Handle Subquery wrapper (e.g., CREATE VIEW AS (SELECT ...))
-        if isinstance(select_node, exp.Subquery):
-            select_node = select_node.this
+        if target_name is None or select_node is None:
+            return
 
         if not isinstance(
             select_node, (exp.Select, exp.Union, exp.Intersect, exp.Except)
@@ -1822,7 +1832,13 @@ class LineageAnalyzer:
                 if actual_name in self._file_schema:
                     return list(self._file_schema[actual_name].keys())
 
-            # Check JOIN clauses for aliased tables
+            # Check if it's an aliased subquery (e.g., FROM (SELECT ...) sub)
+            if isinstance(source, exp.Subquery) and source.alias == table_name:
+                inner_select = source.this
+                if isinstance(inner_select, exp.Select):
+                    return self._extract_subquery_columns(inner_select)
+
+            # Check JOIN clauses for aliased tables and subqueries
             joins = select_node.args.get("joins")
             if joins:
                 for join in joins:
@@ -1835,6 +1851,13 @@ class LineageAnalyzer:
                             actual_name = self._get_qualified_table_name(join_source)
                             if actual_name in self._file_schema:
                                 return list(self._file_schema[actual_name].keys())
+                        if (
+                            isinstance(join_source, exp.Subquery)
+                            and join_source.alias == table_name
+                        ):
+                            inner_select = join_source.this
+                            if isinstance(inner_select, exp.Select):
+                                return self._extract_subquery_columns(inner_select)
 
         return []
 

{sql_glider-0.1.18 → sql_glider-0.1.20}/tests/sqlglider/lineage/test_analyzer.py

@@ -2949,6 +2949,161 @@ class TestCacheTableStatements:
         assert len(skipped) == 1
         assert "DELETE" in skipped[0].statement_type
 
+    def test_cache_table_star_resolution_in_subsequent_query(self):
+        """SELECT * FROM a cached table should resolve columns from the CACHE statement."""
+        sql = """
+        CACHE TABLE cached_orders AS SELECT customer_id, order_total FROM orders;
+        SELECT * FROM cached_orders;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 2
+        # Second query should have resolved the star
+        star_result = results[1]
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        assert output_names == ["customer_id", "order_total"]
+
+    def test_cache_table_qualified_star_resolution(self):
+        """table.* on a cached table should resolve columns."""
+        sql = """
+        CACHE TABLE cached_orders AS SELECT customer_id, order_total FROM orders;
+        SELECT cached_orders.* FROM cached_orders;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 2
+        star_result = results[1]
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        assert output_names == ["customer_id", "order_total"]
+
+    def test_cache_lazy_table_star_resolution(self):
+        """CACHE LAZY TABLE should also register schema for star resolution."""
+        sql = """
+        CACHE LAZY TABLE cached_users AS SELECT id, name, email FROM users;
+        SELECT * FROM cached_users;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 2
+        star_result = results[1]
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        assert output_names == ["email", "id", "name"]
+
+    def test_cache_table_schema_registered(self):
+        """Cached table schema should be registered for downstream star resolution."""
+        sql = """
+        CACHE TABLE cached_orders AS SELECT customer_id, order_total FROM orders;
+        SELECT * FROM cached_orders;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        schema = analyzer.extract_schema_only()
+        assert "cached_orders" in schema
+        assert "customer_id" in schema["cached_orders"]
+        assert "order_total" in schema["cached_orders"]
+
+
+class TestSubqueryAliasStarResolution:
+    """Tests for resolving qualified star (sub.*) on aliased subqueries."""
+
+    def test_subquery_alias_star_in_from(self):
+        """sub.* on an aliased subquery in FROM should resolve columns."""
+        sql = """
+        SELECT sub.*
+        FROM (SELECT id, SUM(amount) AS total FROM orders GROUP BY id) sub
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 1
+        output_names = sorted(item.output_name for item in results[0].lineage_items)
+        assert output_names == ["id", "total"]
+
+    def test_subquery_alias_star_in_join(self):
+        """sub.* on an aliased subquery in JOIN should resolve columns."""
+        sql = """
+        SELECT t.id, sub.*
+        FROM db.schema.users t
+        JOIN (SELECT user_id, SUM(amount) AS total FROM db.schema.orders GROUP BY user_id) sub
+        ON t.id = sub.user_id
+        """
+        analyzer = LineageAnalyzer(
+            sql,
+            dialect="spark",
+            schema={"db.schema.users": {"id": "INT"}},
+        )
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 1
+        output_names = sorted(item.output_name for item in results[0].lineage_items)
+        assert output_names == ["db.schema.users.id", "total", "user_id"]
+
+    def test_subquery_alias_star_with_generated_columns(self):
+        """sub.* plus generated columns should all resolve."""
+        sql = """
+        SELECT
+            sub.*,
+            sub.total * 0.1 AS commission
+        FROM (SELECT id, SUM(amount) AS total FROM orders GROUP BY id) sub
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 1
+        output_names = sorted(item.output_name for item in results[0].lineage_items)
+        assert output_names == ["commission", "id", "total"]
+
+    def test_cache_table_with_subquery_alias_star(self):
+        """CACHE TABLE using sub.* should register all columns for downstream use."""
+        sql = """
+        CACHE TABLE cached_result AS
+        SELECT
+            sub.*,
+            d.dept_name
+        FROM (
+            SELECT id, SUM(amount) AS total
+            FROM db.schema.orders
+            GROUP BY id
+        ) sub
+        JOIN db.schema.departments d ON sub.id = d.id;
+
+        SELECT * FROM cached_result;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 2
+        # Second query should resolve all columns from the cached table
+        star_result = results[1]
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        assert output_names == ["dept_name", "id", "total"]
+
+    def test_temp_view_with_subquery_alias_star(self):
+        """CREATE TEMP VIEW using sub.* should register all columns for downstream use."""
+        sql = """
+        CREATE TEMPORARY VIEW enriched AS
+        SELECT
+            sub.*,
+            sub.total * 0.1 AS commission
+        FROM (
+            SELECT id, SUM(amount) AS total
+            FROM db.schema.orders
+            GROUP BY id
+        ) sub;
+
+        SELECT * FROM enriched;
+        """
+        analyzer = LineageAnalyzer(sql, dialect="spark")
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+
+        assert len(results) == 2
+        star_result = results[1]
+        output_names = sorted(item.output_name for item in star_result.lineage_items)
+        assert output_names == ["commission", "id", "total"]
+
 
 class TestNoStar:
     """Tests for the --no-star flag that fails on unresolvable SELECT *."""