truthound-dashboard 1.4.3__py3-none-any.whl → 1.5.0__py3-none-any.whl

truthound_dashboard/core/lineage.py

@@ -2,10 +2,16 @@
 
 This module provides services for managing data lineage graphs,
 including node and edge CRUD, impact analysis, and auto-discovery.
+
+Uses truthound.lineage for advanced features when available:
+- LineageTracker: Automatic operation tracking
+- ImpactAnalyzer: What-if analysis and schema change impact
+- OpenLineage integration for industry-standard lineage events
 """
 
 from __future__ import annotations
 
+import logging
 from collections.abc import Sequence
 from datetime import datetime
 from typing import Any, Literal
@@ -16,6 +22,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
 from truthound_dashboard.db import BaseRepository
 from truthound_dashboard.db.models import AnomalyDetection, LineageEdge, LineageNode, Source
 
+logger = logging.getLogger(__name__)
+
 
 class LineageNodeRepository(BaseRepository[LineageNode]):
     """Repository for LineageNode model operations."""
@@ -73,6 +81,25 @@ class LineageNodeRepository(BaseRepository[LineageNode]):
             filters=[LineageNode.node_type == node_type],
         )
 
+    async def get_by_name_and_type(
+        self, name: str, node_type: str
+    ) -> LineageNode | None:
+        """Get a node by name and type combination.
+
+        Args:
+            name: Node name.
+            node_type: Node type (source, transform, sink).
+
+        Returns:
+            LineageNode or None if not found.
+        """
+        result = await self.session.execute(
+            select(LineageNode)
+            .where(LineageNode.name == name, LineageNode.node_type == node_type)
+            .limit(1)
+        )
+        return result.scalar_one_or_none()
+
 
 class LineageEdgeRepository(BaseRepository[LineageEdge]):
     """Repository for LineageEdge model operations."""
@@ -167,9 +194,9 @@ class LineageService:
 
     Provides functionality for:
     - Node and edge CRUD operations
-    - Impact analysis (upstream/downstream)
-    - Auto-discovery from source metadata
+    - Impact analysis (upstream/downstream) using truthound.lineage.ImpactAnalyzer
     - Position management for visualization
+    - Integration with truthound.lineage.LineageTracker for automatic tracking
     """
 
     def __init__(self, session: AsyncSession) -> None:
@@ -182,6 +209,44 @@ class LineageService:
         self.node_repo = LineageNodeRepository(session)
         self.edge_repo = LineageEdgeRepository(session)
 
+        # Initialize truthound lineage components if available
+        self._tracker = None
+        self._impact_analyzer = None
+        self._init_truthound_lineage()
+
+    def _init_truthound_lineage(self) -> None:
+        """Initialize truthound lineage components if available."""
+        try:
+            from truthound.lineage import LineageTracker, ImpactAnalyzer, LineageConfig
+
+            # Create lineage tracker with default configuration
+            config = LineageConfig(
+                auto_track=True,
+                track_column_level=True,
+            )
+            self._tracker = LineageTracker(config)
+            logger.info("truthound.lineage.LineageTracker initialized")
+
+        except ImportError:
+            logger.debug("truthound.lineage not available, using SQLAlchemy-based implementation")
+            self._tracker = None
+            self._impact_analyzer = None
+
+    def _sync_tracker_with_db(self) -> None:
+        """Sync the truthound tracker with database state."""
+        if self._tracker is None:
+            return
+
+        try:
+            from truthound.lineage import ImpactAnalyzer
+
+            # Sync nodes and edges from DB to tracker
+            # This ensures truthound analyzer has current state
+            self._impact_analyzer = ImpactAnalyzer(self._tracker.graph)
+
+        except Exception as e:
+            logger.warning(f"Failed to sync lineage tracker: {e}")
+
     # =========================================================================
     # Graph Operations
     # =========================================================================
@@ -248,6 +313,20 @@ class LineageService:
     # Node Operations
     # =========================================================================
 
