iflow-mcp_niclasolofsson-dbt-core-mcp 1.7.0__py3-none-any.whl

dbt_core_mcp/dbt/runner.py

@@ -0,0 +1,67 @@
+"""
+DBT Runner Protocol.
+
+Defines the interface for running dbt commands, supporting both in-process
+and subprocess execution.
+"""
+
+from pathlib import Path
+from typing import Protocol
+
+
+class DbtRunnerResult:
+    """Result from a dbt command execution."""
+
+    def __init__(self, success: bool, exception: Exception | None = None, stdout: str = "", stderr: str = "", elapsed_time: float | None = None):
+        """
+        Initialize a dbt runner result.
+
+        Args:
+            success: Whether the command succeeded
+            exception: Exception if the command failed
+            stdout: Standard output from the command
+            stderr: Standard error from the command
+            elapsed_time: Optional elapsed time in seconds for the command execution
+        """
+        self.success = success
+        self.exception = exception
+        self.stdout = stdout
+        self.stderr = stderr
+        self.elapsed_time = elapsed_time
+
+
+class DbtRunner(Protocol):
+    """Protocol for dbt command execution."""
+
+    def invoke(self, args: list[str]) -> DbtRunnerResult:
+        """
+        Execute a dbt command.
+
+        Args:
+            args: dbt command arguments (e.g., ['parse'], ['run', '--select', 'model'])
+
+        Returns:
+            Result of the command execution
+        """
+        ...
+
+    def get_manifest_path(self) -> Path:
+        """
+        Get the path to the manifest.json file.
+
+        Returns:
+            Path to target/manifest.json
+        """
+        ...
+
+    def invoke_query(self, sql: str) -> DbtRunnerResult:
+        """
+        Execute a SQL query.
+
+        Args:
+            sql: SQL query to execute (include LIMIT in SQL if needed)
+
+        Returns:
+            Result with query output
+        """
+        ...

dbt_core_mcp/dependencies.py

@@ -0,0 +1,50 @@
+"""Dependency injection providers for dbt Core MCP tools.
+
+This module provides dependency functions that can be used with FastMCP's Depends()
+to inject shared state and context into tool functions.
+"""
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .context import DbtCoreServerContext
+
+# Global state reference - set by server during initialization
+_server_state: "DbtCoreServerContext | None" = None
+
+
+def set_server_state(state: "DbtCoreServerContext") -> None:
+    """Set the global server state reference.
+
+    Called by DbtCoreMcpServer during initialization.
+
+    Args:
+        state: The DbtCoreServerContext instance to make available to tools
+    """
+    global _server_state
+    _server_state = state
+
+
+def get_state() -> "DbtCoreServerContext":
+    """Dependency provider for server state.
+
+    Used with FastMCP's Depends() to inject DbtCoreServerContext into tools.
+
+    Returns:
+        The shared DbtCoreServerContext instance
+
+    Raises:
+        RuntimeError: If server state has not been initialized
+
+    Example:
+        @tool
+        async def my_tool(
+            ctx: Context,
+            state: DbtCoreServerContext = Depends(get_state),
+        ) -> dict:
+            # state is automatically injected
+            pass
+    """
+    if _server_state is None:
+        raise RuntimeError("Server state not initialized")
+    return _server_state

dbt_core_mcp/server.py

