salim0986/graph-bug-ai

1"""2Data Cleanup Service - Handles orphaned data and rollback operations3Ensures data consistency across Neo4j, Qdrant, and PostgreSQL4"""5 6import asyncio7from typing import Optional, Dict, Any8from .logger import setup_logger9from .graph_builder import GraphBuilder10from .vector_builder import VectorBuilder11 12logger = setup_logger(__name__)13 14 15class DataCleanup:16    """17    Manages cleanup operations for failed ingestions and deletions18    Ensures no orphaned data remains in any database19    """20    21    def __init__(self, graph_db: GraphBuilder, vector_db: VectorBuilder):22        self.graph_db = graph_db23        self.vector_db = vector_db24    25    async def cleanup_failed_ingestion(self, repo_id: str, reason: str = "unknown") -> Dict[str, Any]:26        """27        Rollback partial ingestion on failure28        29        This is called when ingestion fails mid-way to ensure no partial data remains.30        31        Args:32            repo_id: Repository identifier33            reason: Reason for failure (for logging)34            35        Returns:36            Dict with cleanup statistics37        """38        logger.warning(f"🧹 Starting cleanup for failed ingestion: {repo_id}")39        logger.warning(f"   Reason: {reason}")40        41        stats = {42            "repo_id": repo_id,43            "reason": reason,44            "vectors_deleted": 0,45            "graph_nodes_deleted": 0,46            "success": False47        }48        49        try:50            # 1. Delete vectors from Qdrant51            try:52                self.vector_db.delete_repo(repo_id)53                stats["vectors_deleted"] = "unknown"  # Qdrant doesn't return count54                logger.info(f"   ✓ Vectors cleaned up for {repo_id}")55            except Exception as e:56                logger.error(f"   ✗ Vector cleanup failed: {e}")57                raise58            59            # 2. Delete graph nodes from Neo4j60            try:61                deleted_count = self.graph_db.delete_repo(repo_id)62                stats["graph_nodes_deleted"] = deleted_count63                logger.info(f"   ✓ Graph nodes cleaned up: {deleted_count} nodes")64            except Exception as e:65                logger.error(f"   ✗ Graph cleanup failed: {e}")66                raise67            68            stats["success"] = True69            logger.info(f"✅ Cleanup complete for {repo_id}")70            71        except Exception as e:72            logger.error(f"❌ Cleanup failed for {repo_id}: {e}")73            stats["error"] = str(e)74        75        return stats76    77    async def verify_consistency(self, repo_id: str) -> Dict[str, Any]:78        """79        Verify data consistency across databases80        81        Checks that repository data exists in both Neo4j and Qdrant,82        or doesn't exist in either (no orphaned data).83        84        Args:85            repo_id: Repository identifier86            87        Returns:88            Dict with consistency status89        """90        logger.info(f"🔍 Verifying data consistency for {repo_id}")91        92        result = {93            "repo_id": repo_id,94            "consistent": False,95            "has_vectors": False,96            "has_graph_data": False,97            "issues": []98        }99        100        try:101            # Check vectors102            vector_stats = self.vector_db.get_repository_stats(repo_id)103            result["has_vectors"] = vector_stats.get("total_snippets", 0) > 0104            105            # Check graph106            graph_stats = self.graph_db.get_repo_node_count(repo_id)107            result["has_graph_data"] = graph_stats > 0108            109            # Consistency check110            if result["has_vectors"] == result["has_graph_data"]:111                result["consistent"] = True112                logger.info(f"   ✓ Data is consistent")113            else:114                result["consistent"] = False115                if result["has_vectors"] and not result["has_graph_data"]:116                    result["issues"].append("Orphaned vectors (no graph data)")117                elif result["has_graph_data"] and not result["has_vectors"]:118                    result["issues"].append("Orphaned graph data (no vectors)")119                logger.warning(f"   ⚠️ Inconsistency detected: {result['issues']}")120        121        except Exception as e:122            logger.error(f"❌ Consistency check failed: {e}")123            result["error"] = str(e)124        125        return result126    127    async def cleanup_orphaned_data(self) -> Dict[str, Any]:128        """129        Scan and cleanup any orphaned data across all repositories130        131        This is a maintenance operation that should be run periodically.132        133        Returns:134            Dict with cleanup statistics135        """136        logger.info("🔍 Scanning for orphaned data across all repositories")137        138        stats = {139            "repos_checked": 0,140            "inconsistencies_found": 0,141            "cleaned_up": 0,142            "errors": []143        }144        145        # This would need to iterate through all repos in PostgreSQL146        # and check consistency. Left as TODO for now.147        148        logger.info("⚠️ Global orphaned data cleanup not yet implemented")149        return stats150 151 152class IngestionCheckpoint:153    """154    Transaction-like behavior for ingestion operations155    Allows rollback to previous state on failure156    """157    158    def __init__(self, repo_id: str, cleanup: DataCleanup):159        self.repo_id = repo_id160        self.cleanup = cleanup161        self.created_at = None162        self.committed = False163    164    async def __aenter__(self):165        """Start checkpoint"""166        logger.info(f"📍 Creating ingestion checkpoint for {self.repo_id}")167        from datetime import datetime168        self.created_at = datetime.utcnow()169        return self170    171    async def __aexit__(self, exc_type, exc_val, exc_tb):172        """Rollback on exception, commit on success"""173        if exc_type is not None and not self.committed:174            # Exception occurred and not committed - rollback175            logger.error(f"❌ Ingestion failed, rolling back: {exc_val}")176            await self.cleanup.cleanup_failed_ingestion(177                self.repo_id,178                reason=f"{exc_type.__name__}: {exc_val}"179            )180            return False  # Re-raise exception181        elif self.committed:182            logger.info(f"✅ Checkpoint committed for {self.repo_id}")183        return False184    185    def commit(self):186        """Mark checkpoint as successful"""187        self.committed = True188        logger.info(f"✓ Committing checkpoint for {self.repo_id}")189