shotgun-sh 0.1.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

shotgun/codebase/core/ingestor.py

@@ -1,5 +1,6 @@
 """Kuzu graph ingestor for building code knowledge graphs."""
 
+import asyncio
 import hashlib
 import os
 import time
@@ -8,6 +9,7 @@ from collections import defaultdict
 from pathlib import Path
 from typing import Any
 
+import aiofiles
 import kuzu
 from tree_sitter import Node, Parser, QueryCursor
 
@@ -18,15 +20,12 @@ from shotgun.logging_config import get_logger
 logger = get_logger(__name__)
 
 
-# Default ignore patterns
-IGNORE_PATTERNS = {
+# Directories that should never be traversed during indexing
+BASE_IGNORE_DIRECTORIES = {
     ".git",
     "venv",
     ".venv",
     "__pycache__",
-    "node_modules",
-    "build",
-    "dist",
     ".eggs",
     ".pytest_cache",
     ".mypy_cache",
@@ -36,6 +35,46 @@ IGNORE_PATTERNS = {
     ".vscode",
 }
 
+# Well-known build output directories to skip when determining source files
+BUILD_ARTIFACT_DIRECTORIES = {
+    "node_modules",
+    ".next",
+    ".nuxt",
+    ".vite",
+    ".yarn",
+    ".svelte-kit",
+    ".output",
+    ".turbo",
+    ".parcel-cache",
+    ".vercel",
+    ".serverless",
+    "build",
+    "dist",
+    "out",
+    "tmp",
+    "coverage",
+}
+
+# Default ignore patterns combines base directories and build artifacts
+IGNORE_PATTERNS = BASE_IGNORE_DIRECTORIES | BUILD_ARTIFACT_DIRECTORIES
+
+# Directory prefixes that should always be ignored
+IGNORED_DIRECTORY_PREFIXES = (".",)
+
+
+def should_ignore_directory(name: str, ignore_patterns: set[str] | None = None) -> bool:
+    """Return True if the directory name should be ignored."""
+    patterns = IGNORE_PATTERNS if ignore_patterns is None else ignore_patterns
+    if name in patterns:
+        return True
+    return name.startswith(IGNORED_DIRECTORY_PREFIXES)
+
+
+def is_path_ignored(path: Path, ignore_patterns: set[str] | None = None) -> bool:
+    """Return True if any part of the path should be ignored."""
+    patterns = IGNORE_PATTERNS if ignore_patterns is None else ignore_patterns
+    return any(should_ignore_directory(part, patterns) for part in path.parts)
+
 
 class Ingestor:
     """Handles all communication and ingestion with the Kuzu database."""
@@ -535,6 +574,7 @@ class SimpleGraphBuilder:
         parsers: dict[str, Parser],
         queries: dict[str, Any],
         exclude_patterns: list[str] | None = None,
+        progress_callback: Any | None = None,
     ):
         self.ingestor = ingestor
         self.repo_path = repo_path
@@ -544,6 +584,7 @@ class SimpleGraphBuilder:
         self.ignore_dirs = IGNORE_PATTERNS
         if exclude_patterns:
             self.ignore_dirs = self.ignore_dirs.union(set(exclude_patterns))
+        self.progress_callback = progress_callback
 
         # Caches
         self.structural_elements: dict[Path, str | None] = {}
@@ -552,7 +593,35 @@ class SimpleGraphBuilder:
         self.simple_name_lookup: dict[str, set[str]] = defaultdict(set)
         self.class_inheritance: dict[str, list[str]] = {}  # class_qn -> [parent_qns]
 
-    def run(self) -> None:
+    def _report_progress(
+        self,
+        phase: str,
+        phase_name: str,
+        current: int,
+        total: int | None = None,
+        phase_complete: bool = False,
+    ) -> None:
+        """Report progress via callback if available."""
+        if not self.progress_callback:
+            return
+
+        try:
+            # Import here to avoid circular dependency
+            from shotgun.codebase.models import IndexProgress, ProgressPhase
+
+            progress = IndexProgress(
+                phase=ProgressPhase(phase),
+                phase_name=phase_name,
+                current=current,
+                total=total,
+                phase_complete=phase_complete,
+            )
+            self.progress_callback(progress)
+        except Exception as e:
+            # Don't let progress callback errors crash the build
+            logger.debug(f"Progress callback error: {e}")
+
+    async def run(self) -> None:
         """Run the three-pass graph building process."""
         logger.info(f"Building graph for project: {self.project_name}")
 
@@ -562,7 +631,7 @@ class SimpleGraphBuilder:
 
         # Pass 2: Definitions
         logger.info("Pass 2: Processing files and extracting definitions...")
-        self._process_files()
+        await self._process_files()
 
         # Pass 3: Relationships
         logger.info("Pass 3: Processing relationships (calls, imports)...")
@@ -575,8 +644,11 @@ class SimpleGraphBuilder:
 
     def _identify_structure(self) -> None:
         """First pass: Walk directory to find packages and folders."""
+        dir_count = 0
         for root_str, dirs, _ in os.walk(self.repo_path, topdown=True):
-            dirs[:] = [d for d in dirs if d not in self.ignore_dirs]
+            dirs[:] = [
+                d for d in dirs if not should_ignore_directory(d, self.ignore_dirs)
+            ]
             root = Path(root_str)
             relative_root = root.relative_to(self.repo_path)
 
@@ -584,6 +656,13 @@ class SimpleGraphBuilder:
             if root == self.repo_path:
                 continue
 
+            dir_count += 1
+            # Report progress every 10 directories
+            if dir_count % 10 == 0:
+                self._report_progress(
+                    "structure", "Identifying packages and folders", dir_count
+                )
+
             parent_rel_path = relative_root.parent
             parent_container_qn = self.structural_elements.get(parent_rel_path)
 
@@ -686,14 +765,40 @@ class SimpleGraphBuilder:
 
                 self.structural_elements[relative_root] = None
 
-    def _process_files(self) -> None:
+        # Report phase completion
+        self._report_progress(
+            "structure",
+            "Identifying packages and folders",
+            dir_count,
+            phase_complete=True,
+        )
+
+    async def _process_files(self) -> None:
         """Second pass: Process files and extract definitions."""
+        # First pass: Count total files
+        total_files = 0
+        for root_str, _, files in os.walk(self.repo_path):
+            root = Path(root_str)
+
+            # Skip ignored directories
+            if is_path_ignored(root, self.ignore_dirs):
+                continue
+
+            for filename in files:
+                filepath = root / filename
+                ext = filepath.suffix
+                lang_config = get_language_config(ext)
+
+                if lang_config and lang_config.name in self.parsers:
+                    total_files += 1
+
+        # Second pass: Process files with progress reporting
         file_count = 0
         for root_str, _, files in os.walk(self.repo_path):
             root = Path(root_str)
 
             # Skip ignored directories
-            if any(part in self.ignore_dirs for part in root.parts):
+            if is_path_ignored(root, self.ignore_dirs):
                 continue
 
             for filename in files:
@@ -704,15 +809,32 @@ class SimpleGraphBuilder:
                 lang_config = get_language_config(ext)
 
                 if lang_config and lang_config.name in self.parsers:
-                    self._process_single_file(filepath, lang_config.name)
+                    await self._process_single_file(filepath, lang_config.name)
                     file_count += 1
 
+                    # Report progress after each file
+                    self._report_progress(
+                        "definitions",
+                        "Processing files and extracting definitions",
+                        file_count,
+                        total_files,
+                    )
+
                     if file_count % 100 == 0:
-                        logger.info(f"  Processed {file_count} files...")
+                        logger.info(f"  Processed {file_count}/{total_files} files...")
+
+        logger.info(f"  Total files processed: {file_count}/{total_files}")
 
-        logger.info(f"  Total files processed: {file_count}")
+        # Report phase completion
+        self._report_progress(
+            "definitions",
+            "Processing files and extracting definitions",
+            file_count,
+            total_files,
+            phase_complete=True,
+        )
 
-    def _process_single_file(self, filepath: Path, language: str) -> None:
+    async def _process_single_file(self, filepath: Path, language: str) -> None:
         """Process a single file."""
         relative_path = filepath.relative_to(self.repo_path)
         relative_path_str = str(relative_path).replace(os.sep, "/")
@@ -753,8 +875,8 @@ class SimpleGraphBuilder:
 
         # Parse file
         try:
-            with open(filepath, "rb") as f:
-                content = f.read()
+            async with aiofiles.open(filepath, "rb") as f:
+                content = await f.read()
 
             parser = self.parsers[language]
             tree = parser.parse(content)
@@ -1143,7 +1265,8 @@ class SimpleGraphBuilder:
         self._process_inheritance()
 
         # Then process function calls
-        logger.info(f"Processing function calls for {len(self.ast_cache)} files...")
+        total_files = len(self.ast_cache)
+        logger.info(f"Processing function calls for {total_files} files...")
         logger.info(f"Function registry has {len(self.function_registry)} entries")
         logger.info(
             f"Simple name lookup has {len(self.simple_name_lookup)} unique names"
@@ -1157,10 +1280,29 @@ class SimpleGraphBuilder:
                     f"  Example: '{name}' -> {list(self.simple_name_lookup[name])[:3]}"
                 )
 
+        file_count = 0
         for filepath, (root_node, language) in self.ast_cache.items():
             self._process_calls(filepath, root_node, language)
             # NOTE: Add import processing. wtf does this mean?
 
+            file_count += 1
+            # Report progress after each file
+            self._report_progress(
+                "relationships",
+                "Processing relationships (calls, imports)",
+                file_count,
+                total_files,
+            )
+
+        # Report phase completion
+        self._report_progress(
+            "relationships",
+            "Processing relationships (calls, imports)",
+            file_count,
+            total_files,
+            phase_complete=True,
+        )
+
     def _process_inheritance(self) -> None:
         """Process inheritance relationships between classes."""
         logger.info("Processing inheritance relationships...")
@@ -1444,6 +1586,7 @@ class CodebaseIngestor:
         db_path: str,
         project_name: str | None = None,
         exclude_patterns: list[str] | None = None,
+        progress_callback: Any | None = None,
     ):
         """Initialize the ingestor.
 
@@ -1451,10 +1594,12 @@ class CodebaseIngestor:
             db_path: Path to Kuzu database
             project_name: Optional project name
             exclude_patterns: Patterns to exclude from processing
+            progress_callback: Optional callback for progress reporting
         """
         self.db_path = Path(db_path)
         self.project_name = project_name
         self.exclude_patterns = exclude_patterns or []