@@ -0,0 +1,381 @@
+"""
+dbt Core MCP Server Implementation.
+
+This server provides tools for interacting with dbt projects via the Model Context Protocol.
+"""
+
+import asyncio
+import logging
+import os
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Any, AsyncIterator
+from urllib.parse import unquote
+from urllib.request import url2pathname
+
+from fastmcp import FastMCP
+from fastmcp.server.context import Context
+from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
+from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
+
+from .context import DbtCoreServerContext
+from .dbt.bridge_runner import BridgeRunner
+from .dbt.manifest import ManifestLoader
+from .dependencies import set_server_state
+
+# Import tools for static registration
+from .utils.env_detector import detect_python_command
+
+# Re-export DbtCoreServerContext for tools
+__all__ = ["DbtCoreServerContext", "DbtCoreMcpServer", "create_server"]
+
+logger = logging.getLogger(__name__)
+
+
+class DbtCoreMcpServer:
+    """
+    dbt Core MCP Server.
+
+    Provides tools for interacting with dbt projects.
+    """
+
+    def __init__(self, project_dir: str | None = None, timeout: float | None = None) -> None:
+        """Initialize the server.
+
+        Args:
+            project_dir: Optional path to dbt project directory. If not provided,
+                        automatically detects from MCP workspace roots or falls back to cwd.
+            timeout: Optional timeout in seconds for dbt commands (default: None for no timeout).
+        """
+        # FastMCP initialization with recommended arguments
+        from . import __version__
+
+        @asynccontextmanager
+        async def lifespan(app: FastMCP) -> AsyncIterator[None]:
+            """Lifespan context manager for server startup and shutdown."""
+            # Startup
+            logger.info("dbt Core MCP Server starting up...")
+            yield
+            # Shutdown
+            logger.info("dbt Core MCP Server shutting down...")
+            if self.runner:
+                logger.info("Stopping persistent dbt process...")
+                await self.runner.shutdown()
+                logger.info("Persistent dbt process stopped successfully")
+            logger.info("dbt Core MCP Server shutdown complete")
+
+        self.app = FastMCP(
+            version=__version__,
+            name="dbt Core MCP",
+            lifespan=lifespan,
+            instructions="""dbt Core MCP Server for interacting with dbt projects.
+
+            This server provides tools to:
+            - Query dbt project metadata
+            - Run dbt commands
+            - Inspect models, sources, and tests
+            - View compiled SQL
+            - Access dbt documentation
+
+            Usage:
+            - Use the tools to interact with your dbt project
+            - Query model lineage and dependencies
+            - Run and test dbt models
+            
+            ## Subagent Delegation Protocol
+            
+            ALWAYS use #runSubagent for dbt data exploration and analysis. Your context window fills quickly with query results - avoid this by delegating complex workflows to subagents.
+            
+            When to delegate to #runSubagent:
+            - Data exploration across multiple models/sources
+            - Multi-query analysis workflows  
+            - Large result sets (>100 rows expected)
+            - Impact analysis requiring sample queries
+            - Any workflow involving: schema inspection → queries → analysis
+            
+            The subagent handles all the heavy lifting and returns only a concise summary, keeping your main context clean.
+            
+            Single simple queries with known schema: Execute directly using the tools.
+            """,
+        )  # type: ignore[arg-type]
+
+        # Store the explicit project_dir if provided, otherwise will detect from workspace roots
+        _explicit_project_dir = Path(project_dir) if project_dir else None
+        project_dir_resolved: Path | None = None
+        profiles_dir = os.path.expanduser("~/.dbt")
+
+        # Parse experimental features flag from environment
+        experimental_features = os.getenv("EXPERIMENTAL_FEATURES", "false").lower() == "true"
+
+        # Create shared state with all dbt components
+        self.state = DbtCoreServerContext(
+            app=self.app,
+            project_dir=project_dir_resolved,
+            profiles_dir=profiles_dir,
+            timeout=timeout,
+            runner=None,
+            manifest=None,
+            adapter_type=None,
+            force_fresh_runner=False,  # Set to False to reuse runners for performance
+            experimental_features=experimental_features,
+            _init_lock=asyncio.Lock(),
+            _explicit_project_dir=_explicit_project_dir,
+        )
+
+        # Set back-reference for delegation
+        self.state.server = self
+        set_server_state(self.state)
+
+        # Keep references for backward compatibility with existing helper methods
+        self.project_dir = project_dir_resolved
+        self.profiles_dir = profiles_dir
+        self.timeout = timeout
+        self._explicit_project_dir = _explicit_project_dir
+        self.runner = None
+        self.manifest = None
+        self.adapter_type = None
+        self.force_fresh_runner = False
+        self._init_lock = asyncio.Lock()
+
+        # Add built-in FastMCP middleware (2.11.0)
+        self.app.add_middleware(ErrorHandlingMiddleware())  # Handle errors first
+        self.app.add_middleware(RateLimitingMiddleware(max_requests_per_second=50))
+        # TimingMiddleware and LoggingMiddleware removed - they use structlog with column alignment
+        # which causes formatting issues in VS Code's output panel
+
+        # Register tools dynamically
+        self._register_tools()
+        # Register resources
+        self._register_resources()
+
+        logger.info("dbt Core MCP Server initialized")
+        logger.info(f"Profiles directory: {self.profiles_dir}")
+
+    # Public wrappers for DbtCoreServerContext to avoid private-member access warnings
+    async def ensure_initialized_with_context(self, ctx: Context | None, force_parse: bool = False) -> None:
+        await self._ensure_initialized_with_context(ctx, force_parse=force_parse)
+
+    async def get_runner(self) -> BridgeRunner:
+        return await self._get_runner()
+
+    def _register_tools(self) -> None:
+        """Dynamically register all dbt Core MCP tools."""
+        from .tools import discover_tools_in_package, get_tool_metadata
+
+        tool_functions = discover_tools_in_package("dbt_core_mcp.tools")
+        for tool_func in tool_functions:
+            metadata = get_tool_metadata(tool_func, default=None)
+            if metadata:
+                allowed_keys = {
+                    "name",
+                    "description",
+                    "tags",
+                    "enabled",
+                    "icons",
+                    "annotations",
+                    "meta",
+                }
+                tool_kwargs = {key: value for key, value in metadata.items() if key in allowed_keys}
+                self.app.tool(**tool_kwargs)(tool_func)
+                logger.info("Registered tool metadata for %s: %s", tool_func.__name__, metadata)
+            else:
+                self.app.tool()(tool_func)
+
+    def _register_resources(self) -> None:
+        demo_html_path = Path(__file__).resolve().parent / "tools" / "demo" / "hello.html"
+
+        # Register MCP UI resource for demo_ui tool.
+        # It is CRITICAL for vscode that it is resource:// and not ui://
+        @self.app.resource(
+            uri="resource://demo/hello",
+            name="Demo UI",
+            mime_type="text/html;profile=mcp-app",
+        )
+        def demo_ui() -> str:
+            logger.info("Serving MCP UI resource: resource://demo/hello")
+            return demo_html_path.read_text(encoding="utf-8")
+
+        # Also register with ui:// for compatibility (VS Code may request either)
+        @self.app.resource(
+            uri="ui://demo/hello",
+            name="Demo UI (Legacy)",
+            mime_type="text/html;profile=mcp-app",
+        )
+        def demo_ui_legacy() -> str:
+            logger.info("Serving MCP UI resource: ui://demo/hello")
+            return demo_html_path.read_text(encoding="utf-8")
+
+    # (Removed) Legacy toolImpl_* wrappers: tools are now auto-discovered via FileSystemProvider
+
+    def _detect_project_dir(self) -> Path:
+        """Detect the dbt project directory.
+
+        Resolution order:
+        1. Use explicit project_dir if provided during initialization
+        2. Fall back to current working directory
+
+        Note: Workspace roots detection happens in _detect_workspace_roots()
+        which is called asynchronously from tool contexts.
+
+        Returns:
+            Path to the dbt project directory
+        """
+        # Use explicit project_dir if provided
+        if self._explicit_project_dir:
+            logger.debug(f"Using explicit project directory: {self._explicit_project_dir}")
+            return self._explicit_project_dir
+
+        # Fall back to current working directory
+        cwd = Path.cwd()
+        logger.info(f"Using current working directory: {cwd}")
+        return cwd
+
+    async def _detect_workspace_roots(self, ctx: Any) -> Path | None:
+        """Attempt to detect workspace roots from MCP context.
+
+        Args:
+            ctx: FastMCP Context object
+
+        Returns:
+            Path to first workspace root, or None if unavailable
+        """
+        try:
+            if isinstance(ctx, Context):
+                roots = await ctx.list_roots()
+                if roots:
+                    # Convert file:// URL to platform-appropriate path
+                    # First unquote to decode %XX sequences, then url2pathname for platform conversion
+                    uri_path = roots[0].uri.path if hasattr(roots[0].uri, "path") else str(roots[0].uri)
+                    if uri_path:
+                        workspace_root = Path(url2pathname(unquote(uri_path)))
+                        logger.info(f"Detected workspace root from MCP client: {workspace_root}")
+                        return workspace_root
+        except Exception as e:
+            logger.debug(f"Could not access workspace roots: {e}")
+
+        return None
+
+    async def _get_runner(self) -> BridgeRunner:
+        """Get BridgeRunner instance with explicit control over creation.
+
+        Uses self.force_fresh_runner to determine behavior:
+        - If True, always create a fresh BridgeRunner instance
+        - If False, reuse existing runner if available
+
+        Returns:
+            BridgeRunner instance
+        """
+        if self.force_fresh_runner or not self.runner:
+            if not self.project_dir:
+                raise RuntimeError("Project directory not set")
+
+            # Detect Python command for user's environment
+            python_cmd = detect_python_command(self.project_dir)
+            logger.info(f"Creating {'fresh' if self.force_fresh_runner else 'new'} BridgeRunner with command: {python_cmd}")
+
+            # Create bridge runner with persistent process for better performance
+            self.runner = BridgeRunner(self.project_dir, python_cmd, timeout=self.timeout, use_persistent_process=True)
+
+        return self.runner
+
+    async def _initialize_dbt_components(self, needs_parse: bool = True, force_parse: bool = False) -> None:
+        """Initialize dbt runner and manifest loader.
+
+        Args:
+            needs_parse: Whether to run dbt parse. If False, assumes manifest already exists and is fresh.
+            force_parse: If True, force parsing even if manifest exists (for tools needing fresh data).
+        """
+
+        if not self.project_dir:
+            raise RuntimeError("Project directory not set")
+
+        # Get runner (fresh or reused based on self.force_fresh_runner)
+        runner = await self._get_runner()
+
+        # Run parse if needed and not skipped via force_parse=False
+        if needs_parse or force_parse:
+            logger.info("Running dbt parse...")
+            parse_args = ["parse"]  # Use partial parse for efficiency
+            result = await runner.invoke(parse_args)
+            if not result.success:
+                error_msg = str(result.exception) if result.exception else "Unknown error"
+                raise RuntimeError(f"Failed to parse dbt project: {error_msg}")
+
+        # Initialize or reload manifest loader
+        manifest_path = runner.get_manifest_path()
+        if not self.manifest:
+            self.manifest = ManifestLoader(manifest_path)
+        await self.manifest.load()
+
+        # Keep shared state in sync with server-owned components
+        self._sync_shared_state()
+
+        logger.info("dbt components initialized successfully")
+
+    def _sync_shared_state(self) -> None:
+        """Keep DbtCoreServerContext references aligned with server fields."""
+        self.state.project_dir = self.project_dir
+        self.state.runner = self.runner
+        self.state.manifest = self.manifest
+        self.state.adapter_type = self.adapter_type
+        self.state.force_fresh_runner = self.force_fresh_runner
+
+    async def _ensure_initialized_with_context(self, ctx: Any, force_parse: bool = False) -> None:
+        """Ensure dbt components are initialized, with optional workspace root detection.
+
+        Uses async lock to prevent concurrent initialization races when multiple tools
+        are called simultaneously.
+
+        Args:
+            ctx: FastMCP Context for accessing workspace roots
+            force_parse: If True, force parsing even if manifest exists (for tools needing fresh data)
+        """
+        async with self._init_lock:
+            # Always check for workspace changes, even if previously initialized
+            detected_workspace: Path | None = None
+
+            if not self._explicit_project_dir:
+                detected_workspace = await self._detect_workspace_roots(ctx)
+
+            # If workspace changed, reinitialize everything
+            if detected_workspace and detected_workspace != self.project_dir:
+                logger.info(f"Workspace changed from {self.project_dir} to {detected_workspace}, reinitializing...")
+                self.project_dir = detected_workspace
+                self.runner = None
+                self.manifest = None
+
+            # Ensure project directory is set (first time or after workspace change)
+            if not self.project_dir:
+                if detected_workspace:
+                    self.project_dir = detected_workspace
+                else:
+                    self.project_dir = self._detect_project_dir()
+                    logger.info(f"dbt project directory: {self.project_dir}")
+
+            if not self.project_dir:
+                raise RuntimeError("dbt project directory not set. The MCP server requires a workspace with a dbt_project.yml file.")
+
+            await self._initialize_dbt_components(needs_parse=not self.state.manifest_exists(), force_parse=force_parse)
+
+    def run(self, stateless: bool = False) -> None:
+        """Run the MCP server.
+
+        Args:
+            stateless: Enable stateless mode for seamless restarts (auto-enabled with --reload).
+        """
+        self.app.run(show_banner=False, stateless=stateless)
+
+
+def create_server(project_dir: str | None = None, timeout: float | None = None) -> DbtCoreMcpServer:
+    """Create a new dbt Core MCP server instance.
+
+    Args:
+        project_dir: Optional path to dbt project directory.
+                     If not provided, automatically detects from MCP workspace roots
+                     or falls back to current working directory.
+        timeout: Optional timeout in seconds for dbt commands (default: None for no timeout).
+
+    Returns:
+        DbtCoreMcpServer instance
+    """
+    return DbtCoreMcpServer(project_dir=project_dir, timeout=timeout)

