shotgun-sh 0.3.3.dev1__py3-none-any.whl → 0.6.2__py3-none-any.whl

shotgun/codebase/core/manager.py

@@ -5,22 +5,34 @@ from __future__ import annotations
 import asyncio
 import hashlib
 import json
+import shutil
 import time
 import uuid
 from collections.abc import Awaitable, Callable
 from pathlib import Path
-from typing import Any, ClassVar
+from typing import TYPE_CHECKING, Any, ClassVar
 
 import anyio
-import kuzu
 from watchdog.events import FileSystemEvent, FileSystemEventHandler
 from watchdog.observers import Observer
 
+from shotgun.codebase.core.kuzu_compat import get_kuzu
+
+if TYPE_CHECKING:
+    import real_ladybug as kuzu
+
+from shotgun.codebase.core.errors import (
+    DatabaseIssue,
+    KuzuErrorType,
+    classify_kuzu_error,
+)
 from shotgun.codebase.models import (
     CodebaseGraph,
     FileChange,
     GraphStatus,
+    NodeLabel,
     OperationStats,
+    RelationshipType,
 )
 from shotgun.logging_config import get_logger
 
@@ -190,7 +202,7 @@ class CodebaseGraphManager:
         return cls._lock
 
     @classmethod
-    def _generate_graph_id(cls, repo_path: str) -> str:
+    def generate_graph_id(cls, repo_path: str) -> str:
         """Generate deterministic graph ID from repository path."""
         normalized = str(Path(repo_path).resolve())
         return hashlib.sha256(normalized.encode()).hexdigest()[:12]
@@ -265,7 +277,8 @@ class CodebaseGraphManager:
         """
         graph_path = self.storage_dir / f"{graph_id}.kuzu"
 
-        # Create database and connection
+        # Create database and connection (lazy import for Windows compatibility)
+        kuzu = get_kuzu()
         lock = await self._get_lock()
         async with lock:
             db = kuzu.Database(str(graph_path))
@@ -348,7 +361,7 @@ class CodebaseGraphManager:
             Created graph metadata
         """
         repo_path = str(Path(repo_path).resolve())
-        graph_id = self._generate_graph_id(repo_path)
+        graph_id = self.generate_graph_id(repo_path)
 
         # Use repository name as default name
         if not name:
@@ -412,7 +425,8 @@ class CodebaseGraphManager:
         # Run build in thread pool
         await anyio.to_thread.run_sync(ingestor.build_graph_from_directory, repo_path)
 
-        # Get statistics
+        # Get statistics (lazy import for Windows compatibility)
+        kuzu = get_kuzu()
         lock = await self._get_lock()
         async with lock:
             db = kuzu.Database(str(graph_path))
@@ -554,6 +568,8 @@ class CodebaseGraphManager:
             "relationships_removed": 0,
         }
 
+        # Lazy import for Windows compatibility
+        kuzu = get_kuzu()
         lock = await self._get_lock()
         async with lock:
             if graph_id not in self._connections:
@@ -622,6 +638,8 @@ class CodebaseGraphManager:
         languages = build_options.get("languages")
         exclude_patterns = build_options.get("exclude_patterns")
 
+        # Lazy import for Windows compatibility
+        kuzu = get_kuzu()
         lock = await self._get_lock()
         async with lock:
             if graph_id not in self._connections:
@@ -1044,6 +1062,8 @@ class CodebaseGraphManager:
         self, graph_id: str, query: str, parameters: dict[str, Any] | None = None
     ) -> list[dict[str, Any]]:
         """Internal query execution with connection management."""
+        # Lazy import for Windows compatibility
+        kuzu = get_kuzu()
         lock = await self._get_lock()
         async with lock:
             if graph_id not in self._connections:
@@ -1216,8 +1236,19 @@ class CodebaseGraphManager:
                 indexed_from_cwds=indexed_from_cwds,
             )
         except Exception as e:
