sql-glider 0.1.15__tar.gz → 0.1.16__tar.gz

{sql_glider-0.1.15 → sql_glider-0.1.16}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-glider
-Version: 0.1.15
+Version: 0.1.16
 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.15 → sql_glider-0.1.16}/src/sqlglider/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '0.1.15'
-__version_tuple__ = version_tuple = (0, 1, 15)
+__version__ = version = '0.1.16'
+__version_tuple__ = version_tuple = (0, 1, 16)
 
 __commit_id__ = commit_id = None

{sql_glider-0.1.15 → sql_glider-0.1.16}/src/sqlglider/lineage/analyzer.py

@@ -11,6 +11,48 @@ from sqlglot.lineage import Node, lineage
 from sqlglider.global_models import AnalysisLevel
 
 
+def _flat_schema_to_nested(
+    schema: Dict[str, Dict[str, str]],
+) -> Dict[str, object]:
+    """Convert flat dot-notation schema keys to the nested dict structure sqlglot expects.
+
+    sqlglot's MappingSchema requires consistent nesting depth across all tables.
+    Flat keys like ``"db.table"`` are split on dots and nested accordingly.
+    Shorter keys are padded with empty-string prefixes to match the max depth.
+
+    Examples::
+
+        {"users": {"id": "UNKNOWN"}}
+        → {"users": {"id": "UNKNOWN"}}  (depth 1, no change)
+
+        {"db.users": {"id": "UNKNOWN"}, "my_view": {"x": "UNKNOWN"}}
+        → {"db": {"users": {"id": "UNKNOWN"}}, "": {"my_view": {"x": "UNKNOWN"}}}
+    """
+    if not schema:
+        return {}
+
+    # Split all keys into parts
+    entries = [(key.split("."), cols) for key, cols in schema.items()]
+    max_depth = max(len(parts) for parts, _ in entries)
+
+    # If all keys are single-part (unqualified), return as-is
+    if max_depth == 1:
+        return schema  # type: ignore[return-value]
+
+    # Pad shorter keys with empty-string prefixes to match max depth
+    nested: Dict[str, object] = {}
+    for parts, cols in entries:
+        while len(parts) < max_depth:
+            parts.insert(0, "")
+        d: Dict[str, object] = nested
+        for part in parts[:-1]:
+            if part not in d:
+                d[part] = {}
+            d = d[part]  # type: ignore[assignment]
+        d[parts[-1]] = cols
+    return nested
+
+
 class StarResolutionError(Exception):
     """Raised when SELECT * cannot be resolved and no_star mode is enabled."""
 
@@ -860,8 +902,10 @@ class LineageAnalyzer:
             current_query_sql = self.expr.sql(dialect=self.dialect)
 
         # Prune schema to only tables referenced in this query to avoid
-        # sqlglot.lineage() performance degradation with large schema dicts
-        pruned_schema: Optional[Dict[str, Dict[str, str]]] = None
+        # sqlglot.lineage() performance degradation with large schema dicts.
+        # Then convert from flat dot-notation keys to the nested dict structure
+        # that sqlglot's MappingSchema expects.
+        lineage_schema: Optional[Dict[str, object]] = None
         if self._file_schema:
             referenced = {t.lower() for t in self._get_query_tables()}
             pruned_schema = {
@@ -869,8 +913,8 @@ class LineageAnalyzer:
                 for table, cols in self._file_schema.items()
                 if table.lower() in referenced
             }
-            if not pruned_schema:
-                pruned_schema = None
+            if pruned_schema:
+                lineage_schema = _flat_schema_to_nested(pruned_schema)
 
         for col in columns_to_analyze:
             try:
@@ -883,7 +927,7 @@ class LineageAnalyzer:
                     lineage_col,
                     current_query_sql,
                     dialect=self.dialect,
-                    schema=pruned_schema,
+                    schema=lineage_schema,
                 )
 
                 # Collect all source columns

{sql_glider-0.1.15 → sql_glider-0.1.16}/tests/sqlglider/lineage/test_analyzer.py

@@ -3,7 +3,11 @@
 import pytest
 
 from sqlglider.global_models import AnalysisLevel