dbt_core_mcp/tools/__init__.py

@@ -0,0 +1,77 @@
+"""dbt-core-mcp tools package.
+
+Each tool module implements a single MCP tool with:
+- setup(app, state): Registers the tool with FastMCP
+- _implementation(...): Pure logic function (testable)
+"""
+
+import importlib
+import inspect
+import pkgutil
+from pathlib import Path
+from typing import Any, Callable, TypeVar
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+def dbtTool(**metadata: Any) -> Callable[[F], F]:
+    """Decorator to mark dbt MCP tool functions.
+
+    Stores optional metadata on the function for later use during registration.
+    Actual tool registration happens in server.py.
+    """
+
+    def decorator(func: F) -> F:
+        func.metadata = metadata  # type: ignore[attr-defined]
+        return func
+
+    return decorator
+
+
+def get_tool_metadata(func: Callable[..., Any], default: Any = None) -> Any:
+    return getattr(func, "metadata", default)
+
+
+def has_tool_metadata(func: Callable[..., Any]) -> bool:
+    return hasattr(func, "metadata")
+
+
+def set_tool_metadata(func: Callable[..., Any], metadata: Any) -> None:
+    func.metadata = metadata  # type: ignore[attr-defined]
+
+
+def discover_tools(module: Any) -> list[Callable[..., Any]]:
+    tools: list[Callable[..., Any]] = []
+    for _, obj in inspect.getmembers(module):
+        if (inspect.isfunction(obj) or inspect.iscoroutinefunction(obj)) and has_tool_metadata(obj):
+            tools.append(obj)
+    return tools
+
+
+def discover_tools_in_path(path: str | Path, package_prefix: str) -> list[Callable[..., Any]]:
+    tools: list[Callable[..., Any]] = []
+    base_path = Path(path)
+    for module_info in pkgutil.walk_packages([str(base_path)], prefix=f"{package_prefix}."):
+        module = importlib.import_module(module_info.name)
+        tools.extend(discover_tools(module))
+    return tools
+
+
+def discover_tools_in_package(package_name: str) -> list[Callable[..., Any]]:
+    package = importlib.import_module(package_name)
+    package_file = getattr(package, "__file__", None)
+    if not package_file:
+        raise RuntimeError(f"Package {package_name} has no __file__ path to scan")
+    package_path = Path(package_file).parent
+    return discover_tools_in_path(package_path, package_name)
+
+
+__all__ = [
+    "dbtTool",
+    "get_tool_metadata",
+    "has_tool_metadata",
+    "set_tool_metadata",
+    "discover_tools",
+    "discover_tools_in_path",
+    "discover_tools_in_package",
+]