+            # Classify the error to determine if we should re-raise
+            error_type = classify_kuzu_error(e)
+
+            if error_type == KuzuErrorType.LOCKED:
+                # Don't mask lock errors - let caller handle them
+                logger.warning(
+                    f"Database locked - graph_id: {graph_id}, error: {str(e)}"
+                )
+                raise
+
             logger.error(
-                f"Failed to get graph metadata - graph_id: {graph_id}, error: {str(e)}"
+                f"Failed to get graph metadata - graph_id: {graph_id}, "
+                f"error_type: {error_type.value}, error: {str(e)}"
             )
             return None
 
@@ -1264,6 +1295,7 @@ class CodebaseGraphManager:
 
                         # Try to open the database
                         def _open_and_query(g: str = gid, p: Path = db_path) -> bool:
+                            kuzu = get_kuzu()
                             db = kuzu.Database(str(p))
                             conn = kuzu.Connection(db)
                             try:
@@ -1336,6 +1368,174 @@ class CodebaseGraphManager:
 
         return removed_graphs
 
+    async def _try_open_database(self, graph_id: str, db_path: Path) -> bool:
+        """Try to open a database and verify it has a Project node.
+
+        Args:
+            graph_id: The graph identifier
+            db_path: Path to the database file
+
+        Returns:
+            True if database opened and has Project node, False otherwise
+        """
+        lock = await self._get_lock()
+        async with lock:
+            # Close existing connections if any
+            if graph_id in self._connections:
+                try:
+                    self._connections[graph_id].close()
+                except Exception as e:
+                    logger.debug(f"Failed to close connection for {graph_id}: {e}")
+                del self._connections[graph_id]
+            if graph_id in self._databases:
+                try:
+                    self._databases[graph_id].close()
+                except Exception as e:
+                    logger.debug(f"Failed to close database for {graph_id}: {e}")
+                del self._databases[graph_id]
+
+            def _open_and_query() -> bool:
+                kuzu = get_kuzu()
+                db = kuzu.Database(str(db_path))
+                conn = kuzu.Connection(db)
+                try:
+                    result = conn.execute(
+                        "MATCH (p:Project {graph_id: $graph_id}) RETURN p",
+                        {"graph_id": graph_id},
+                    )
+                    return result.has_next() if hasattr(result, "has_next") else False
+                finally:
+                    conn.close()
+                    db.close()
+
+            return await anyio.to_thread.run_sync(_open_and_query)
+
+    async def _check_single_database(
+        self, graph_id: str, path: Path, timeout_seconds: float
+    ) -> DatabaseIssue | None:
+        """Check a single database for issues.
+
+        Args:
+            graph_id: The graph identifier
+            path: Path to the database file
+            timeout_seconds: How long to wait for the database to respond
+
+        Returns:
+            DatabaseIssue if problem found, None if database is healthy
+        """
+        try:
+            has_project = await asyncio.wait_for(
+                self._try_open_database(graph_id, path), timeout=timeout_seconds
+            )
+            if not has_project:
+                return DatabaseIssue(
+                    graph_id=graph_id,
+                    graph_path=path,
+                    error_type=KuzuErrorType.SCHEMA,
+                    message="Database has no Project node (incomplete build)",
+                )
+            return None
+
+        except asyncio.TimeoutError:
+            return DatabaseIssue(
+                graph_id=graph_id,
+                graph_path=path,
+                error_type=KuzuErrorType.TIMEOUT,
+                message=f"Database operation timed out after {timeout_seconds}s",
+            )
+
+        except Exception as e:
+            error_type = classify_kuzu_error(e)
+            logger.debug(f"Detected {error_type.value} issue with {graph_id}: {e}")
+            return DatabaseIssue(
+                graph_id=graph_id,
+                graph_path=path,
+                error_type=error_type,
+                message=str(e),
+            )
+
+    async def detect_database_issues(
+        self, timeout_seconds: float = 10.0
+    ) -> list[DatabaseIssue]:
+        """Detect issues with Kuzu databases without deleting them.
+
+        This method iterates through all .kuzu files in the storage directory,
+        attempts to open them, and returns information about any issues found.
+        Unlike cleanup_corrupted_databases(), this method does NOT delete anything -
+        it only detects and reports issues for the caller to handle.
+
+        Args:
+            timeout_seconds: How long to wait for each database to respond.
+                            Default is 10s; use 90s for retry with large codebases.
+
+        Returns:
+            List of DatabaseIssue objects describing any problems found
+        """
+        issues: list[DatabaseIssue] = []
+
+        for path in self.storage_dir.glob("*.kuzu"):
+            graph_id = path.stem
+            issue = await self._check_single_database(graph_id, path, timeout_seconds)
+            if issue:
+                issues.append(issue)
+
+        return issues
+
+    async def delete_database(self, graph_id: str) -> bool:
+        """Delete a database file and its WAL file.
+
+        Args:
+            graph_id: The ID of the graph to delete
+
+        Returns:
+            True if deletion was successful, False otherwise
+        """
+        graph_path = self.storage_dir / f"{graph_id}.kuzu"
+
+        try:
+            # Clean up any open connections
+            lock = await self._get_lock()
+            async with lock:
+                if graph_id in self._connections:
+                    try:
+                        self._connections[graph_id].close()
+                    except Exception as e:
+                        logger.debug(
+                            f"Failed to close connection during delete for {graph_id}: {e}"
+                        )
+                    del self._connections[graph_id]
+                if graph_id in self._databases:
+                    try:
+                        self._databases[graph_id].close()
+                    except Exception as e:
+                        logger.debug(
+                            f"Failed to close database during delete for {graph_id}: {e}"
+                        )
+                    del self._databases[graph_id]
+
+            # Remove the database (could be file or directory)
+            if graph_path.exists():
+                if graph_path.is_dir():
+                    await anyio.to_thread.run_sync(shutil.rmtree, graph_path)
+                else:
+                    await anyio.to_thread.run_sync(graph_path.unlink)
+
+                # Also delete WAL file if it exists
+                wal_path = graph_path.with_suffix(graph_path.suffix + ".wal")
+                if wal_path.exists():
+                    await anyio.to_thread.run_sync(wal_path.unlink)
+                    logger.debug(f"Deleted WAL file: {wal_path}")
+
+                logger.info(f"Deleted database: {graph_id}")
+                return True
+            else:
+                logger.warning(f"Database file not found for deletion: {graph_id}")
+                return False
+
+        except Exception as e:
+            logger.error(f"Failed to delete database {graph_id}: {e}")
+            return False
+
     async def list_graphs(self) -> list[CodebaseGraph]:
         """List all available graphs.
 
@@ -1477,20 +1677,20 @@ class CodebaseGraphManager:
         Returns:
             Tuple of (node_stats, relationship_stats)
         """