+        self.progress_callback = progress_callback
 
     def build_graph_from_directory(self, repo_path: str) -> None:
         """Build a code knowledge graph from a directory.
@@ -1484,11 +1629,16 @@ class CodebaseIngestor:
 
             # Build graph
             builder = SimpleGraphBuilder(
-                ingestor, repo_path_obj, parsers, queries, self.exclude_patterns
+                ingestor,
+                repo_path_obj,
+                parsers,
+                queries,
+                self.exclude_patterns,
+                self.progress_callback,
             )
             if self.project_name:
                 builder.project_name = self.project_name
-            builder.run()
+            asyncio.run(builder.run())
 
             logger.info(f"Graph successfully created at: {self.db_path}")
 

shotgun/codebase/core/manager.py

@@ -51,9 +51,13 @@ class CodebaseFileHandler(FileSystemEventHandler):
         self.pending_changes: list[FileChange] = []
         self._lock = anyio.Lock()
         # Import default ignore patterns from ingestor
-        from shotgun.codebase.core.ingestor import IGNORE_PATTERNS
+        from shotgun.codebase.core.ingestor import (
+            IGNORE_PATTERNS,
+            should_ignore_directory,
+        )
 
         self.ignore_patterns = ignore_patterns or IGNORE_PATTERNS
+        self._should_ignore_directory = should_ignore_directory
 
     def on_any_event(self, event: FileSystemEvent) -> None:
         """Handle any file system event."""
@@ -71,7 +75,7 @@ class CodebaseFileHandler(FileSystemEventHandler):
 
         # Check if any parent directory should be ignored
         for parent in path.parents:
-            if parent.name in self.ignore_patterns:
+            if self._should_ignore_directory(parent.name, self.ignore_patterns):
                 logger.debug(
                     f"Ignoring file in ignored directory: {parent.name} - path: {src_path_str}"
                 )
@@ -106,7 +110,7 @@ class CodebaseFileHandler(FileSystemEventHandler):
             )
             dest_path = Path(dest_path_str)
             for parent in dest_path.parents:
-                if parent.name in self.ignore_patterns:
+                if self._should_ignore_directory(parent.name, self.ignore_patterns):
                     logger.debug(
                         f"Ignoring move to ignored directory: {parent.name} - dest_path: {dest_path_str}"
                     )
@@ -329,6 +333,7 @@ class CodebaseGraphManager:
         languages: list[str] | None = None,
         exclude_patterns: list[str] | None = None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Build a new code knowledge graph.
 
@@ -337,6 +342,7 @@ class CodebaseGraphManager:
             name: Optional human-readable name
             languages: Languages to parse (default: all supported)
             exclude_patterns: Patterns to exclude
+            progress_callback: Optional callback for progress reporting
 
         Returns:
             Created graph metadata
@@ -353,7 +359,28 @@ class CodebaseGraphManager:
 
         # Check if graph already exists
         if graph_path.exists():
-            raise CodebaseAlreadyIndexedError(repo_path)
+            # Verify it's not corrupted by checking if we can load the Project node
+            existing_graph = await self.get_graph(graph_id)
+            if existing_graph:
+                # Valid existing graph
+                raise CodebaseAlreadyIndexedError(repo_path)
+            else:
+                # Corrupted database - remove and re-index
+                logger.warning(
+                    f"Found corrupted database at {graph_path}, removing for re-indexing..."
+                )
+                import shutil
+
+                # Handle both files and directories (kuzu v0.11.2+ uses files)
+                if graph_path.is_file():
+                    graph_path.unlink()  # Delete file
+                    # Also delete WAL file if it exists
+                    wal_path = graph_path.with_suffix(graph_path.suffix + ".wal")
+                    if wal_path.exists():
+                        wal_path.unlink()
+                        logger.debug(f"Deleted WAL file: {wal_path}")
+                else:
+                    shutil.rmtree(graph_path)  # Delete directory
 
         # Import the builder from local core module
         from shotgun.codebase.core import CodebaseIngestor
@@ -379,6 +406,7 @@ class CodebaseGraphManager:
             db_path=str(graph_path),
             project_name=name,
             exclude_patterns=exclude_patterns or [],
+            progress_callback=progress_callback,
         )
 
         # Run build in thread pool
@@ -741,7 +769,7 @@ class CodebaseGraphManager:
 
                     lang_config = get_language_config(full_path.suffix)
                     if lang_config and lang_config.name in parsers:
-                        builder._process_single_file(full_path, lang_config.name)
+                        await builder._process_single_file(full_path, lang_config.name)
                         stats["nodes_modified"] += 1  # Approximate
 
             # Process additions
@@ -756,7 +784,7 @@ class CodebaseGraphManager:
 
                     lang_config = get_language_config(full_path.suffix)
                     if lang_config and lang_config.name in parsers:
-                        builder._process_single_file(full_path, lang_config.name)
+                        await builder._process_single_file(full_path, lang_config.name)
                         stats["nodes_added"] += 1  # Approximate
 
             # Flush all pending operations
@@ -1193,6 +1221,121 @@ class CodebaseGraphManager:
             )
             return None
 
+    async def cleanup_corrupted_databases(self) -> list[str]:
+        """Detect and remove corrupted Kuzu databases.
+
+        This method iterates through all .kuzu files in the storage directory,
+        attempts to open them, and removes any that are corrupted or unreadable.
+
+        Returns:
+            List of graph_ids that were removed due to corruption
+        """
+        import shutil
+
+        removed_graphs = []
+
+        # Find all .kuzu databases (files in v0.11.2, directories in newer versions)
+        for path in self.storage_dir.glob("*.kuzu"):
+            graph_id = path.stem
+
+            # Try to open and validate the database
+            try:
+                # Try to open the database with a timeout to prevent hanging
+                async def try_open_database(
+                    gid: str = graph_id, db_path: Path = path
+                ) -> bool:
+                    lock = await self._get_lock()
+                    async with lock:
+                        # Close existing connections if any
+                        if gid in self._connections:
+                            try:
+                                self._connections[gid].close()
+                            except Exception as e:
+                                logger.debug(
+                                    f"Failed to close connection for {gid}: {e}"
+                                )
+                            del self._connections[gid]
+                        if gid in self._databases:
+                            try:
+                                self._databases[gid].close()
+                            except Exception as e:
+                                logger.debug(f"Failed to close database for {gid}: {e}")
+                            del self._databases[gid]
+
+                        # Try to open the database
+                        def _open_and_query(g: str = gid, p: Path = db_path) -> bool:
+                            db = kuzu.Database(str(p))
+                            conn = kuzu.Connection(db)
+                            try:
+                                result = conn.execute(
+                                    "MATCH (p:Project {graph_id: $graph_id}) RETURN p",
+                                    {"graph_id": g},
+                                )
+                                has_results = (
+                                    result.has_next()
+                                    if hasattr(result, "has_next")
+                                    else False
+                                )
+                                return has_results
+                            finally:
+                                conn.close()
+                                db.close()
+
+                        return await anyio.to_thread.run_sync(_open_and_query)
+
+                # Try to open with 5 second timeout
+                has_project = await asyncio.wait_for(try_open_database(), timeout=5.0)
+
+                if not has_project:
+                    # Database exists but has no Project node - consider it corrupted
+                    raise ValueError("No Project node found in database")
+
+            except (Exception, asyncio.TimeoutError) as e:
+                # Database is corrupted or timed out - remove it
+                error_type = (
+                    "timed out" if isinstance(e, asyncio.TimeoutError) else "corrupted"
+                )
+                logger.warning(
+                    f"Detected {error_type} database at {path}, removing it. "
+                    f"Error: {str(e) if not isinstance(e, asyncio.TimeoutError) else 'Operation timed out after 5 seconds'}"
+                )
+
+                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 cleanup 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 cleanup for {graph_id}: {e}"
+                                )
+                            del self._databases[graph_id]
+
+                    # Remove the database (could be file or directory)
+                    if path.is_dir():
+                        await anyio.to_thread.run_sync(shutil.rmtree, path)
+                    else:
+                        await anyio.to_thread.run_sync(path.unlink)
+                    removed_graphs.append(graph_id)
+                    logger.info(f"Removed {error_type} database: {graph_id}")
+
+                except Exception as cleanup_error:
+                    logger.error(
+                        f"Failed to remove corrupted database {graph_id}: {cleanup_error}"
+                    )
+
+        return removed_graphs
+
     async def list_graphs(self) -> list[CodebaseGraph]:
         """List all available graphs.
 
@@ -1201,7 +1344,7 @@ class CodebaseGraphManager:
         """
         graphs = []
 
-        # Find all .kuzu files
+        # Find all .kuzu database files (Kuzu v0.11.2 creates files, not directories)
         for path in self.storage_dir.glob("*.kuzu"):
             if path.is_file():
                 graph_id = path.stem
@@ -1294,6 +1437,8 @@ class CodebaseGraphManager:
         Args:
             graph_id: Graph to delete
         """
+        import shutil
+
         # Stop watcher if running
         if graph_id in self._watchers:
             await self.stop_watcher(graph_id)
@@ -1308,11 +1453,14 @@ class CodebaseGraphManager:
                 self._databases[graph_id].close()
                 del self._databases[graph_id]
 
-        # Delete files
+        # Delete database (files in v0.11.2, directories in newer versions)
         graph_path = self.storage_dir / f"{graph_id}.kuzu"
         if graph_path.exists():
-            # Delete the database file
-            await anyio.to_thread.run_sync(graph_path.unlink)
+            if graph_path.is_dir():
+                await anyio.to_thread.run_sync(shutil.rmtree, graph_path)
+            else:
+                # File-based database (Kuzu v0.11.2)
+                await anyio.to_thread.run_sync(graph_path.unlink)
 
         # Also delete the WAL file if it exists
         wal_path = self.storage_dir / f"{graph_id}.kuzu.wal"