dbt_core_mcp/tools/analyze_impact.py

@@ -0,0 +1,78 @@
+"""Analyze the impact of changing any dbt resource.
+
+This module implements the analyze_impact tool for dbt Core MCP.
+"""
+
+import logging
+from typing import Any
+
+from fastmcp.dependencies import Depends  # type: ignore[reportAttributeAccessIssue]
+from fastmcp.server.context import Context
+
+from ..context import DbtCoreServerContext
+from ..dependencies import get_state
+from . import dbtTool
+
+logger = logging.getLogger(__name__)
+
+
+async def _implementation(ctx: Context | None, name: str, resource_type: str | None, state: DbtCoreServerContext, force_parse: bool = True) -> dict[str, Any]:
+    """Implementation function for analyze_impact tool.
+
+    Separated for testing purposes - tests call this directly with explicit state.
+    The @tool() decorated analyze_impact() function calls this with injected dependencies.
+    """
+    # Initialize state if needed (metadata tool uses force_parse=True)
+    await state.ensure_initialized(ctx, force_parse)
+
+    # Delegate to manifest helper for downstream impact calculation
+    try:
+        return state.manifest.analyze_impact(name, resource_type)  # type: ignore
+    except ValueError as e:
+        raise ValueError(f"Impact analysis error: {e}")
+
+
+@dbtTool()
+async def analyze_impact(
+    ctx: Context,
+    name: str,
+    resource_type: str | None = None,
+    state: DbtCoreServerContext = Depends(get_state),
+) -> dict[str, Any]:
+    """Analyze the impact of changing any dbt resource with auto-detection.
+
+    This unified tool works across all resource types (models, sources, seeds, snapshots, etc.)
+    showing all downstream dependencies that would be affected by changes. Provides actionable
+    recommendations for running affected resources.
+
+    Args:
+        name: Resource name. For sources, use "source_name.table_name" or just "table_name"
+            Examples: "stg_customers", "jaffle_shop.orders", "raw_customers"
+        resource_type: Optional filter to narrow search:
+            - "model": Data transformation models
+            - "source": External data sources
+            - "seed": CSV reference data files
+            - "snapshot": SCD Type 2 historical tables
+            - "test": Data quality tests
+            - "analysis": Ad-hoc analysis queries
+            - None: Auto-detect (searches all types)
+
+    Returns:
+        Impact analysis with:
+        - List of affected models by distance
+        - Count of affected tests and other resources
+        - Total impact statistics
+        - Resources grouped by distance from changed resource
+        - Recommended dbt command to run affected resources
+        - Human-readable impact assessment message
+        If multiple matches found, returns all matches for LLM to process.
+
+    Raises:
+        ValueError: If resource not found
+
+    Examples:
+        analyze_impact("stg_customers") -> auto-detect and show impact
+        analyze_impact("jaffle_shop.orders", "source") -> impact of source change
+        analyze_impact("raw_customers", "seed") -> impact of seed data change
+    """
+    return await _implementation(ctx, name, resource_type, state)