fix: Address mechanical review findings (security, performance, code hygiene)

pyproject.toml

@@ -32,7 +32,6 @@ dependencies = [
     "sqlglot>=28.0.0",
     "graphviz>=0.20.0",
     "jinja2>=3.0.0",
-    "cloudpickle>=3.1.2",
 ]
 
 [project.optional-dependencies]

src/clgraph/__init__.py

@@ -5,10 +5,10 @@
 """
 
 try:
-    from importlib.metadata import version
+    from importlib.metadata import PackageNotFoundError, version
 
     __version__ = version("clgraph")
-except Exception:
+except PackageNotFoundError:
     # Fallback when package metadata is not available (e.g., mounted as volume)
     __version__ = "dev"
 

src/clgraph/agent.py

@@ -20,6 +20,7 @@
     print(result.data["sql"])
 """
 
+import logging
 import re
 from dataclasses import dataclass
 from enum import Enum
@@ -30,6 +31,8 @@
 if TYPE_CHECKING:
     from .pipeline import Pipeline
 
+logger = logging.getLogger(__name__)
+
 
 class QuestionType(Enum):
     """Types of questions the agent can handle."""
@@ -199,8 +202,7 @@ def query(self, question: str) -> AgentResult:
         # Classify question
         question_type = self._classify_question(question)
 
-        if self.verbose:
-            print(f"Question type: {question_type}")
+        logger.debug("Question type: %s", question_type)
 
         # Route to appropriate handler
         try:
@@ -225,6 +227,7 @@ def query(self, question: str) -> AgentResult:
             else:
                 return self._handle_general(question)
         except Exception as e:
+            logger.error("Agent query failed: %s", e, exc_info=True)
             return AgentResult(
                 answer=f"Sorry, I encountered an error: {e}",
                 question_type=question_type,
@@ -267,7 +270,9 @@ def _extract_table_column(self, question: str) -> tuple:
         """Extract table and column references from question."""
         # First, try to match against known table names (handles schema.table.column)
         # Sort by length descending to match longer table names first
-        known_tables = sorted(self.pipeline.table_graph.tables.keys(), key=len, reverse=True)
+        known_tables = sorted(
+            self.pipeline.table_graph.tables.keys(), key=lambda t: len(t), reverse=True
+        )
 
         for table_name in known_tables:
             # Check for table.column pattern
@@ -564,7 +569,7 @@ def _handle_sql_explain(self, question: str) -> AgentResult:
         result = self.registry.run("explain_query", sql=sql)
 
         return AgentResult(
-            answer=result.message if result.success else result.error,
+            answer=(result.message if result.success else result.error) or "",
             question_type=QuestionType.SQL_EXPLAIN,
             tool_used="explain_query",
             tool_result=result,

src/clgraph/column.py

@@ -6,8 +6,9 @@
 """
 
 import logging
+from collections import deque
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, Dict, List, Set
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set
 
 from .models import (
     ColumnEdge,
@@ -64,8 +65,9 @@ def generate_description(column: ColumnNode, llm: Any, pipeline: "Pipeline"):
 
         column.description = response.content.strip()
         column.description_source = DescriptionSource.GENERATED
-    except Exception:
-        # Fallback to simple rule-based description
+    except (ImportError, ValueError, AttributeError, RuntimeError) as e:
+        # Fallback to simple rule-based description if LLM fails
+        logger.debug("LLM description generation failed: %s", e)
         _generate_fallback_description(column)
 
 
@@ -233,6 +235,11 @@ class PipelineLineageGraph:
     edges: List[ColumnEdge] = field(default_factory=list)
     issues: List[ValidationIssue] = field(default_factory=list)  # Validation issues
 
+    # Adjacency indices: full_name -> list of edges
+    _outgoing_index: Dict[str, List[ColumnEdge]] = field(default_factory=dict, repr=False)
+    _incoming_index: Dict[str, List[ColumnEdge]] = field(default_factory=dict, repr=False)
+    _column_deps_cache: Optional[Dict[str, Set[str]]] = field(default=None, repr=False)
+
     def add_column(self, column: ColumnNode) -> ColumnNode:
         """Add a column node to the graph"""
         self.columns[column.full_name] = column
@@ -241,6 +248,9 @@ def add_column(self, column: ColumnNode) -> ColumnNode:
     def add_edge(self, edge: ColumnEdge):
         """Add a lineage edge"""
         self.edges.append(edge)
+        self._outgoing_index.setdefault(edge.from_node.full_name, []).append(edge)
+        self._incoming_index.setdefault(edge.to_node.full_name, []).append(edge)
+        self._column_deps_cache = None  # Invalidate cache
 
     def add_issue(self, issue: ValidationIssue):
         """Add a validation issue and log it"""
@@ -263,9 +273,14 @@ def _build_column_dependencies(self) -> Dict[str, Set[str]]:
         Build dependency map: column_full_name -> set of column_full_names it depends on.
         This is the column-level equivalent of TableDependencyGraph._build_table_dependencies.
 
+        Returns cached result when available; invalidated by add_edge().
+
         Returns:
             Dict mapping column full_name to set of upstream column full_names
         """
+        if self._column_deps_cache is not None:
+            return self._column_deps_cache
+
         deps: Dict[str, Set[str]] = {}
 
         for full_name in self.columns:
@@ -278,6 +293,7 @@ def _build_column_dependencies(self) -> Dict[str, Set[str]]:
             if to_name in deps:
                 deps[to_name].add(from_name)
 
+        self._column_deps_cache = deps
         return deps
 
     def get_upstream(self, full_name: str) -> List[ColumnNode]:
@@ -398,10 +414,10 @@ def to_simplified(self) -> "PipelineLineageGraph":
         for col_name, col in table_columns.items():
             # BFS backward to find all reachable table columns
             visited: Set[str] = set()
-            queue = [col_name]
+            queue = deque([col_name])
 
             while queue:
-                current = queue.pop(0)
+                current = queue.popleft()
                 if current in visited:
                     continue
                 visited.add(current)

src/clgraph/execution.py

@@ -6,10 +6,13 @@
 """
 
 import asyncio
+import logging
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from typing import TYPE_CHECKING, Any, Awaitable, Callable, Dict, List, Tuple
 
+logger = logging.getLogger(__name__)
+
 if TYPE_CHECKING:
     from .pipeline import Pipeline
 
@@ -118,9 +121,7 @@ def execute_sql(sql: str):
             result = executor.run(execute_sql, max_workers=4)
             print(f"Completed {len(result['completed'])} queries")
         """
-        if verbose:
-            print(f"🚀 Starting pipeline execution ({len(self.table_graph.queries)} queries)")
-            print()
+        logger.info("Starting pipeline execution (%d queries)", len(self.table_graph.queries))
 
         # Track completed queries
         completed = set()
@@ -132,8 +133,7 @@ def execute_sql(sql: str):
 
         # Execute level by level
         for level_num, level_queries in enumerate(levels, 1):
-            if verbose:
-                print(f"📊 Level {level_num}: {len(level_queries)} queries")
+            logger.info("Level %d: %d queries", level_num, len(level_queries))
 
             # Execute queries in this level concurrently
             with ThreadPoolExecutor(max_workers=max_workers) as pool:
@@ -151,31 +151,21 @@ def execute_sql(sql: str):
                     try:
                         future.result()
                         completed.add(query_id)
-
-                        if verbose:
-                            print(f"  ✅ {query_id}")
+                        logger.info("Completed: %s", query_id)
                     except Exception as e:
                         failed.append((query_id, str(e)))
-
-                        if verbose:
-                            print(f"  ❌ {query_id}: {e}")
-
-            if verbose:
-                print()
+                        logger.debug("Query %s execution failed", query_id, exc_info=True)
+                        logger.warning("Failed: %s: %s", query_id, e)
 
         elapsed = time.time() - start_time
 
         # Summary
-        if verbose:
-            print("=" * 60)
-            print(f"✅ Pipeline completed in {elapsed:.2f}s")
-            print(f"   Successful: {len(completed)}")
-            print(f"   Failed: {len(failed)}")
-            if failed:
-                print("\n⚠️  Failed queries:")
-                for query_id, error in failed:
-                    print(f"   - {query_id}: {error}")
-            print("=" * 60)
+        logger.info("Pipeline completed in %.2fs", elapsed)
+        logger.info("Successful: %d", len(completed))
+        logger.info("Failed: %d", len(failed))
+        if failed:
+            for query_id, error in failed:
+                logger.warning("Failed query - %s: %s", query_id, error)
 
         return {
             "completed": list(completed),
@@ -214,9 +204,7 @@ async def execute_sql(sql: str):
             result = await executor.async_run(execute_sql, max_workers=4)
             print(f"Completed {len(result['completed'])} queries")
         """
-        if verbose:
-            print(f"🚀 Starting async pipeline execution ({len(self.table_graph.queries)} queries)")
-            print()
+        logger.info("Starting async pipeline execution (%d queries)", len(self.table_graph.queries))
 
         # Track completed queries
         completed = set()
@@ -231,21 +219,19 @@ async def execute_sql(sql: str):
 
         # Execute level by level
         for level_num, level_queries in enumerate(levels, 1):
-            if verbose:
-                print(f"📊 Level {level_num}: {len(level_queries)} queries")
+            logger.info("Level %d: %d queries", level_num, len(level_queries))
 
             async def execute_with_semaphore(query_id: str, sql: str):
                 """Execute query with semaphore for concurrency control"""
                 async with semaphore:
                     try:
                         await executor(sql)
                         completed.add(query_id)
-                        if verbose:
-                            print(f"  ✅ {query_id}")
+                        logger.info("Completed: %s", query_id)
                     except Exception as e:
                         failed.append((query_id, str(e)))
-                        if verbose:
-                            print(f"  ❌ {query_id}: {e}")
+                        logger.debug("Async query %s execution failed", query_id, exc_info=True)
+                        logger.warning("Failed: %s: %s", query_id, e)
 
             # Execute queries in this level concurrently
             tasks = []
@@ -257,22 +243,15 @@ async def execute_with_semaphore(query_id: str, sql: str):
             # Wait for all tasks in this level to complete
             await asyncio.gather(*tasks)
 
-            if verbose:
-                print()
-
         elapsed = time.time() - start_time
 
         # Summary
-        if verbose:
-            print("=" * 60)
-            print(f"✅ Pipeline completed in {elapsed:.2f}s")
-            print(f"   Successful: {len(completed)}")
-            print(f"   Failed: {len(failed)}")
-            if failed:
-                print("\n⚠️  Failed queries:")
-                for query_id, error in failed:
-                    print(f"   - {query_id}: {error}")
-            print("=" * 60)
+        logger.info("Pipeline completed in %.2fs", elapsed)
+        logger.info("Successful: %d", len(completed))
+        logger.info("Failed: %d", len(failed))
+        if failed:
+            for query_id, error in failed:
+                logger.warning("Failed query - %s: %s", query_id, error)
 
         return {
             "completed": list(completed),

src/clgraph/lineage_builder.py

@@ -5,6 +5,7 @@
 Includes SQLColumnTracer wrapper for backward compatibility.
 """
 
+from collections import deque
 from typing import Any, Dict, List, Optional, Set, Tuple, TypedDict
 
 import sqlglot
@@ -558,7 +559,7 @@ def _qualify_sql_with_schema(
         # Return the qualified SQL
         return qualified.sql(dialect=dialect)
 
-    except Exception:
+    except (sqlglot.errors.SqlglotError, KeyError, ValueError, TypeError):
         # If qualification fails, return original SQL
         # The lineage builder will handle unqualified columns as before
         return sql_query
@@ -982,7 +983,7 @@ def _create_lateral_correlation_edges(self, unit: QueryUnit):
         For each correlated column (reference to outer table), create an edge
         showing the correlation relationship.
         """
-        lateral_alias = unit.name
+        lateral_alias = unit.name or ""
 
         for correlated_col in unit.correlated_columns:
             # Parse table.column format
@@ -1505,7 +1506,7 @@ def _extract_columns_from_expr(
                     )
                 col_name = col.name
                 result.append((table_ref, col_name))
-        except Exception:
+        except (sqlglot.errors.SqlglotError, ValueError, TypeError):
             # If parsing fails, try simple extraction for "table.column" format
             if "." in expr_str:
                 parts = expr_str.split(".")
@@ -2650,7 +2651,7 @@ def _parse_aggregate_spec(self, ast_node: Optional[exp.Expression]) -> Optional[
 
         # Extract ORDER BY within aggregate (fallback for standard syntax)
         if not order_by and hasattr(agg_func, "order") and agg_func.order:
-            for order_expr in agg_func.order.expressions:
+            for order_expr in agg_func.order.expressions:  # type: ignore[union-attr]
                 col_name = ""
                 direction = "asc"
                 nulls = None
@@ -2723,7 +2724,7 @@ def _get_aggregate_func_name(self, node: exp.Expression) -> str:
         elif isinstance(node, exp.Max):
             return "MAX"
         elif hasattr(node, "sql_name"):
-            return node.sql_name().upper()
+            return node.sql_name().upper()  # type: ignore[union-attr]
         elif hasattr(node, "name") and node.name:
             return node.name.upper()
         return "AGGREGATE"
@@ -3096,7 +3097,7 @@ def _validate_qualified_columns_in_joins(
                             category=IssueCategory.UNQUALIFIED_COLUMN,
                             message=(
                                 f"Unqualified column '{col_name}' in expression for '{output_col_name}'. "
-                                f"With multiple tables ({', '.join(available_tables)}), "
+                                f"With multiple tables ({', '.join(str(t) for t in available_tables)}), "
                                 f"the source table is ambiguous."
                             ),
                             query_id=self.query_id,
@@ -3219,10 +3220,10 @@ def get_forward_lineage(self, input_columns: List[str]) -> Dict[str, Any]:
             # BFS forward from each start node
             for start_node in start_nodes:
                 visited = set()
-                queue = [(start_node, [start_node.full_name], [])]
+                queue = deque([(start_node, [start_node.full_name], [])])
 
                 while queue:
-                    current, path, transformations = queue.pop(0)
+                    current, path, transformations = queue.popleft()
 
                     if current.full_name in visited:
                         continue
@@ -3290,10 +3291,10 @@ def get_backward_lineage(self, output_columns: List[str]) -> BackwardLineageResu
             # BFS backward from each start node
             for start_node in start_nodes:
                 visited = set()
-                queue = [(start_node, [start_node.full_name], [])]
+                queue = deque([(start_node, [start_node.full_name], [])])
 
                 while queue:
-                    current, path, transformations = queue.pop(0)
+                    current, path, transformations = queue.popleft()
 
                     if current.full_name in visited:
                         continue