-from sqlglider.lineage.analyzer import LineageAnalyzer, StarResolutionError
+from sqlglider.lineage.analyzer import (
+    LineageAnalyzer,
+    StarResolutionError,
+    _flat_schema_to_nested,
+)
 
 
 class TestCaseInsensitiveForwardLineage:
@@ -3181,3 +3185,97 @@ class TestSchemaPruning:
         output_names = {item.output_name for r in results for item in r.lineage_items}
         assert "id" in output_names
         assert "email" in output_names
+
+
+class TestFlatSchemaToNested:
+    """Tests for _flat_schema_to_nested conversion utility."""
+
+    def test_empty(self):
+        assert _flat_schema_to_nested({}) == {}
+
+    def test_unqualified_passthrough(self):
+        schema = {"users": {"id": "UNKNOWN"}}
+        assert _flat_schema_to_nested(schema) == schema
+
+    def test_two_part_keys(self):
+        schema = {"db.users": {"id": "UNKNOWN"}}
+        result = _flat_schema_to_nested(schema)
+        assert result == {"db": {"users": {"id": "UNKNOWN"}}}
+
+    def test_three_part_keys(self):
+        schema = {"cat.db.users": {"id": "UNKNOWN"}}
+        result = _flat_schema_to_nested(schema)
+        assert result == {"cat": {"db": {"users": {"id": "UNKNOWN"}}}}
+
+    def test_mixed_depth_pads_shorter_keys(self):
+        schema = {
+            "my_view": {"x": "UNKNOWN"},
+            "db.users": {"id": "UNKNOWN"},
+        }
+        result = _flat_schema_to_nested(schema)
+        assert result == {
+            "": {"my_view": {"x": "UNKNOWN"}},
+            "db": {"users": {"id": "UNKNOWN"}},
+        }
+
+
+class TestQualifiedSchemaKeys:
+    """Tests for schema with qualified (dotted) table names."""
+
+    def test_qualified_star_expansion(self):
+        """SELECT * resolves correctly with qualified schema keys."""
+        sql = "SELECT * FROM mydb.users"
+        schema = {"mydb.users": {"id": "UNKNOWN", "name": "UNKNOWN"}}
+        analyzer = LineageAnalyzer(sql, dialect="spark", schema=schema)
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+        items = {
+            (item.source_name, item.output_name)
+            for r in results
+            for item in r.lineage_items
+        }
+        assert ("mydb.users.id", "id") in items
+        assert ("mydb.users.name", "name") in items
+
+    def test_qualified_explicit_columns(self):
+        """Explicit columns trace sources correctly with qualified schema keys."""
+        sql = "SELECT id, name FROM mydb.users"
+        schema = {"mydb.users": {"id": "UNKNOWN", "name": "UNKNOWN"}}
+        analyzer = LineageAnalyzer(sql, dialect="spark", schema=schema)
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+        items = {
+            (item.source_name, item.output_name)
+            for r in results
+            for item in r.lineage_items
+        }
+        assert ("mydb.users.id", "mydb.users.id") in items
+        assert ("mydb.users.name", "mydb.users.name") in items
+
+    def test_three_part_qualified(self):
+        """3-part qualified names (catalog.db.table) work correctly."""
+        sql = "SELECT id FROM catalog.mydb.users"
+        schema = {"catalog.mydb.users": {"id": "UNKNOWN"}}
+        analyzer = LineageAnalyzer(sql, dialect="spark", schema=schema)
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+        items = [
+            (item.source_name, item.output_name)
+            for r in results
+            for item in r.lineage_items
+        ]
+        assert len(items) == 1
+        assert items[0] == ("catalog.mydb.users.id", "catalog.mydb.users.id")
+
+    def test_mixed_qualified_and_unqualified(self):
+        """Mix of qualified and unqualified table names in schema."""
+        sql = "SELECT * FROM my_view"
+        schema = {
+            "my_view": {"id": "UNKNOWN"},
+            "mydb.users": {"id": "UNKNOWN", "name": "UNKNOWN"},
+        }
+        analyzer = LineageAnalyzer(sql, dialect="spark", schema=schema)
+        results = analyzer.analyze_queries(level=AnalysisLevel.COLUMN)
+        items = {
+            (item.source_name, item.output_name)
+            for r in results
+            for item in r.lineage_items
+        }
+        assert ("my_view.id", "id") in items