@@ -1464,6 +1612,7 @@ class CodebaseGraphManager:
         languages: list[str] | None,
         exclude_patterns: list[str] | None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Internal implementation of graph building (runs in background)."""
         operation_id = str(uuid.uuid4())
@@ -1487,7 +1636,13 @@ class CodebaseGraphManager:
 
             # Do the actual build work
             graph = await self._do_build_graph(
-                graph_id, repo_path, name, languages, exclude_patterns, indexed_from_cwd
+                graph_id,
+                repo_path,
+                name,
+                languages,
+                exclude_patterns,
+                indexed_from_cwd,
+                progress_callback,
             )
 
             # Update operation stats
@@ -1536,6 +1691,7 @@ class CodebaseGraphManager:
         languages: list[str] | None,
         exclude_patterns: list[str] | None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> CodebaseGraph:
         """Execute the actual graph building logic (extracted from original build_graph)."""
         # The database and Project node already exist from _initialize_graph_metadata
@@ -1591,10 +1747,11 @@ class CodebaseGraphManager:
                 parsers=parsers,
                 queries=queries,
                 exclude_patterns=exclude_patterns,
+                progress_callback=progress_callback,
             )
 
             # Build the graph
-            builder.run()
+            asyncio.run(builder.run())
 
         # Run build in thread pool
         await anyio.to_thread.run_sync(_build_graph)
@@ -1616,6 +1773,7 @@ class CodebaseGraphManager:
         languages: list[str] | None = None,
         exclude_patterns: list[str] | None = None,
         indexed_from_cwd: str | None = None,
+        progress_callback: Any | None = None,
     ) -> str:
         """Start building a new code knowledge graph asynchronously.
 
@@ -1654,7 +1812,13 @@ class CodebaseGraphManager:
         # Start the build operation in background
         task = asyncio.create_task(
             self._build_graph_impl(
-                graph_id, repo_path, name, languages, exclude_patterns, indexed_from_cwd
+                graph_id,
+                repo_path,
+                name,
+                languages,
+                exclude_patterns,
+                indexed_from_cwd,
+                progress_callback,
             )
         )
         self._operations[graph_id] = task

shotgun/codebase/core/nl_query.py

@@ -34,7 +34,7 @@ async def llm_cypher_prompt(
     Returns:
         CypherGenerationResponse with cypher_query, can_generate flag, and reason if not
     """
-    model_config = get_provider_model()
+    model_config = await get_provider_model()
 
     # Create an agent with structured output for Cypher generation
     cypher_agent = Agent(