-        node_stats = {}
+        node_stats: dict[str, int] = {}
 
-        # Count each node type
+        # Count each node type (excluding ExternalPackage which is rarely needed)
         node_types = [
-            "Project",
-            "Package",
-            "Module",
-            "Class",
-            "Function",
-            "Method",
-            "File",
-            "Folder",
-            "FileMetadata",
-            "DeletionLog",
+            NodeLabel.PROJECT,
+            NodeLabel.PACKAGE,
+            NodeLabel.MODULE,
+            NodeLabel.CLASS,
+            NodeLabel.FUNCTION,
+            NodeLabel.METHOD,
+            NodeLabel.FILE,
+            NodeLabel.FOLDER,
+            NodeLabel.FILE_METADATA,
+            NodeLabel.DELETION_LOG,
         ]
 
         for node_type in node_types:
@@ -1505,14 +1705,14 @@ class CodebaseGraphManager:
                 logger.debug(f"Failed to count {node_type} nodes: {e}")
 
         # Count relationships - need to handle multiple tables for each type
-        rel_counts = {}
+        rel_counts: dict[str, int] = {}
 
         # CONTAINS relationships
         for prefix in [
-            "CONTAINS_PACKAGE",
-            "CONTAINS_FOLDER",
-            "CONTAINS_FILE",
-            "CONTAINS_MODULE",
+            RelationshipType.CONTAINS_PACKAGE,
+            RelationshipType.CONTAINS_FOLDER,
+            RelationshipType.CONTAINS_FILE,
+            RelationshipType.CONTAINS_MODULE,
         ]:
             count = 0
             for suffix in ["", "_PKG", "_FOLDER"]:
@@ -1530,13 +1730,13 @@ class CodebaseGraphManager:
 
         # Other relationships
         for rel_type in [
-            "DEFINES",
-            "DEFINES_FUNC",
-            "DEFINES_METHOD",
-            "INHERITS",
-            "OVERRIDES",
-            "DEPENDS_ON_EXTERNAL",
-            "IMPORTS",
+            RelationshipType.DEFINES,
+            RelationshipType.DEFINES_FUNC,
+            RelationshipType.DEFINES_METHOD,
+            RelationshipType.INHERITS,
+            RelationshipType.OVERRIDES,
+            RelationshipType.DEPENDS_ON_EXTERNAL,
+            RelationshipType.IMPORTS,
         ]:
             try:
                 result = await self._execute_query(
@@ -1549,7 +1749,12 @@ class CodebaseGraphManager:
 
         # CALLS relationships (multiple tables)
         calls_count = 0
-        for table in ["CALLS", "CALLS_FM", "CALLS_MF", "CALLS_MM"]:
+        for table in [
+            RelationshipType.CALLS,
+            RelationshipType.CALLS_FM,
+            RelationshipType.CALLS_MF,
+            RelationshipType.CALLS_MM,
+        ]:
             try:
                 result = await self._execute_query(
                     graph_id, f"MATCH ()-[r:{table}]->() RETURN COUNT(r) as count"
@@ -1563,16 +1768,21 @@ class CodebaseGraphManager:
 
         # TRACKS relationships
         tracks_count = 0
-        for entity in ["Module", "Class", "Function", "Method"]:
+        for tracks_rel in [
+            RelationshipType.TRACKS_MODULE,
+            RelationshipType.TRACKS_CLASS,
+            RelationshipType.TRACKS_FUNCTION,
+            RelationshipType.TRACKS_METHOD,
+        ]:
             try:
                 result = await self._execute_query(
                     graph_id,
-                    f"MATCH ()-[r:TRACKS_{entity}]->() RETURN COUNT(r) as count",
+                    f"MATCH ()-[r:{tracks_rel}]->() RETURN COUNT(r) as count",
                 )
                 if result:
                     tracks_count += result[0]["count"]
             except Exception as e:
-                logger.debug(f"Failed to count TRACKS_{entity} relationships: {e}")
+                logger.debug(f"Failed to count {tracks_rel} relationships: {e}")
         if tracks_count > 0:
             rel_counts["TRACKS (total)"] = tracks_count
 
@@ -1586,16 +1796,16 @@ class CodebaseGraphManager:
 
         # Print node stats
         for node_type in [
-            "Project",
-            "Package",
-            "Module",
-            "Class",
-            "Function",
-            "Method",
-            "File",
-            "Folder",
-            "FileMetadata",
-            "DeletionLog",
+            NodeLabel.PROJECT,
+            NodeLabel.PACKAGE,
+            NodeLabel.MODULE,
+            NodeLabel.CLASS,
+            NodeLabel.FUNCTION,
+            NodeLabel.METHOD,
+            NodeLabel.FILE,
+            NodeLabel.FOLDER,
+            NodeLabel.FILE_METADATA,
+            NodeLabel.DELETION_LOG,
         ]:
             count = node_stats.get(node_type, 0)
             logger.info(f"{node_type}: {count}")
@@ -1781,7 +1991,7 @@ class CodebaseGraphManager:
             Graph ID of the graph being built
         """
         repo_path = str(Path(repo_path).resolve())
-        graph_id = self._generate_graph_id(repo_path)
+        graph_id = self.generate_graph_id(repo_path)
 
         # Use repository name as default name
         if not name: