shotgun-sh 0.2.29.dev2__py3-none-any.whl → 0.6.1.dev1__py3-none-any.whl

shotgun/cli/spec/pull_service.py

@@ -1,16 +1,18 @@
 """Shared spec pull service for CLI and TUI."""
 
+import time
 from collections.abc import Callable
 from dataclasses import dataclass
 from datetime import datetime, timezone
 from pathlib import Path
 
 from shotgun.logging_config import get_logger
+from shotgun.posthog_telemetry import track_event
 from shotgun.shotgun_web.specs_client import SpecsClient
 from shotgun.shotgun_web.supabase_client import download_file_from_url
 
 from .backup import clear_shotgun_dir, create_backup
-from .models import SpecMeta
+from .models import PullPhase, PullSource, SpecMeta
 
 logger = get_logger(__name__)
 
@@ -53,6 +55,7 @@ class SpecPullService:
         shotgun_dir: Path,
         on_progress: Callable[[PullProgress], None] | None = None,
         is_cancelled: Callable[[], bool] | None = None,
+        source: PullSource = PullSource.CLI,
     ) -> PullResult:
         """Pull a spec version to the local directory.
 
@@ -61,10 +64,14 @@ class SpecPullService:
             shotgun_dir: Target directory (typically .shotgun/)
             on_progress: Optional callback for progress updates
             is_cancelled: Optional callback to check if cancelled
+            source: Source of the pull request (CLI or TUI)
 
         Returns:
             PullResult with success status and details
         """
+        start_time = time.time()
+        current_phase: PullPhase = PullPhase.STARTING
+        track_event("spec_pull_started", {"source": source.value})
 
         def report(
             phase: str,
@@ -83,83 +90,130 @@ class SpecPullService:
                 )
 
         def check_cancelled() -> None:
+            nonlocal current_phase
             if is_cancelled and is_cancelled():
+                track_event(
+                    "spec_pull_cancelled",
+                    {"source": source.value, "phase": current_phase.value},
+                )
                 raise CancelledError()
 
-        # Phase 1: Fetch version metadata
-        report("Fetching version info...")
-        check_cancelled()
+        try:
+            # Phase 1: Fetch version metadata
+            current_phase = PullPhase.FETCHING
+            report("Fetching version info...")
+            check_cancelled()
 
-        response = await self._client.get_version_with_files(version_id)
-        spec_name = response.spec_name
-        files = response.files
+            response = await self._client.get_version_with_files(version_id)
+            spec_name = response.spec_name
+            files = response.files
+
+            if not files:
+                track_event(
+                    "spec_pull_failed",
+                    {
+                        "source": source.value,
+                        "error_type": "EmptyVersion",
+                        "phase": current_phase.value,
+                    },
+                )
+                return PullResult(
+                    success=False,
+                    spec_name=spec_name,
+                    error="No files in this version.",
+                )
 
-        if not files:
-            return PullResult(
-                success=False,
-                spec_name=spec_name,
-                error="No files in this version.",
-            )
+            # Phase 2: Backup existing content
+            current_phase = PullPhase.BACKUP
+            backup_path: str | None = None
+            if shotgun_dir.exists():
+                report("Backing up existing files...")
+                check_cancelled()
+
+                backup_path = await create_backup(shotgun_dir)
+                if backup_path:
+                    clear_shotgun_dir(shotgun_dir)
+
+            # Ensure directory exists
+            shotgun_dir.mkdir(parents=True, exist_ok=True)
+
+            # Phase 3: Download files
+            current_phase = PullPhase.DOWNLOADING
+            total_files = len(files)
+            total_bytes = 0
+            for idx, file_info in enumerate(files):
+                check_cancelled()
+
+                report(
+                    f"Downloading files ({idx + 1}/{total_files})...",
+                    file_index=idx,
+                    total_files=total_files,
+                    current_file=file_info.relative_path,
+                )
 
-        # Phase 2: Backup existing content
-        backup_path: str | None = None
-        if shotgun_dir.exists():
-            report("Backing up existing files...")
-            check_cancelled()
+                if not file_info.download_url:
+                    logger.warning(
+                        "Skipping file without download URL: %s",
+                        file_info.relative_path,
+                    )
+                    continue
 
-            backup_path = await create_backup(shotgun_dir)
-            if backup_path:
-                clear_shotgun_dir(shotgun_dir)
+                content = await download_file_from_url(file_info.download_url)
+                total_bytes += file_info.size_bytes
 
-        # Ensure directory exists
-        shotgun_dir.mkdir(parents=True, exist_ok=True)
+                local_path = shotgun_dir / file_info.relative_path
+                local_path.parent.mkdir(parents=True, exist_ok=True)
+                local_path.write_bytes(content)
 
-        # Phase 3: Download files
-        total_files = len(files)
-        for idx, file_info in enumerate(files):
+            # Phase 4: Write meta.json
+            current_phase = PullPhase.FINALIZING
+            report("Finalizing...")
             check_cancelled()
 
-            report(
-                f"Downloading files ({idx + 1}/{total_files})...",
-                file_index=idx,
-                total_files=total_files,
-                current_file=file_info.relative_path,
+            meta = SpecMeta(
+                version_id=response.version.id,
+                spec_id=response.spec_id,
+                spec_name=response.spec_name,
+                workspace_id=response.workspace_id,
+                is_latest=response.version.is_latest,
+                pulled_at=datetime.now(timezone.utc),
+                backup_path=backup_path,
+                web_url=response.web_url,
+            )
+            meta_path = shotgun_dir / "meta.json"
+            meta_path.write_text(meta.model_dump_json(indent=2))
+
+            # Track successful completion
+            duration = time.time() - start_time
+            track_event(
+                "spec_pull_completed",
+                {
+                    "source": source.value,
+                    "file_count": total_files,
+                    "total_bytes": total_bytes,
+                    "duration_seconds": round(duration, 2),
+                    "had_backup": backup_path is not None,
+                },
             )
 
-            if not file_info.download_url:
-                logger.warning(
-                    "Skipping file without download URL: %s",
-                    file_info.relative_path,
-                )
-                continue
-
-            content = await download_file_from_url(file_info.download_url)
-
-            local_path = shotgun_dir / file_info.relative_path
-            local_path.parent.mkdir(parents=True, exist_ok=True)
-            local_path.write_bytes(content)
-
-        # Phase 4: Write meta.json
-        report("Finalizing...")
-        check_cancelled()
-
-        meta = SpecMeta(
-            version_id=response.version.id,
-            spec_id=response.spec_id,
-            spec_name=response.spec_name,
-            workspace_id=response.workspace_id,
-            is_latest=response.version.is_latest,
-            pulled_at=datetime.now(timezone.utc),
-            backup_path=backup_path,
-            web_url=response.web_url,
-        )
-        meta_path = shotgun_dir / "meta.json"
-        meta_path.write_text(meta.model_dump_json(indent=2))
-
-        return PullResult(
-            success=True,
-            spec_name=spec_name,
-            file_count=total_files,
-            backup_path=backup_path,
-            web_url=response.web_url,
-        )
+            return PullResult(
+                success=True,
+                spec_name=spec_name,
+                file_count=total_files,
+                backup_path=backup_path,
+                web_url=response.web_url,
+            )
+
+        except CancelledError:
+            # Already tracked in check_cancelled()
+            raise
+        except Exception as e:
+            track_event(
+                "spec_pull_failed",
+                {
+                    "source": source.value,
+                    "error_type": type(e).__name__,
+                    "phase": current_phase.value,
+                },
+            )
+            raise

shotgun/codebase/__init__.py

@@ -1,5 +1,6 @@
 """Shotgun codebase analysis and graph management."""
 
+from shotgun.codebase.indexing_state import IndexingState
 from shotgun.codebase.models import CodebaseGraph, GraphStatus, QueryResult, QueryType
 from shotgun.codebase.service import CodebaseService
 
@@ -7,6 +8,7 @@ __all__ = [
     "CodebaseService",
     "CodebaseGraph",
     "GraphStatus",
+    "IndexingState",
     "QueryResult",
     "QueryType",
 ]

shotgun/codebase/benchmarks/__init__.py

@@ -0,0 +1,35 @@
+"""Benchmark system for codebase indexing performance analysis.
+
+This package provides tools for running benchmarks and reporting metrics
+for the codebase indexing pipeline.
+"""
+
+from shotgun.codebase.benchmarks.benchmark_runner import BenchmarkRunner
+from shotgun.codebase.benchmarks.exporters import MetricsExporter
+from shotgun.codebase.benchmarks.formatters import (
+    JsonFormatter,
+    MarkdownFormatter,
+    MetricsDisplayOptions,
+    get_formatter,
+)
+from shotgun.codebase.benchmarks.models import (
+    BenchmarkConfig,
+    BenchmarkMode,
+    BenchmarkResults,
+    BenchmarkRun,
+    OutputFormat,
+)
+
+__all__ = [
+    "BenchmarkConfig",
+    "BenchmarkMode",
+    "BenchmarkResults",
+    "BenchmarkRun",
+    "BenchmarkRunner",
+    "JsonFormatter",
+    "MarkdownFormatter",
+    "MetricsDisplayOptions",
+    "MetricsExporter",
+    "OutputFormat",
+    "get_formatter",
+]

shotgun/codebase/benchmarks/benchmark_runner.py

@@ -0,0 +1,309 @@
+"""Benchmark runner for codebase indexing performance analysis.
+
+This module provides the BenchmarkRunner class for running benchmark iterations
+and collecting performance statistics.
+"""
+
+from __future__ import annotations
+
+import gc
+import hashlib
+import json
+import shutil
+import time
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+from shotgun.codebase.benchmarks.models import (
+    BenchmarkConfig,
+    BenchmarkMode,
+    BenchmarkResults,
+    BenchmarkRun,
+)
+from shotgun.codebase.core import Ingestor, SimpleGraphBuilder
+from shotgun.codebase.core.kuzu_compat import get_kuzu
+from shotgun.codebase.core.metrics_collector import MetricsCollector
+from shotgun.codebase.core.parser_loader import load_parsers
+from shotgun.logging_config import get_logger
+from shotgun.sdk.services import get_codebase_service
+from shotgun.utils.file_system_utils import get_shotgun_home
+
+logger = get_logger(__name__)
+
+
+def _compute_graph_id(codebase_path: Path) -> str:
+    """Compute a unique graph ID from the codebase path.
+
+    Args:
+        codebase_path: Path to the codebase
+
+    Returns:
+        A 12-character hex string identifying this codebase
+    """
+    return hashlib.sha256(str(codebase_path).encode()).hexdigest()[:12]
+
+
+class BenchmarkRunner:
+    """Runs benchmark iterations and collects statistics."""
+
+    def __init__(
+        self,
+        codebase_path: Path,
+        codebase_name: str,
+        iterations: int = 1,
+        warmup_iterations: int = 0,
+        parallel: bool = True,
+        worker_count: int | None = None,
+        collect_file_metrics: bool = True,
+        collect_worker_metrics: bool = True,
+        progress_callback: Callable[..., Any] | None = None,
+    ) -> None:
+        """Initialize benchmark runner.
+
+        Args:
+            codebase_path: Path to repository to benchmark
+            codebase_name: Human-readable name for the codebase
+            iterations: Number of measured benchmark runs
+            warmup_iterations: Number of warmup runs (not measured)
+            parallel: Whether to use parallel execution
+            worker_count: Number of workers (None = auto)
+            collect_file_metrics: Whether to collect per-file metrics
+            collect_worker_metrics: Whether to collect per-worker metrics
+            progress_callback: Optional callback for progress updates
+        """
+        self.codebase_path = codebase_path.resolve()
+        self.codebase_name = codebase_name
+        self.iterations = iterations
+        self.warmup_iterations = warmup_iterations
+        self.parallel = parallel
+        self.worker_count = worker_count
+        self.collect_file_metrics = collect_file_metrics
+        self.collect_worker_metrics = collect_worker_metrics
+        self.progress_callback = progress_callback
+
+        # Configuration object
+        self.config = BenchmarkConfig(
+            mode=BenchmarkMode.PARALLEL if parallel else BenchmarkMode.SEQUENTIAL,
+            worker_count=worker_count,
+            iterations=iterations,
+            warmup_iterations=warmup_iterations,
+            collect_file_metrics=collect_file_metrics,
+            collect_worker_metrics=collect_worker_metrics,
+        )
+
+        # Storage for database operations
+        self._storage_dir = get_shotgun_home() / "codebases"
+        self._service = get_codebase_service(self._storage_dir)
+
+    async def run(self) -> BenchmarkResults:
+        """Run all benchmark iterations and return aggregated results.
+
+        Returns:
+            BenchmarkResults with all run data and statistics
+        """
+        results = BenchmarkResults(
+            codebase_name=self.codebase_name,
+            codebase_path=str(self.codebase_path),
+            config=self.config,
+        )
+
+        # Run warmup iterations
+        for i in range(self.warmup_iterations):
+            logger.info(f"Running warmup iteration {i + 1}/{self.warmup_iterations}...")
+            if self.progress_callback:
+                self.progress_callback(
+                    f"Warmup {i + 1}/{self.warmup_iterations}", None, None
+                )
+
+            run = await self._run_single_iteration(
+                run_id=i,
+                is_warmup=True,
+            )
+            results.add_run(run)
+            await self._cleanup_database()
+
+        # Run measured iterations
+        for i in range(self.iterations):
+            logger.info(f"Running benchmark iteration {i + 1}/{self.iterations}...")
+            if self.progress_callback:
+                self.progress_callback(
+                    f"Benchmark {i + 1}/{self.iterations}", None, None
+                )
+
+            run = await self._run_single_iteration(
+                run_id=i,
+                is_warmup=False,
+            )
+            results.add_run(run)
+
+            # Clean up between iterations (but not after the last one)
+            if i < self.iterations - 1:
+                await self._cleanup_database()
+
+        # Register the codebase so it persists after benchmark
+        await self._register_codebase()
+
+        # Calculate statistics
+        results.calculate_statistics()
+
+        logger.info(
+            f"Benchmark complete: {self.iterations} iterations, "
+            f"avg {results.avg_duration_seconds:.2f}s"
+        )
+
+        return results
+
+    async def _run_single_iteration(
+        self,
+        run_id: int,
+        is_warmup: bool,
+    ) -> BenchmarkRun:
+        """Run a single benchmark iteration.
+
+        Args:
+            run_id: Run number
+            is_warmup: Whether this is a warmup run
+
+        Returns:
+            BenchmarkRun with collected metrics
+        """
+        # Create metrics collector
+        metrics_collector = MetricsCollector(
+            codebase_name=self.codebase_name,
+            collect_file_metrics=self.collect_file_metrics,
+            collect_worker_metrics=self.collect_worker_metrics,
+        )
+
+        # Generate unique graph ID for this run
+        graph_id = _compute_graph_id(self.codebase_path)
+
+        # Create database
+        kuzu = get_kuzu()
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+
+        # Ensure clean state
+        if graph_path.exists():
+            if graph_path.is_dir():
+                shutil.rmtree(graph_path)
+            else:
+                graph_path.unlink()
+        wal_path = self._storage_dir / f"{graph_id}.kuzu.wal"
+        if wal_path.exists():
+            wal_path.unlink()
+
+        # Create database and connection
+        db = kuzu.Database(str(graph_path))
+        conn = kuzu.Connection(db)
+
+        # Load parsers
+        parsers, queries = load_parsers()
+
+        # Create ingestor and builder
+        ingestor = Ingestor(conn)
+        ingestor.create_schema()
+
+        builder = SimpleGraphBuilder(
+            ingestor=ingestor,
+            repo_path=self.codebase_path,
+            parsers=parsers,
+            queries=queries,
+            metrics_collector=metrics_collector,
+            enable_parallel=self.parallel,
+            progress_callback=None,  # Disable TUI progress in benchmark mode
+        )
+
+        # Run indexing
+        await builder.run()
+
+        # Get metrics
+        metrics = metrics_collector.get_metrics()
+
+        # Close connection
+        del conn
+        del db
+
+        return BenchmarkRun(
+            run_id=run_id,
+            is_warmup=is_warmup,
+            metrics=metrics,
+        )
+
+    async def _cleanup_database(self) -> None:
+        """Delete database files and clear caches between runs."""
+        graph_id = _compute_graph_id(self.codebase_path)
+
+        # Delete database file
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+        if graph_path.exists():
+            if graph_path.is_dir():
+                shutil.rmtree(graph_path)
+            else:
+                graph_path.unlink()
+            logger.debug(f"Deleted database: {graph_path}")
+
+        # Delete WAL file
+        wal_path = self._storage_dir / f"{graph_id}.kuzu.wal"
+        if wal_path.exists():
+            wal_path.unlink()
+            logger.debug(f"Deleted WAL: {wal_path}")
+
+        # Force garbage collection
+        gc.collect()
+
+    async def _register_codebase(self) -> None:
+        """Register the codebase so it appears in `shotgun codebase list`.
+
+        This creates a Project node in the database with metadata that
+        identifies the indexed codebase.
+        """
+        graph_id = _compute_graph_id(self.codebase_path)
+        graph_path = self._storage_dir / f"{graph_id}.kuzu"
+
+        if not graph_path.exists():
+            logger.warning("Cannot register codebase: database not found")
+            return
+
+        kuzu = get_kuzu()
+        db = kuzu.Database(str(graph_path))
+        conn = kuzu.Connection(db)
+
+        try:
+            # Check if Project node already exists
+            result = conn.execute("MATCH (p:Project) RETURN p.graph_id LIMIT 1")
+            if result.has_next():
+                logger.debug("Project node already exists, skipping registration")
+                return
+
+            # Create Project node with metadata
+            current_time = int(time.time())
+            conn.execute(
+                """
+                CREATE (p:Project {
+                    name: $name,
+                    repo_path: $repo_path,
+                    graph_id: $graph_id,
+                    created_at: $created_at,
+                    updated_at: $updated_at,
+                    schema_version: $schema_version,
+                    build_options: $build_options,
+                    indexed_from_cwds: $indexed_from_cwds
+                })
+                """,
+                {
+                    "name": self.codebase_name,
+                    "repo_path": str(self.codebase_path),
+                    "graph_id": graph_id,
+                    "created_at": current_time,
+                    "updated_at": current_time,
+                    "schema_version": "1.0.0",
+                    "build_options": json.dumps({}),
+                    "indexed_from_cwds": json.dumps([str(Path.cwd())]),
+                },
+            )
+            logger.info(
+                f"Registered codebase '{self.codebase_name}' with graph_id: {graph_id}"
+            )
+        finally:
+            del conn
+            del db