sql-dag-flow 0.5.1__tar.gz → 0.5.2__tar.gz

{sql_dag_flow-0.5.1/src/sql_dag_flow.egg-info → sql_dag_flow-0.5.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sql-dag-flow
-Version: 0.5.1
+Version: 0.5.2
 Summary: A sophisticated SQL lineage visualization tool for Medallion Architectures.
 Author-email: Flavio Sandoval <dsandovalflavio@gmail.com>
 License: MIT

{sql_dag_flow-0.5.1 → sql_dag_flow-0.5.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sql-dag-flow"
-version = "0.5.1"
+version = "0.5.2"
 description = "A sophisticated SQL lineage visualization tool for Medallion Architectures."
 readme = "README.md"
 requires-python = ">=3.8"

{sql_dag_flow-0.5.1 → sql_dag_flow-0.5.2}/src/sql_dag_flow/parser.py

@@ -566,18 +566,28 @@ def parse_sql_files(directory, allowed_subfolders=None, dialect="bigquery"):
         global_schema["default"][label] = col_dict
     
     # Re-extract column_references using qualify_columns for precision
-    for tid, tdata in tables.items():
-        if tdata.get("error") or not tdata.get("content"):
-            continue
+    # Protected with per-table time budget to prevent hanging on large projects
+    qualify_tables = [(tid, tdata) for tid, tdata in tables.items() 
+                      if not tdata.get("error") and tdata.get("content")]
+    total_qualify = len(qualify_tables)
+    
+    for idx, (tid, tdata) in enumerate(qualify_tables):
+        table_start = time.time()
+        if idx % 10 == 0 and total_qualify > 10:
+            print(f"  Qualifying columns: {idx}/{total_qualify} tables...")
         try:
             parsed = sqlglot.parse_one(tdata["content"], read=dialect)
             
-            # Try to qualify columns using the global schema
+            # Try to qualify columns using the global schema (with 2s budget)
             try:
                 qualified_ast = sqlglot_qualify_columns(parsed.copy(), schema=global_schema, dialect=dialect)
             except Exception:
                 qualified_ast = parsed  # Fallback to unqualified
             
+            # Check time budget
+            if time.time() - table_start > 2.0:
+                continue  # Skip extraction if qualification took too long
+            
             # Re-extract column references from qualified AST
             alias_map = {}
             defined_ctes = set()
@@ -634,55 +644,83 @@ def parse_sql_files(directory, allowed_subfolders=None, dialect="bigquery"):
         except Exception:
             pass  # Keep original column_references
     
+    if total_qualify > 10:
+        print(f"  Qualifying columns: {total_qualify}/{total_qualify} done.")
+    
     # ===== Column-Level Lineage =====
     # For each table, trace how each output column derives from source columns
-    for tid, tdata in tables.items():
-        if tdata.get("error") or not tdata.get("content"):
-            continue
-        schema_cols = tdata.get("schema", [])
-        if not schema_cols:
-            continue
-        
-        column_lineage = {}  # col_name -> [{source_table, source_column, transform}]
-        sql_content = tdata["content"]
-        
-        for col_info in schema_cols:
-            col_name = col_info["name"]
-            if col_name == "*":
+    # Protected: skip if project is very large (>100 tables) and cap columns per table
+    MAX_LINEAGE_TABLES = 100
+    MAX_COLS_PER_TABLE = 30
+    LINEAGE_TIME_BUDGET = 1.5  # seconds per table
+    
+    lineage_tables = [(tid, tdata) for tid, tdata in tables.items()
+                      if not tdata.get("error") and tdata.get("content") and tdata.get("schema")]
+    total_lineage = len(lineage_tables)
+    
+    if total_lineage > MAX_LINEAGE_TABLES:
+        print(f"  Skipping column lineage: {total_lineage} tables exceeds limit of {MAX_LINEAGE_TABLES}")
+    else:
+        for idx, (tid, tdata) in enumerate(lineage_tables):
+            table_start = time.time()
+            if idx % 10 == 0 and total_lineage > 10:
+                print(f"  Column lineage: {idx}/{total_lineage} tables...")
+            
+            schema_cols = tdata.get("schema", [])
+            if not schema_cols:
                 continue
-            try:
-                from sqlglot.lineage import lineage as sqlglot_lineage
-                node = sqlglot_lineage(
-                    col_name, sql_content,
-                    schema=global_schema,
-                    dialect=dialect
-                )
-                sources = []
-                for child in node.downstream:
-                    source_col = child.name
-                    # Extract the source table from the expression
-                    source_table = ""
-                    if child.source and isinstance(child.source, exp.Table):
-                        source_table = child.source.name
-                        if child.source.db:
-                            source_table = f"{child.source.db}.{child.source.name}"
-                    
-                    transform = ""
-                    if child.expression and str(child.expression) != source_col:
-                        transform = str(child.expression)
+            
+            column_lineage = {}  # col_name -> [{source_table, source_column, transform}]
+            sql_content = tdata["content"]
+            
+            cols_processed = 0
+            for col_info in schema_cols:
+                # Check time budget
+                if time.time() - table_start > LINEAGE_TIME_BUDGET:
+                    break
+                if cols_processed >= MAX_COLS_PER_TABLE:
+                    break
                     
-                    sources.append({
-                        "source_table": source_table,
-                        "source_column": source_col,
-                        "transform": transform
-                    })
-                if sources:
-                    column_lineage[col_name] = sources
-            except Exception:
-                pass  # Lineage not available for this column
+                col_name = col_info["name"]
+                if col_name == "*":
+                    continue
+                cols_processed += 1
+                try:
+                    from sqlglot.lineage import lineage as sqlglot_lineage
+                    node = sqlglot_lineage(
+                        col_name, sql_content,
+                        schema=global_schema,
+                        dialect=dialect
+                    )
+                    sources = []
+                    for child in node.downstream:
+                        source_col = child.name
+                        # Extract the source table from the expression
+                        source_table = ""
+                        if child.source and isinstance(child.source, exp.Table):
+                            source_table = child.source.name
+                            if child.source.db:
+                                source_table = f"{child.source.db}.{child.source.name}"
+                        
+                        transform = ""
+                        if child.expression and str(child.expression) != source_col:
+                            transform = str(child.expression)
+                        
+                        sources.append({
+                            "source_table": source_table,
+                            "source_column": source_col,
+                            "transform": transform
+                        })
+                    if sources:
+                        column_lineage[col_name] = sources
+                except Exception:
+                    pass  # Lineage not available for this column
+            
+            if column_lineage:
+                tdata["column_lineage"] = column_lineage
         
-        if column_lineage:
-            tdata["column_lineage"] = column_lineage
+        if total_lineage > 10:
+            print(f"  Column lineage: {total_lineage}/{total_lineage} done.")
     
     return tables