+    async def get_node_by_name_and_type(
+        self, name: str, node_type: str
+    ) -> LineageNode | None:
+        """Get a node by name and type.
+
+        Args:
+            name: Node name.
+            node_type: Node type (source, transform, sink).
+
+        Returns:
+            LineageNode or None if not found.
+        """
+        return await self.node_repo.get_by_name_and_type(name, node_type)
+
     async def create_node(
         self,
         *,
@@ -270,7 +349,17 @@ class LineageService:
 
         Returns:
             Created node.
+
+        Raises:
+            ValueError: If a node with same name and type already exists.
         """
+        # Check for existing node with same name and type
+        existing = await self.get_node_by_name_and_type(name, node_type)
+        if existing:
+            raise ValueError(
+                f"Node with name '{name}' and type '{node_type}' already exists"
+            )
+
         node = await self.node_repo.create(
             name=name,
             node_type=node_type,
@@ -281,6 +370,43 @@ class LineageService:
         )
         return node
 
+    async def get_or_create_node(
+        self,
+        *,
+        name: str,
+        node_type: str,
+        source_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        position_x: float | None = None,
+        position_y: float | None = None,
+    ) -> tuple[LineageNode, bool]:
+        """Get an existing node or create a new one.
+
+        Args:
+            name: Node name.
+            node_type: Node type (source, transform, sink).
+            source_id: Optional linked data source ID.
+            metadata: Optional additional metadata.
+            position_x: X coordinate for visualization.
+            position_y: Y coordinate for visualization.
+
+        Returns:
+            Tuple of (node, created) where created is True if new node was created.
+        """
+        existing = await self.get_node_by_name_and_type(name, node_type)
+        if existing:
+            return existing, False
+
+        node = await self.node_repo.create(
+            name=name,
+            node_type=node_type,
+            source_id=source_id,
+            metadata_json=metadata,
+            position_x=position_x,
+            position_y=position_y,
+        )
+        return node, True
+
     async def get_node(self, node_id: str) -> LineageNode | None:
         """Get a node by ID.
 
@@ -352,7 +478,7 @@ class LineageService:
         target_node_id: str,
         edge_type: str = "derives_from",
         metadata: dict[str, Any] | None = None,
-    ) -> LineageEdge:
+    ) -> tuple[LineageEdge, LineageNode, LineageNode]:
         """Create a new lineage edge.
 
         Args:
@@ -362,7 +488,7 @@ class LineageService:
             metadata: Optional additional metadata.
 
         Returns:
-            Created edge.
+            Tuple of (created edge, source node, target node).
 
         Raises:
             ValueError: If source or target node not found, or edge already exists.
@@ -386,7 +512,7 @@ class LineageService:
             edge_type=edge_type,
             metadata_json=metadata,
         )
-        return edge
+        return edge, source_node, target_node
 
     async def get_edge(self, edge_id: str) -> LineageEdge | None:
         """Get an edge by ID.
@@ -411,7 +537,7 @@ class LineageService:
         return await self.edge_repo.delete(edge_id)
 
     # =========================================================================
