codegraphcontext-rust 0.5.5__cp311-cp311-win_amd64.whl

codegraphcontext/core/cgc_bundle.py

@@ -0,0 +1,858 @@
+# src/codegraphcontext/core/cgc_bundle.py
+"""
+This module handles the creation and loading of .cgc (CodeGraphContext Bundle) files.
+
+A .cgc file is a portable, pre-indexed graph snapshot that can be distributed and loaded
+instantly without re-indexing. This enables:
+- Pre-indexing famous repositories once
+- Distributing graph knowledge as artifacts
+- Instant context loading for LLMs
+- Version-controlled code knowledge
+
+Bundle Structure:
+    .cgc (ZIP archive)
+    ├── metadata.json       # Repository and indexing metadata
+    ├── schema.json         # Graph schema definition
+    ├── nodes.jsonl         # All nodes (one JSON object per line)
+    ├── edges.jsonl         # All relationships (one JSON object per line)
+    ├── stats.json          # Graph statistics
+    └── README.md           # Human-readable description
+"""
+
+import json
+import zipfile
+import tempfile
+from pathlib import Path
+from typing import Dict, List, Optional, Any, Tuple
+from datetime import datetime, date
+import subprocess
+
+from codegraphcontext.utils.debug_log import debug_log, info_logger, error_logger, warning_logger
+
+
+class _BundleEncoder(json.JSONEncoder):
+    """Handles Neo4j DateTime and other non-standard types for bundle serialization."""
+    def default(self, obj):
+        if isinstance(obj, (datetime, date)):
+            return obj.isoformat()
+        if hasattr(obj, 'isoformat'):
+            return obj.isoformat()
+        if hasattr(obj, 'iso_format'):
+            return obj.iso_format()
+        if isinstance(obj, Path):
+            return str(obj)
+        if isinstance(obj, set):
+            return list(obj)
+        if isinstance(obj, bytes):
+            return obj.decode('utf-8', errors='replace')
+        return super().default(obj)
+
+
+class CGCBundle:
+    """Handles creation and loading of .cgc bundle files."""
+    
+    VERSION = "0.1.0"  # CGC bundle format version
+    
+    def __init__(self, db_manager):
+        """
+        Initialize the CGC Bundle handler.
+        
+        Args:
+            db_manager: DatabaseManager instance for graph queries
+        """
+        self.db_manager = db_manager
+    
+    def _get_id_function(self) -> str:
+        """
+        Get the appropriate ID function based on the database backend.
+        
+        Returns:
+            str: 'elementId' for Neo4j, 'id' for FalkorDB
+        """
+        # Check if we're using Neo4j or FalkorDB
+        backend = self.db_manager.get_backend_type()
+        if backend == 'neo4j':
+            return 'elementId'
+        else:  # FalkorDB or other backends
+            return 'id'
+
+    
+    def export_to_bundle(
+        self,
+        output_path: Path,
+        repo_path: Optional[Path] = None,
+        include_stats: bool = True
+    ) -> Tuple[bool, str]:
+        """
+        Export the current graph (or a specific repository) to a .cgc bundle.
+        
+        Args:
+            output_path: Path where the .cgc file should be saved
+            repo_path: Optional specific repository path to export (None = export all)
+            include_stats: Whether to include detailed statistics
+            
+        Returns:
+            Tuple[bool, str]: (success, message)
+        """
+        try:
+            info_logger(f"Starting export to {output_path}")
+            
+            # Ensure output path has .cgc extension
+            if not str(output_path).endswith('.cgc'):
+                output_path = Path(str(output_path) + '.cgc')
+            
+            # Create temporary directory for bundle contents
+            with tempfile.TemporaryDirectory() as temp_dir:
+                temp_path = Path(temp_dir)
+                
+                # Step 1: Extract metadata
+                info_logger("Extracting metadata...")
+                metadata = self._extract_metadata(repo_path)
+                with open(temp_path / "metadata.json", 'w') as f:
+                    json.dump(metadata, f, indent=2, cls=_BundleEncoder)
+                
+                # Step 2: Extract schema
+                info_logger("Extracting schema...")
+                schema = self._extract_schema()
+                with open(temp_path / "schema.json", 'w') as f:
+                    json.dump(schema, f, indent=2, cls=_BundleEncoder)
+                
+                # Step 3: Extract nodes
+                info_logger("Extracting nodes...")
+                node_count = self._extract_nodes(temp_path / "nodes.jsonl", repo_path)
+                
+                # Step 4: Extract edges
+                info_logger("Extracting edges...")
+                edge_count = self._extract_edges(temp_path / "edges.jsonl", repo_path)
+                
+                # Step 5: Generate statistics
+                if include_stats:
+                    info_logger("Generating statistics...")
+                    stats = self._generate_stats(repo_path, node_count, edge_count)
+                    with open(temp_path / "stats.json", 'w') as f:
+                        json.dump(stats, f, indent=2, cls=_BundleEncoder)
+                
+                # Step 6: Create README
+                self._create_readme(temp_path / "README.md", metadata, stats if include_stats else None)
+                
+                # Step 7: Create ZIP archive
+                info_logger("Creating bundle archive...")
+                self._create_zip(temp_path, output_path)
+            
+            success_msg = f"✅ Successfully exported to {output_path}\n"
+            success_msg += f"   Nodes: {node_count:,} | Edges: {edge_count:,}"
+            info_logger(success_msg)
+            return True, success_msg
+            
+        except Exception as e:
+            import traceback
+            error_msg = f"Failed to export bundle: {str(e)}"
+            error_logger(error_msg)
+            # Print full traceback for debugging
+            traceback.print_exc()
+            return False, error_msg
+    
+    def import_from_bundle(
+        self,
+        bundle_path: Path,
+        clear_existing: bool = False,
+        readonly: bool = False
+    ) -> Tuple[bool, str]:
+        """
+        Import a .cgc bundle into the current database.
+        
+        Args:
+            bundle_path: Path to the .cgc file
+            clear_existing: Whether to clear existing graph data first
+            readonly: If True, mount as read-only (future feature)
+            
+        Returns:
+            Tuple[bool, str]: (success, message)
+        """
+        try:
+            info_logger(f"Starting import from {bundle_path}")
+            
+            if not bundle_path.exists():
+                return False, f"Bundle file not found: {bundle_path}"
+            
+            # Extract bundle to temporary directory
+            with tempfile.TemporaryDirectory() as temp_dir:
+                temp_path = Path(temp_dir)
+                
+                # Step 1: Extract ZIP (with Zip Slip protection)
+                info_logger("Extracting bundle...")
+                with zipfile.ZipFile(bundle_path, 'r') as zip_ref:
+                    for entry in zip_ref.namelist():
+                        resolved = (temp_path / entry).resolve()
+                        if not str(resolved).startswith(str(temp_path.resolve())):
+                            return False, f"Zip Slip detected: entry '{entry}' escapes target directory"
+                    zip_ref.extractall(temp_path)
+                
+                # Step 2: Validate bundle
+                info_logger("Validating bundle...")
+                is_valid, validation_msg = self._validate_bundle(temp_path)
+                if not is_valid:
+                    return False, f"Invalid bundle: {validation_msg}"
+                
+                # Step 3: Load metadata
+                with open(temp_path / "metadata.json", 'r') as f:
+                    metadata = json.load(f)
+                
+                info_logger(f"Loading bundle: {metadata.get('repo', 'unknown')}")
+                info_logger(f"Bundle version: {metadata.get('cgc_version', 'unknown')}")
+                
+                # Step 4: Handle existing data
+                repo_name = metadata.get('repo', 'unknown')
+                repo_path = metadata.get('repo_path')
+                
+                if clear_existing:
+                    # User explicitly wants to clear - remove everything
+                    info_logger("Clearing all existing graph data...")
+                    self._clear_graph()
+                else:
+                    # Check if this repository already exists (only when NOT clearing)
+                    existing_repo = self._check_existing_repository(repo_name, repo_path)
+                    
+                    if existing_repo:
+                        return False, f"Repository '{repo_name}' already exists in the database. Use clear_existing=True to replace it."
+                
+                
+                # Step 5: Create schema
+                info_logger("Creating schema...")
+                self._import_schema(temp_path / "schema.json")
+                
+                # Step 6: Import nodes
+                info_logger("Importing nodes...")
+                node_count = self._import_nodes(temp_path / "nodes.jsonl")
+                
+                # Step 7: Import edges
+                info_logger("Importing edges...")
+                edge_count = self._import_edges(temp_path / "edges.jsonl")
+            
+            success_msg = f"✅ Successfully imported {bundle_path.name}\n"
+            success_msg += f"   Repository: {metadata.get('repo', 'unknown')}\n"
+            success_msg += f"   Nodes: {node_count:,} | Edges: {edge_count:,}"
+            info_logger(success_msg)
+            return True, success_msg
+            
+        except Exception as e:
+            error_msg = f"Failed to import bundle: {str(e)}"
+            error_logger(error_msg)
+            return False, error_msg
+    
+    # ========================================================================
+    # EXPORT HELPERS
+    # ========================================================================
+    
+    def _extract_metadata(self, repo_path: Optional[Path]) -> Dict[str, Any]:
+        """Extract metadata about the repository and indexing process."""
+        metadata = {
+            "cgc_version": self.VERSION,
+            "exported_at": datetime.now().isoformat(),
+            "format_version": "1.0"
+        }
+        
+        # Get repository information
+        with self.db_manager.get_driver().session() as session:
+            if repo_path:
+                # Specific repository
+                result = session.run(
+                    "MATCH (r:Repository {path: $path}) RETURN r",
+                    path=str(repo_path.resolve())
+                )
+                repo_node = result.single()
+                if repo_node:
+                    node = repo_node['r']
+                    # Convert Node to dict (handle both Neo4j and FalkorDB)
+                    try:
+                        repo = dict(node)
+                    except TypeError:
+                        # FalkorDB nodes - access properties directly
+                        repo = {}
+                        if hasattr(node, '_properties'):
+                            repo = dict(node._properties)
+                        elif hasattr(node, 'properties'):
+                            repo = dict(node.properties)
+                        else:
+                            # Fallback: try to get individual properties
+                            for attr in ['name', 'path', 'is_dependency']:
+                                if hasattr(node, attr):
+                                    repo[attr] = getattr(node, attr)
+                    
+                    metadata["repo"] = repo.get('name', str(repo_path))
+                    metadata["repo_path"] = repo.get('path')
+                    metadata["is_dependency"] = repo.get('is_dependency', False)
+            else:
+                # All repositories
+                result = session.run(
+                    "MATCH (r:Repository) RETURN r.name as name, r.path as path"
+                )
+                repos = [{"name": record["name"], "path": record["path"]} for record in result]
+                metadata["repositories"] = repos
+                metadata["repo"] = "multiple" if len(repos) > 1 else repos[0]["name"] if repos else "unknown"
+            
+            # Try to get git information if available
+            if repo_path and repo_path.exists():
+                try:
+                    commit = subprocess.check_output(
+                        ['git', 'rev-parse', 'HEAD'],
+                        cwd=repo_path,
+                        stderr=subprocess.DEVNULL
+                    ).decode().strip()
+                    metadata["commit"] = commit[:8]
+                except (subprocess.CalledProcessError, FileNotFoundError):
+                    pass
+
+                try:
+                    result = session.run("""
+                        MATCH (f:File)
+                        WHERE f.path STARTS WITH $repo_path
+                        RETURN f.language as language, count(*) as count
+                        ORDER BY count DESC
+                    """, repo_path=str(repo_path.resolve()))
+                    languages = {record["language"]: record["count"] for record in result if record["language"]}
+                    metadata["languages"] = list(languages.keys())
+                except Exception:
+                    pass
+        
+        return metadata
+    
+    def _extract_schema(self) -> Dict[str, Any]:
+        """Extract the graph schema (node labels, relationship types, constraints)."""
+        schema = {
+            "node_labels": [],
+            "relationship_types": [],
+            "constraints": [],
+            "indexes": []
+        }
+        
+        with self.db_manager.get_driver().session() as session:
+            # Get node labels
+            try:
+                result = session.run("CALL db.labels()")
+                labels = []
+                for record in result:
+                    try:
+                        labels.append(record[0])
+                    except (KeyError, TypeError):
+                        if hasattr(record, 'values'):
+                            vals = list(record.values())
+                            if vals:
+                                labels.append(vals[0])
+                schema["node_labels"] = labels
+            except Exception:
+                schema["node_labels"] = []
+            
+            # Get relationship types
+            try:
+                result = session.run("CALL db.relationshipTypes()")
+                rel_types = []
+                for record in result:
+                    try:
+                        rel_types.append(record[0])
+                    except (KeyError, TypeError):
+                        if hasattr(record, 'values'):
+                            vals = list(record.values())
+                            if vals:
+                                rel_types.append(vals[0])
+                schema["relationship_types"] = rel_types
+            except Exception:
+                schema["relationship_types"] = []
+            
+            # Get constraints (Neo4j specific, may not work on all backends)
+            try:
+                result = session.run("SHOW CONSTRAINTS")
+                schema["constraints"] = [dict(record) for record in result]
+            except:
+                pass
+            
+            # Get indexes
+            try:
+                result = session.run("SHOW INDEXES")
+                schema["indexes"] = [dict(record) for record in result]
+            except:
+                pass
+        
+        return schema
+    
+    def _extract_nodes(self, output_file: Path, repo_path: Optional[Path]) -> int:
+        """Extract all nodes to JSONL format."""
+        count = 0
+        
+        with self.db_manager.get_driver().session() as session:
+            # Build query based on repo_path
+            if repo_path:
+                query = """
+                    MATCH (n)
+                    WHERE n.path STARTS WITH $repo_path
+                    RETURN n, labels(n) as labels
+                """
+                params = {"repo_path": str(repo_path.resolve())}
+            else:
+                query = "MATCH (n) RETURN n, labels(n) as labels"
+                params = {}
+            
+            # Run query with proper parameter handling for both Neo4j and FalkorDB
+            try:
+                result = session.run(query, **params)
+            except TypeError:
+                # FalkorDB might not support **params, try without
+                result = session.run(query)
+            
+            with open(output_file, 'w') as f:
+                for record in result:
+                    node = record['n']
+                    labels = record['labels']
+                    
+                    # Convert node to dict (handle both Neo4j and FalkorDB)
+                    try:
+                        node_dict = dict(node)
+                    except TypeError:
+                        # FalkorDB nodes might not be directly convertible
+                        node_dict = {}
+                        if hasattr(node, '_properties'):
+                            node_dict = dict(node._properties)
+                        elif hasattr(node, 'properties'):
+                            node_dict = dict(node.properties)
+                    
+                    node_dict['_labels'] = labels
+                    
+                    # Store internal ID for reference
+                    if hasattr(node, 'element_id'):
+                        node_dict['_id'] = node.element_id
+                    elif hasattr(node, 'id'):
+                        node_dict['_id'] = str(node.id)
+                    
+                    f.write(json.dumps(node_dict, cls=_BundleEncoder) + '\n')
+                    count += 1
+        
+        return count
+    
+    def _extract_edges(self, output_file: Path, repo_path: Optional[Path]) -> int:
+        """Extract all relationships to JSONL format."""
+        count = 0
+        
+        with self.db_manager.get_driver().session() as session:
+            # Build query based on repo_path
+            if repo_path:
+                query = """
+                    MATCH (n)-[r]->(m)
+                    WHERE (n.path STARTS WITH $repo_path)
+                       OR (m.path STARTS WITH $repo_path)
+                    RETURN n, r, m, type(r) as rel_type
+                """
+                params = {"repo_path": str(repo_path.resolve())}
+            else:
+                query = "MATCH (n)-[r]->(m) RETURN n, r, m, type(r) as rel_type"
+                params = {}
+            
+            # Run query with proper parameter handling for both Neo4j and FalkorDB
+            try:
+                result = session.run(query, **params)
+            except TypeError:
+                # FalkorDB might not support **params, try without
+                result = session.run(query)
+            
+            with open(output_file, 'w') as f:
+                for record in result:
+                    source = record['n']
+                    target = record['m']
+                    rel = record['r']
+                    rel_type = record['rel_type']
+                    
+                    # Get source and target IDs (handle both Neo4j and FalkorDB)
+                    if hasattr(source, 'element_id'):
+                        from_id = source.element_id
+                    elif hasattr(source, 'id'):
+                        from_id = str(source.id)
+                    else:
+                        from_id = str(id(source))  # Fallback
+                    
+                    if hasattr(target, 'element_id'):
+                        to_id = target.element_id
+                    elif hasattr(target, 'id'):
+                        to_id = str(target.id)
+                    else:
+                        to_id = str(id(target))  # Fallback
+                    
+                    # Get relationship properties
+                    try:
+                        rel_props = dict(rel)
+                    except TypeError:
+                        rel_props = {}
+                        if hasattr(rel, '_properties'):
+                            rel_props = dict(rel._properties)
+                        elif hasattr(rel, 'properties'):
+                            rel_props = dict(rel.properties)
+                    
+                    # Create edge representation
+                    edge_dict = {
+                        'from': from_id,
+                        'to': to_id,
+                        'type': rel_type,
+                        'properties': rel_props
+                    }
+                    
+                    f.write(json.dumps(edge_dict, cls=_BundleEncoder) + '\n')
+                    count += 1
+        
+        return count
+    
+    def _generate_stats(self, repo_path: Optional[Path], node_count: int, edge_count: int) -> Dict[str, Any]:
+        """Generate statistics about the graph."""
+        stats = {
+            "total_nodes": node_count,
+            "total_edges": edge_count,
+            "generated_at": datetime.now().isoformat()
+        }
+        
+        with self.db_manager.get_driver().session() as session:
+            # Count by node type
+            result = session.run("""
+                MATCH (n)
+                RETURN labels(n)[0] as label, count(*) as count
+                ORDER BY count DESC
+            """)
+            stats["nodes_by_type"] = {record["label"]: record["count"] for record in result if record["label"]}
+            
+            # Count by relationship type
+            result = session.run("""
+                MATCH ()-[r]->()
+                RETURN type(r) as type, count(*) as count
+                ORDER BY count DESC
+            """)
+            stats["edges_by_type"] = {record["type"]: record["count"] for record in result}
+            
+            # File count
+            if repo_path:
+                result = session.run(
+                    "MATCH (f:File) WHERE f.path STARTS WITH $repo_path RETURN count(f) as count",
+                    repo_path=str(repo_path.resolve())
+                )
+            else:
+                result = session.run("MATCH (f:File) RETURN count(f) as count")
+            
+            file_count = result.single()
+            stats["files"] = file_count["count"] if file_count else 0
+        
+        return stats
+    
+    def _create_readme(self, output_file: Path, metadata: Dict, stats: Optional[Dict]):
+        """Create a human-readable README for the bundle."""
+        readme_content = f"""# CodeGraphContext Bundle
+
+## Repository Information
+- **Repository**: {metadata.get('repo', 'Unknown')}
+- **Exported**: {metadata.get('exported_at', 'Unknown')}
+- **CGC Version**: {metadata.get('cgc_version', 'Unknown')}
+"""
+        
+        if 'commit' in metadata:
+            readme_content += f"- **Commit**: {metadata['commit']}\n"
+        
+        if 'languages' in metadata:
+            readme_content += f"- **Languages**: {', '.join(metadata['languages'])}\n"
+        
+        if stats:
+            readme_content += f"""
+## Statistics
+- **Total Nodes**: {stats.get('total_nodes', 0):,}
+- **Total Edges**: {stats.get('total_edges', 0):,}
+- **Files**: {stats.get('files', 0):,}
+
+### Nodes by Type
+"""
+            for label, count in stats.get('nodes_by_type', {}).items():
+                readme_content += f"- {label}: {count:,}\n"
+            
+            readme_content += "\n### Edges by Type\n"
+            for rel_type, count in stats.get('edges_by_type', {}).items():
+                readme_content += f"- {rel_type}: {count:,}\n"
+        
+        readme_content += """
+## Usage
+
+Load this bundle with:
+```bash
+cgc load <bundle-file>.cgc
+```
+
+Or import into existing graph:
+```bash
+cgc import <bundle-file>.cgc
+```
+"""
+        
+        with open(output_file, 'w') as f:
+            f.write(readme_content)
+    
+    def _create_zip(self, source_dir: Path, output_file: Path):
+        """Create a ZIP archive from the bundle directory."""
+        with zipfile.ZipFile(output_file, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            for path in source_dir.rglob('*'):
+                if path.is_file():
+                    arcname = path.relative_to(source_dir)
+                    zipf.write(path, arcname)
+    
+    # ========================================================================
+    # IMPORT HELPERS
+    # ========================================================================
+    
+    def _validate_bundle(self, bundle_dir: Path) -> Tuple[bool, str]:
+        """Validate that the bundle contains all required files."""
+        required_files = ['metadata.json', 'schema.json', 'nodes.jsonl', 'edges.jsonl']
+        
+        for file_name in required_files:
+            if not (bundle_dir / file_name).exists():
+                return False, f"Missing required file: {file_name}"
+        
+        # Validate metadata
+        try:
+            with open(bundle_dir / "metadata.json", 'r') as f:
+                metadata = json.load(f)
+                if 'cgc_version' not in metadata:
+                    return False, "Invalid metadata: missing cgc_version"
+        except json.JSONDecodeError as e:
+            return False, f"Invalid metadata.json: {e}"
+        
+        return True, "Valid bundle"
+    
+    def _check_existing_repository(self, repo_name: str, repo_path: Optional[str]) -> bool:
+        """Check if a repository already exists in the database."""
+        with self.db_manager.get_driver().session() as session:
+            # Try to find by name first
+            result = session.run(
+                "MATCH (r:Repository {name: $name}) RETURN r LIMIT 1",
+                name=repo_name
+            )
+            if result.single():
+                return True
+            
+            # If repo_path is provided, also check by path
+            if repo_path:
+                result = session.run(
+                    "MATCH (r:Repository {path: $path}) RETURN r LIMIT 1",
+                    path=repo_path
+                )
+                if result.single():
+                    return True
+        
+        return False
+    
+    def _delete_repository(self, repo_identifier: str):
+        """Delete a specific repository and all its related nodes from the graph."""
+        with self.db_manager.get_driver().session() as session:
+            # First, try to find the repository by name or path
+            result = session.run("""
+                MATCH (r:Repository)
+                WHERE r.name = $identifier OR r.path = $identifier
+                RETURN r.path as path
+                LIMIT 1
+            """, identifier=repo_identifier)
+            
+            record = result.single()
+            if not record:
+                warning_logger(f"Repository '{repo_identifier}' not found for deletion")
+                return
+            
+            repo_path = record['path']
+            
+            # Delete all nodes that belong to this repository
+            # Files, Functions, Classes, Modules all have paths that start with repo_path
+            session.run("""
+                MATCH (n)
+                WHERE n.path STARTS WITH $repo_path
+                DETACH DELETE n
+            """, repo_path=repo_path)
+            
+            # Delete the repository node itself
+            session.run("""
+                MATCH (r:Repository)
+                WHERE r.path = $repo_path
+                DELETE r
+            """, repo_path=repo_path)
+            
+            info_logger(f"Deleted repository: {repo_identifier}")
+    
+    def _clear_graph(self):
+        """Clear all nodes and relationships from the graph in batches."""
+        with self.db_manager.get_driver().session() as session:
+            while True:
+                result = session.run(
+                    "MATCH (n) WITH n LIMIT 500 DETACH DELETE n RETURN count(n) as deleted"
+                )
+                record = result.single()
+                if not record or record["deleted"] == 0:
+                    break
+    
+    def _import_schema(self, schema_file: Path):
+        """Import schema (constraints and indexes)."""
+        with open(schema_file, 'r') as f:
+            schema = json.load(f)
+        
+        # Note: Schema import is complex and database-specific
+        # For now, we'll rely on the application to create the schema
+        # This is a placeholder for future enhancement
+        debug_log("Schema import not yet implemented - relying on application schema")
+    
+    def _import_nodes(self, nodes_file: Path) -> int:
+        """Import nodes from JSONL file."""
+        count = 0
+        batch_size = 1000
+        batch = []
+        
+        # Create a mapping from old IDs to new IDs
+        id_mapping = {}
+        
+        with self.db_manager.get_driver().session() as session:
+            with open(nodes_file, 'r') as f:
+                for line in f:
+                    node_data = json.loads(line)
+                    
+                    # Extract labels and old ID (handle both Neo4j and KuzuDB formats)
+                    labels = node_data.pop('_labels', None) or node_data.pop('_label', None) or []
+                    if isinstance(labels, str):
+                        labels = [labels]
+                    old_id = node_data.pop('_id', None)
+                    # Convert dict IDs to hashable tuples for mapping
+                    if isinstance(old_id, dict):
+                        old_id = (old_id.get('table', 0), old_id.get('offset', 0))
+                    
+                    # Remove internal properties
+                    node_data.pop('_element_id', None)
+                    
+                    batch.append((labels, node_data, old_id))
+                    
+                    if len(batch) >= batch_size:
+                        count += self._import_node_batch(session, batch, id_mapping)
+                        batch = []
+                
+                # Import remaining nodes
+                if batch:
+                    count += self._import_node_batch(session, batch, id_mapping)
+        
+        # Store ID mapping for edge import
+        self._id_mapping = id_mapping
+        
+        return count
+    
+    _PK_MAP = {
+        'Repository': 'path', 'File': 'path', 'Directory': 'path',
+        'Module': 'name',
+        'Function': 'uid', 'Class': 'uid', 'Variable': 'uid',
+        'Trait': 'uid', 'Interface': 'uid', 'Macro': 'uid',
+        'Struct': 'uid', 'Enum': 'uid', 'Union': 'uid',
+        'Annotation': 'uid', 'Record': 'uid', 'Property': 'uid',
+        'Parameter': 'uid',
+    }
+    _UID_PARTS = {
+        'Function': ['name', 'path', 'line_number'],
+        'Class': ['name', 'path', 'line_number'],
+        'Variable': ['name', 'path', 'line_number'],
+        'Trait': ['name', 'path', 'line_number'],
+        'Interface': ['name', 'path', 'line_number'],
+        'Macro': ['name', 'path', 'line_number'],
+        'Struct': ['name', 'path', 'line_number'],
+        'Enum': ['name', 'path', 'line_number'],
+        'Union': ['name', 'path', 'line_number'],
+        'Annotation': ['name', 'path', 'line_number'],
+        'Record': ['name', 'path', 'line_number'],
+        'Property': ['name', 'path', 'line_number'],
+        'Parameter': ['name', 'path', 'function_line_number'],
+    }
+
+    def _import_node_batch(self, session, batch: List[Tuple], id_mapping: Dict) -> int:
+        """Import a batch of nodes."""
+        id_function = self._get_id_function()
+        
+        for labels, properties, old_id in batch:
+            if not labels:
+                continue
+            
+            if isinstance(labels, str):
+                labels = [labels]
+            label_str = ':'.join(labels)
+            primary_label = labels[0]
+
+            pk_field = self._PK_MAP.get(primary_label)
+            if pk_field == 'uid' and 'uid' not in properties:
+                parts = self._UID_PARTS.get(primary_label, [])
+                properties['uid'] = ''.join(str(properties.get(p, '')) for p in parts)
+
+            if pk_field and pk_field in properties:
+                pk_val = properties[pk_field]
+                remaining = {k: v for k, v in properties.items() if k != pk_field}
+                query = (
+                    f"MERGE (n:{label_str} {{{pk_field}: $pk_val}}) "
+                    f"SET n += $props RETURN {id_function}(n) as new_id"
+                )
+                result = session.run(query, pk_val=pk_val, props=remaining)
+            else:
+                query = f"CREATE (n:{label_str}) SET n = $props RETURN {id_function}(n) as new_id"
+                result = session.run(query, props=properties)
+
+            record = result.single()
+            if record and old_id:
+                id_mapping[old_id] = record['new_id']
+        
+        return len(batch)
+    
+    def _import_edges(self, edges_file: Path) -> int:
+        """Import edges from JSONL file."""
+        count = 0
+        batch_size = 1000
+        batch = []
+        
+        with self.db_manager.get_driver().session() as session:
+            with open(edges_file, 'r') as f:
+                for line in f:
+                    edge_data = json.loads(line)
+                    batch.append(edge_data)
+                    
+                    if len(batch) >= batch_size:
+                        count += self._import_edge_batch(session, batch)
+                        batch = []
+                
+                # Import remaining edges
+                if batch:
+                    count += self._import_edge_batch(session, batch)
+        
+        return count
+    
+    def _import_edge_batch(self, session, batch: List[Dict]) -> int:
+        """Import a batch of edges."""
+        id_mapping = getattr(self, '_id_mapping', {})
+        # Detect database backend to use appropriate ID function
+        id_function = self._get_id_function()
+        
+        for edge in batch:
+            old_from = edge.get('from')
+            old_to = edge.get('to')
+            # Convert dict IDs to hashable tuples (matches node import conversion)
+            if isinstance(old_from, dict):
+                old_from = (old_from.get('table', 0), old_from.get('offset', 0))
+            if isinstance(old_to, dict):
+                old_to = (old_to.get('table', 0), old_to.get('offset', 0))
+            rel_type = edge.get('type')
+            properties = edge.get('properties', {})
+            
+            # Map old IDs to new IDs
+            new_from = id_mapping.get(old_from)
+            new_to = id_mapping.get(old_to)
+            
+            if not new_from or not new_to:
+                warning_logger(f"Skipping edge: node IDs not found in mapping")
+                continue
+            
+            # Create relationship
+            query = f"""
+                MATCH (a), (b)
+                WHERE {id_function}(a) = $from_id AND {id_function}(b) = $to_id
+                CREATE (a)-[r:{rel_type}]->(b)
+                SET r = $props
+            """
+            
+            session.run(query, from_id=new_from, to_id=new_to, props=properties)
+        
+        return len(batch)
+