-    # Impact Analysis
+    # Impact Analysis (using truthound.lineage.ImpactAnalyzer when available)
     # =========================================================================
 
     async def analyze_impact(
@@ -422,6 +548,9 @@ class LineageService:
     ) -> dict[str, Any]:
         """Analyze upstream/downstream impact from a node.
 
+        Uses truthound.lineage.ImpactAnalyzer when available for advanced analysis
+        including schema change impact and what-if scenarios.
+
         Args:
             node_id: Starting node ID.
             direction: Analysis direction.
@@ -437,20 +566,72 @@ class LineageService:
         if root_node is None:
             raise ValueError(f"Node '{node_id}' not found")
 
-        upstream_nodes: list[dict[str, Any]] = []
-        downstream_nodes: list[dict[str, Any]] = []
+        # Try to use truthound's ImpactAnalyzer for enhanced analysis
+        try:
+            from truthound.lineage import ImpactAnalyzer, LineageGraph
+
+            # Build lineage graph from database
+            graph = await self._build_truthound_graph()
+
+            # Create impact analyzer
+            analyzer = ImpactAnalyzer(graph)
+
+            # Perform impact analysis
+            impact_result = analyzer.analyze_impact(
+                node_id=node_id,
+                max_depth=max_depth,
+            )
+
+            # Convert truthound result to our format
+            # ImpactResult has: source_node, affected_nodes (downstream), total_affected, max_depth
+            downstream_nodes = [
+                {
+                    "id": getattr(an.node, 'id', None),
+                    "name": getattr(an.node, 'name', None),
+                    "node_type": getattr(an.node, 'node_type', None),
+                    "source_id": None,
+                    "depth": getattr(an, 'depth', 0),
+                    "impact_level": str(getattr(an, 'impact_level', '')),
+                }
+                for an in impact_result.affected_nodes
+            ]
+
+            # For upstream, use fallback traversal (ImpactAnalyzer only does downstream)
+            upstream_nodes_list: list[dict[str, Any]] = []
+            if direction in ("upstream", "both"):
+                upstream = await self._traverse_upstream(node_id, max_depth)
+                upstream_nodes_list = [self._node_summary(n) for n in upstream]
+
+            return {
+                "root_node_id": node_id,
+                "root_node_name": root_node.name,
+                "direction": direction,
+                "upstream_nodes": upstream_nodes_list,
+                "downstream_nodes": downstream_nodes,
+                "affected_sources": [],
+                "upstream_count": len(upstream_nodes_list),
+                "downstream_count": len(downstream_nodes),
+                "total_affected": len(upstream_nodes_list) + len(downstream_nodes),
+            }
+
+        except ImportError:
+            logger.debug("truthound.lineage not available, using fallback implementation")
+
+        # Fallback to SQLAlchemy-based implementation
+        upstream_nodes_list: list[dict[str, Any]] = []
+        downstream_nodes_list: list[dict[str, Any]] = []
         affected_sources: set[str] = set()
 
         if direction in ("upstream", "both"):
             upstream = await self._traverse_upstream(node_id, max_depth)
-            upstream_nodes = [self._node_summary(n) for n in upstream]
+            upstream_nodes_list = [self._node_summary(n) for n in upstream]
             for n in upstream:
                 if n.source_id:
                     affected_sources.add(n.source_id)
 
         if direction in ("downstream", "both"):
             downstream = await self._traverse_downstream(node_id, max_depth)
-            downstream_nodes = [self._node_summary(n) for n in downstream]
+            downstream_nodes_list = [self._node_summary(n) for n in downstream]
             for n in downstream:
                 if n.source_id:
                     affected_sources.add(n.source_id)
@@ -459,12 +640,223 @@ class LineageService:
             "root_node_id": node_id,
             "root_node_name": root_node.name,
             "direction": direction,
-            "upstream_nodes": upstream_nodes,
-            "downstream_nodes": downstream_nodes,
+            "upstream_nodes": upstream_nodes_list,
+            "downstream_nodes": downstream_nodes_list,
             "affected_sources": list(affected_sources),
-            "upstream_count": len(upstream_nodes),
-            "downstream_count": len(downstream_nodes),
-            "total_affected": len(upstream_nodes) + len(downstream_nodes),
+            "upstream_count": len(upstream_nodes_list),
+            "downstream_count": len(downstream_nodes_list),
+            "total_affected": len(upstream_nodes_list) + len(downstream_nodes_list),
+        }
+
+    async def _build_truthound_graph(self) -> Any:
+        """Build a truthound LineageGraph from database state."""
+        try:
+            from truthound.lineage import LineageGraph, LineageNode as TruthoundNode, LineageEdge as TruthoundEdge
+            from truthound.lineage.base import NodeType, EdgeType
+
+            # Map dashboard node types to truthound NodeType enum
+            node_type_map = {
+                "source": NodeType.SOURCE,
+                "transform": NodeType.TRANSFORMATION,
+                "sink": NodeType.EXTERNAL,
+            }
+
+            # Map dashboard edge types to truthound EdgeType enum
+            edge_type_map = {
+                "derives_from": EdgeType.DERIVED_FROM,
+                "transforms_to": EdgeType.TRANSFORMED_TO,
+                "joins_with": EdgeType.JOINED_WITH,
+                "filters_from": EdgeType.FILTERED_TO,
+            }
+
+            graph = LineageGraph()
+
+            # Add all nodes
+            nodes = await self.node_repo.get_all_nodes()
+            for node in nodes:
+                th_node_type = node_type_map.get(node.node_type, NodeType.EXTERNAL)
+                th_node = TruthoundNode(
+                    id=node.id,
+                    name=node.name,
+                    node_type=th_node_type,
+                )
+                graph.add_node(th_node)
+
+            # Add all edges
+            edges = await self.edge_repo.get_all_edges()
+            for edge in edges:
+                th_edge_type = edge_type_map.get(edge.edge_type, EdgeType.DERIVED_FROM)
+                th_edge = TruthoundEdge(
+                    source=edge.source_node_id,
+                    target=edge.target_node_id,
+                    edge_type=th_edge_type,
+                )
+                graph.add_edge(th_edge)
+
+            return graph
+
+        except ImportError:
+            return None
+
+    def _truthound_node_to_summary(self, node: Any) -> dict[str, Any]:
+        """Convert truthound LineageNode to summary dict."""
+        node_type = getattr(node, 'node_type', None)
+        return {
+            "id": getattr(node, 'id', None),
+            "name": getattr(node, 'name', None),
+            "node_type": str(node_type.value) if hasattr(node_type, 'value') else str(node_type),
+            "source_id": None,
+        }
+
+    async def analyze_schema_change_impact(
+        self,
+        node_id: str,
+        schema_changes: list[dict[str, Any]],
+    ) -> dict[str, Any]:
+        """Analyze impact of schema changes on downstream nodes.
+
+        Uses truthound.lineage.ImpactAnalyzer for what-if analysis.
+
+        Args:
+            node_id: Node ID where schema change occurs.
+            schema_changes: List of schema changes (e.g., column removal, type change).
+
+        Returns:
+            Schema change impact analysis.
+        """
+        try:
+            from truthound.lineage import ImpactAnalyzer
+
+            # Build lineage graph
+            graph = await self._build_truthound_graph()
+            if graph is None:
+                raise ImportError("truthound.lineage not available")
+
+            analyzer = ImpactAnalyzer(graph)
+
+            # Analyze schema change impact
+            impact_result = analyzer.analyze_schema_change_impact(
+                node_id=node_id,
+                changes=schema_changes,
+            )
+
+            return {
+                "node_id": node_id,
+                "schema_changes": schema_changes,
+                "breaking_changes": impact_result.breaking_changes,
+                "affected_downstream_nodes": [
+                    self._truthound_node_to_summary(n)
+                    for n in impact_result.affected_nodes
+                ],
+                "impact_severity": impact_result.severity,
+                "recommendations": impact_result.recommendations,
+            }
+
+        except ImportError:
+            # Fallback: basic downstream analysis without schema-aware logic
+            downstream = await self._traverse_downstream(node_id, max_depth=10)
+            return {
+                "node_id": node_id,
+                "schema_changes": schema_changes,
+                "breaking_changes": [],
+                "affected_downstream_nodes": [self._node_summary(n) for n in downstream],
+                "impact_severity": "unknown",
+                "recommendations": ["Truthound lineage module not available for detailed analysis"],
+            }
+
+    async def track_operation(
+        self,
+        operation_type: str,
+        source_nodes: list[str],
+        target_nodes: list[str],
+        metadata: dict[str, Any] | None = None,
+    ) -> dict[str, Any]:
+        """Track a data operation for automatic lineage recording.
+
+        Uses truthound.lineage.LineageTracker for automatic operation tracking.
+
+        Args:
+            operation_type: Type of operation (e.g., 'transform', 'aggregate', 'join').
+            source_nodes: Source node IDs.
+            target_nodes: Target node IDs.
+            metadata: Optional operation metadata.
+
+        Returns:
+            Tracking result with created edges.
+        """
+        try:
+            from truthound.lineage import LineageTracker, OperationType
+
+            if self._tracker is None:
+                self._init_truthound_lineage()
+
+            if self._tracker is not None:
+                # Map string to enum
+                op_type_map = {
+                    "transform": OperationType.TRANSFORM,
+                    "aggregate": OperationType.AGGREGATE,
+                    "join": OperationType.JOIN,
+                    "filter": OperationType.FILTER,
+                    "derive": OperationType.DERIVE,
+                }
+                op_type = op_type_map.get(operation_type, OperationType.TRANSFORM)
+
+                # Track operation
+                result = self._tracker.track_operation(
+                    operation_type=op_type,
+                    source_nodes=source_nodes,
+                    target_nodes=target_nodes,
+                    metadata=metadata or {},
+                )
+
+                # Also persist to database
+                created_edges = []
+                for source_id in source_nodes:
+                    for target_id in target_nodes:
+                        try:
+                            edge = await self.create_edge(
+                                source_node_id=source_id,
+                                target_node_id=target_id,
+                                edge_type=operation_type,
+                                metadata=metadata,
+                            )
+                            created_edges.append(self._edge_to_dict(edge))
+                        except ValueError:
+                            # Edge might already exist
+                            pass
+
+                return {
+                    "operation_type": operation_type,
+                    "source_nodes": source_nodes,
+                    "target_nodes": target_nodes,
+                    "created_edges": created_edges,
+                    "tracking_id": getattr(result, 'tracking_id', None),
+                }
+
+        except ImportError:
+            pass
+
+        # Fallback: just create edges in database
+        created_edges = []
+        for source_id in source_nodes:
+            for target_id in target_nodes:
+                try:
+                    edge = await self.create_edge(
+                        source_node_id=source_id,
+                        target_node_id=target_id,
+                        edge_type=operation_type,
+                        metadata=metadata,
+                    )
+                    created_edges.append(self._edge_to_dict(edge))
+                except ValueError:
+                    pass
+
+        return {
+            "operation_type": operation_type,
+            "source_nodes": source_nodes,
+            "target_nodes": target_nodes,
+            "created_edges": created_edges,
+            "tracking_id": None,
         }
 
     async def _traverse_upstream(
@@ -527,62 +919,6 @@ class LineageService:
 
         return result
 
-    # =========================================================================
-    # Auto-Discovery
-    # =========================================================================
-
-    async def auto_discover(
-        self,
-        source_id: str,
-        include_fk_relations: bool = True,
-        max_depth: int = 3,
-    ) -> dict[str, Any]:
-        """Auto-discover lineage from a data source.
-
-        This is a placeholder for more sophisticated discovery logic.
-        In a real implementation, this would analyze source metadata,
-        SQL queries, or foreign key relationships.
-
-        Args:
-            source_id: Source ID to discover from.
-            include_fk_relations: Include foreign key relationships (for DB sources).
-            max_depth: Maximum discovery depth.
-
-        Returns:
-            Discovered graph.
-        """
-        # Check if node already exists for this source
-        existing_node = await self.node_repo.get_by_source_id(source_id)
-        if existing_node:
-            return await self.get_graph(source_id)
-
-        # Get source info
-        from truthound_dashboard.db import Source as SourceModel
-
-        result = await self.session.execute(
-            select(SourceModel).where(SourceModel.id == source_id)
-        )
-        source = result.scalar_one_or_none()
-        if source is None:
-            raise ValueError(f"Source '{source_id}' not found")
-
-        # Create a node for this source
-        node = await self.create_node(
-            name=source.name,
-            node_type="source",
-            source_id=source_id,
-            metadata={"auto_discovered": True, "source_type": source.type},
-            position_x=100,
-            position_y=100,
-        )
-
-        return {
-            "source_id": source_id,
-            "discovered_nodes": 1,
-            "discovered_edges": 0,
-            "graph": await self.get_graph(source_id),
-        }
-
     # =========================================================================
     # Position Management
     # =========================================================================