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

dbt_core_mcp/tools/snapshot_models.py

@@ -0,0 +1,120 @@
+"""Snapshot models (capture historical changes).
+
+This module implements the snapshot_models 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,
+    select: str | None,
+    exclude: str | None,
+    state: DbtCoreServerContext,
+) -> dict[str, Any]:
+    """Implementation function for snapshot_models tool.
+
+    Separated for testing purposes - tests call this directly with explicit state.
+    The @tool() decorated snapshot_models() function calls this with injected dependencies.
+    """
+    # Ensure dbt components are initialized
+    await state.ensure_initialized(ctx, force_parse=False)
+
+    # Construct dbt CLI args for snapshot
+    args = ["snapshot"]
+
+    if select:
+        args.extend(["-s", select])
+
+    if exclude:
+        args.extend(["--exclude", exclude])
+
+    logger.info(f"Running DBT snapshot with args: {args}")
+
+    # Clear stale run_results so we parse only fresh output
+    state.clear_stale_run_results()
+
+    runner = await state.get_runner()
+    result = await runner.invoke(args)  # type: ignore
+
+    run_results = state.validate_and_parse_results(result, "snapshot")
+
+    if ctx:
+        if run_results.get("results"):
+            results_list = run_results["results"]
+            total = len(results_list)
+            passed_count = sum(1 for r in results_list if r.get("status") == "success")
+            failed_count = sum(1 for r in results_list if r.get("status") in ("error", "fail"))
+
+            # Summarize final progress back to caller
+            parts = []
+            if passed_count > 0:
+                parts.append(f"✅ {passed_count} passed" if failed_count > 0 else "✅ All passed")
+            if failed_count > 0:
+                parts.append(f"❌ {failed_count} failed")
+
+            summary = f"Snapshot: {total}/{total} snapshots completed ({', '.join(parts)})"
+            await ctx.report_progress(progress=total, total=total, message=summary)
+        else:
+            await ctx.report_progress(progress=0, total=0, message="0 snapshots matched selector")
+
+    if not run_results.get("results"):
+        raise RuntimeError(f"No snapshots matched selector: {select or 'all'}")
+
+    return {
+        "status": "success",
+        "command": " ".join(args),
+        "results": run_results.get("results", []),
+        "elapsed_time": run_results.get("elapsed_time"),
+    }
+
+
+@dbtTool()
+async def snapshot_models(
+    ctx: Context,
+    select: str | None = None,
+    exclude: str | None = None,
+    state: DbtCoreServerContext = Depends(get_state),
+) -> dict[str, Any]:
+    """Snapshot models (capture historical changes - SCD Type 2).
+
+    Snapshots capture historical changes in data, enabling you to track slowly changing
+    dimensions over time. This is particularly useful for maintaining accurate historical
+    records in data warehouses.
+
+    **When to use**: To track changes in slowly changing dimensions (SCD Type 2).
+    For example, tracking customer address changes over time while preserving history.
+
+    **How it works**: dbt compares current source data with existing snapshot table,
+    identifies changes, and inserts new rows with validity timestamps (dbt_valid_from,
+    dbt_valid_to, dbt_updated_at). Original rows are closed by setting dbt_valid_to.
+
+    Args:
+        select: dbt selector syntax (e.g., "snapshot_name", "tag:daily")
+        exclude: Exclude specific snapshots
+        state: Shared state object injected by FastMCP
+
+    Returns:
+        Snapshot results with status and timing info
+
+    Examples:
+        # Run all snapshots
+        snapshot_models()
+
+        # Run specific snapshot
+        snapshot_models(select="customers_snapshot")
+
+        # Run tagged snapshots
+        snapshot_models(select="tag:daily")
+    """
+    return await _implementation(ctx, select, exclude, state)

dbt_core_mcp/tools/test_models.py

@@ -0,0 +1,238 @@
+"""Run dbt tests on models and sources.
+
+This module implements the test_models tool for dbt Core MCP.
+"""
+
+import json
+import logging
+from typing import Any
+
+from fastmcp.dependencies import Depends  # type: ignore[reportAttributeAccessIssue]
+from fastmcp.exceptions import McpError  # type: ignore[attr-defined]
+from fastmcp.server.context import Context
+from mcp.types import ErrorData
+
+from ..context import DbtCoreServerContext
+from ..cte_generator import cleanup_cte_tests, generate_cte_tests
+from ..dependencies import get_state
+from . import dbtTool
+
+logger = logging.getLogger(__name__)
+
+
+async def _implementation(
+    ctx: Context | None,
+    select: str | None,
+    exclude: str | None,
+    select_state_modified: bool,
+    select_state_modified_plus_downstream: bool,
+    fail_fast: bool,
+    keep_cte_tests: bool,
+    state: DbtCoreServerContext,
+) -> dict[str, Any]:
+    """Implementation function for test_models tool.
+
+    Separated for testing purposes - tests call this directly with explicit state.
+    The @tool() decorated test_models() function calls this with injected dependencies.
+    """
+    # Ensure dbt components are initialized
+    await state.ensure_initialized(ctx, force_parse=False)
+
+    # Track if CTE tests were generated for cleanup
+    cte_tests_generated = False
+
+    try:
+        # Generate CTE tests if experimental features enabled
+        if state.experimental_features and state.project_dir:
+            logger.info("Experimental features enabled - generating CTE tests")
+            try:
+                cte_count = generate_cte_tests(state.project_dir)
+                if cte_count > 0:
+                    logger.info(f"Generated {cte_count} CTE tests")
+                    cte_tests_generated = True
+            except Exception as e:
+                logger.warning(f"CTE test generation failed: {e}")
+                # Don't fail the entire test run if CTE generation fails
+
+        # Build state-based selector if requested (avoids redundant parsing when possible)
+        selector = await state.prepare_state_based_selection(select_state_modified, select_state_modified_plus_downstream, select)
+
+        # If user asked for state:modified but no baseline exists, fail fast with clear guidance
+        if select_state_modified and not selector:
+            raise RuntimeError("No previous state found - cannot determine modifications. Run 'dbt run' or 'dbt build' first to create baseline state.")
+
+        args = ["test"]
+
+        # Add selector if we have one (state-based or manual)
+        if selector:
+            target_path = state.get_project_paths().get("target-path", "target")
+            args.extend(["-s", selector, "--state", f"{target_path}/state_last_run"])
+        elif select:
+            args.extend(["-s", select])
+
+        if exclude:
+            args.extend(["--exclude", exclude])
+
+        if fail_fast:
+            args.append("--fail-fast")
+
+        # Execute with progress reporting
+        logger.info(f"Running dbt tests with args: {args}")
+
+        # Define progress callback if context available
+        async def progress_callback(current: int, total: int, message: str) -> None:
+            if ctx:
+                await ctx.report_progress(progress=current, total=total, message=message)
+
+        # Delete stale run_results.json to ensure we only read fresh results
+        state.clear_stale_run_results()
+
+        runner = await state.get_runner()
+        result = await runner.invoke(args, progress_callback=progress_callback if ctx else None)  # type: ignore
+
+        # Parse run_results.json to discriminate system errors from business outcomes
+        run_results = state.validate_and_parse_results(result, "test")
+
+        # Business outcome - dbt executed tests (some may have failed)
+        # Continue with existing run_results parsing
+        results_list = run_results.get("results", [])
+
+        # Remove heavy fields from results to reduce token usage
+        for result in results_list:
+            result.pop("compiled_code", None)
+            result.pop("raw_code", None)
+
+        # Send final progress update with test summary
+        await state.report_final_progress(ctx, results_list, "Test", "tests")
+
+        # Empty results means selector matched nothing - this is an error
+        if not results_list:
+            raise RuntimeError(f"No tests matched selector: {select or selector or 'all'}")
+
+        # Check if any tests failed - if so, return RPC error with structured data
+        failed_tests = [r for r in results_list if r.get("status") == "fail"]
+        if failed_tests:
+            # Check if any failed tests are unit tests (they have diff output)
+            has_unit_test_failures = any("unit_test" in r.get("unique_id", "") for r in failed_tests)
+
+            # Add diff format legend for unit test failures
+            diff_legend = ""
+            if has_unit_test_failures:
+                diff_legend = (
+                    "\n\nUnit test diff format (daff tabular diff):\n"
+                    "  @@    Header row with column names\n"
+                    "  +++   Row present in actual output but missing from expected\n"
+                    "  ---   Row present in expected but missing from actual output\n"
+                    "  →     Row with at least one modified cell (→, -->, etc.)\n"
+                    "  ...   Rows omitted for brevity\n"
+                    "  old_value→new_value   Shows the change in a specific cell\n"
+                    "\nFull specification: https://paulfitz.github.io/daff-doc/spec.html"
+                )
+
+            error_data = {
+                "error": "test_failure",
+                "message": f"{len(failed_tests)} test(s) failed{diff_legend}",
+                "command": " ".join(args),
+                "results": results_list,  # Include ALL results (both passing and failing)
+                "elapsed_time": run_results.get("elapsed_time"),
+                "summary": {
+                    "total": len(results_list),
+                    "passed": len([r for r in results_list if r.get("status") == "pass"]),
+                    "failed": len(failed_tests),
+                },
+            }
+            # Use McpError with custom code for business failures (not system errors)
+            # Code -32000 to -32099 are reserved for implementation-defined server errors
+            raise McpError(
+                ErrorData(
+                    code=-32000,  # Server error (business failure, not system error)
+                    message=json.dumps(error_data, indent=2),
+                )
+            )
+
+        # All tests passed - return success
+        return {
+            "status": "success",
+            "command": " ".join(args),
+            "results": results_list,
+            "elapsed_time": run_results.get("elapsed_time"),
+        }
+    finally:
+        # Clean up CTE tests if they were generated and not debugging
+        if cte_tests_generated and not keep_cte_tests and state.project_dir:
+            try:
+                cleanup_cte_tests(state.project_dir)
+            except Exception as e:
+                logger.warning(f"Failed to cleanup CTE tests: {e}")
+
+
+@dbtTool()
+async def test_models(
+    ctx: Context,
+    select: str | None = None,
+    exclude: str | None = None,
+    select_state_modified: bool = False,
+    select_state_modified_plus_downstream: bool = False,
+    fail_fast: bool = False,
+    keep_cte_tests: bool = False,
+    state: DbtCoreServerContext = Depends(get_state),
+) -> dict[str, Any]:
+    """Run dbt tests on models and sources.
+
+    **When to use**: After running models to validate data quality. Tests check constraints
+    like uniqueness, not-null, relationships, and custom data quality rules.
+
+    **Important**: Ensure seeds and models are built before running tests that depend on them.
+
+    State-based selection modes (uses dbt state:modified selector):
+    - select_state_modified: Test only models modified since last successful run (state:modified)
+    - select_state_modified_plus_downstream: Test modified + downstream dependencies (state:modified+)
+      Note: Requires select_state_modified=True
+
+    Manual selection (alternative to state-based):
+    - select: dbt selector syntax (e.g., "customers", "tag:mart", "test_type:generic")
+    - exclude: Exclude specific tests
+
+    Args:
+        select: Manual selector for tests/models to test
+        exclude: Exclude selector
+        select_state_modified: Use state:modified selector (changed models only)
+        select_state_modified_plus_downstream: Extend to state:modified+ (changed + downstream)
+        fail_fast: Stop execution on first failure
+        keep_cte_tests: Keep generated CTE test files for debugging (default: False)
+        state: Shared state object injected by FastMCP
+
+    Returns:
+        Test results with status and failures
+
+    See also:
+        - run_models(): Execute models before testing them
+        - build_models(): Run models + tests together automatically
+        - load_seeds(): Load seeds if tests reference seed data
+
+    Examples:
+        # After building a model, test it
+        run_models(select="customers")
+        test_models(select="customers")
+
+        # Test only generic tests (not singular)
+        test_models(select="test_type:generic")
+
+        # Test everything that changed
+        test_models(select_state_modified=True)
+
+        # Stop on first failure for quick feedback
+        test_models(fail_fast=True)
+
+        # Keep CTE test files for debugging
+        test_models(keep_cte_tests=True)
+
+    Note: Unit test failures show diffs in the "daff" tabular format:
+        @@ = column headers
+        +++ = row in actual, not in expected (extra row)
+        --- = row in expected, not in actual (missing row)
+        → = row with modified cell(s), shown as old_value→new_value
+        ... = omitted matching rows
+        Full format spec: https://paulfitz.github.io/daff-doc/spec.html
+    """
+    return await _implementation(ctx, select, exclude, select_state_modified, select_state_modified_plus_downstream, fail_fast, keep_cte_tests, state)

dbt_core_mcp/utils/__init__.py

@@ -0,0 +1 @@
+"""Utility modules for dbt-core-mcp."""

dbt_core_mcp/utils/env_detector.py

@@ -0,0 +1,186 @@
+"""
+Environment detection for dbt projects.
+
+Detects the Python environment setup and returns the appropriate command
+and environment variables to run Python in that environment.
+"""
+
+import logging
+import os
+import sys
+from pathlib import Path
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+
+def detect_python_command(project_dir: Path) -> list[str]:
+    """
+    Detect how to run Python in the project's environment.
+
+    Detects common Python environment setups and returns the command prefix
+    needed to run Python in that environment.
+
+    Args:
+        project_dir: Path to the dbt project directory
+
+    Returns:
+        Command prefix to run Python (e.g., ['uv', 'run', 'python'])
+
+    Examples:
+        >>> detect_python_command(Path("/project/with/uv"))
+        ['uv', 'run', 'python']
+        >>> detect_python_command(Path("/project/with/poetry"))
+        ['poetry', 'run', 'python']
+    """
+    # Convert to absolute path
+    project_dir = project_dir.resolve()
+
+    # Check for standard venv (.venv or venv) - prefer this for uv projects to avoid VIRTUAL_ENV conflicts
+    venv_path = _find_venv(project_dir)
+    if venv_path:
+        logger.info(f"Detected venv at {venv_path}")
+        python_exe = _get_venv_python(venv_path)
+        return [str(python_exe)]
+
+    # Check for uv (uv.lock) - only if no venv found
+    if (project_dir / "uv.lock").exists():
+        logger.info(f"Detected uv environment in {project_dir}")
+        return ["uv", "run", "--directory", str(project_dir), "python"]
+
+    # Check for poetry (poetry.lock)
+    if (project_dir / "poetry.lock").exists():
+        logger.info(f"Detected poetry environment in {project_dir}")
+        return ["poetry", "run", "--directory", str(project_dir), "python"]
+
+    # Check for pipenv (Pipfile.lock)
+    if (project_dir / "Pipfile.lock").exists():
+        logger.info(f"Detected pipenv environment in {project_dir}")
+        # pipenv doesn't have --directory, need to cd first
+        return ["pipenv", "run", "python"]
+
+    # Check for conda environment
+    conda_env = os.environ.get("CONDA_DEFAULT_ENV")
+    if conda_env:
+        logger.info(f"Detected conda environment: {conda_env}")
+        return ["conda", "run", "-n", conda_env, "python"]
+
+    # Fall back to system Python
+    logger.warning(f"No virtual environment detected in {project_dir}, using system Python")
+    return [sys.executable]
+
+
+def get_env_vars(python_command: list[str]) -> dict[str, str] | None:
+    """
+    Get environment variables needed for the given Python command.
+
+    This centralizes environment-specific configuration that's needed
+    to properly run commands in different virtual environment managers.
+
+    Args:
+        python_command: The Python command prefix (e.g., ['pipenv', 'run', 'python'])
+
+    Returns:
+        Dictionary of environment variables to set, or None if no special env needed
+
+    Examples:
+        >>> get_env_vars(['pipenv', 'run', 'python'])
+        {'PIPENV_IGNORE_VIRTUALENVS': '1', 'PIPENV_VERBOSITY': '-1'}
+        >>> get_env_vars(['python'])
+        None
+    """
+    if not python_command:
+        return None
+
+    env_tool = python_command[0]
+
+    if env_tool == "pipenv":
+        # Pipenv needs to ignore outer virtualenvs when running inside another env (e.g., uv run)
+        # This prevents pipenv from using the wrong environment
+        return {
+            "PIPENV_IGNORE_VIRTUALENVS": "1",
+            "PIPENV_VERBOSITY": "-1",
+        }
+
+    # Add more env-specific settings here as needed
+    # elif env_tool == "poetry":
+    #     return {"POETRY_VIRTUALENVS_IN_PROJECT": "true"}
+
+    return None
+
+
+def _find_venv(project_dir: Path) -> Optional[Path]:
+    """Find a virtual environment directory."""
+    for venv_name in [".venv", "venv", "env"]:
+        venv_path = project_dir / venv_name
+        if venv_path.is_dir() and _is_venv(venv_path):
+            return venv_path
+    return None
+
+
+def _is_venv(path: Path) -> bool:
+    """Check if a directory is a valid virtual environment."""
+    # Check for pyvenv.cfg (created by venv module)
+    if (path / "pyvenv.cfg").exists():
+        return True
+
+    # Check for Scripts/python.exe (Windows) or bin/python (Unix)
+    if sys.platform == "win32":
+        return (path / "Scripts" / "python.exe").exists()
+    else:
+        return (path / "bin" / "python").exists()
+
+
+def _get_venv_python(venv_path: Path) -> Path:
+    """Get the Python executable path from a venv."""
+    if sys.platform == "win32":
+        return venv_path / "Scripts" / "python.exe"
+    else:
+        return venv_path / "bin" / "python"
+
+
+def detect_dbt_adapter(project_dir: Path) -> str:
+    """
+    Detect the dbt adapter type from profiles.yml.
+
+    Args:
+        project_dir: Path to the dbt project directory
+
+    Returns:
+        Adapter type (e.g., 'duckdb', 'postgres', 'snowflake')
+
+    Raises:
+        FileNotFoundError: If dbt_project.yml or profiles.yml not found
+        KeyError: If required keys not found in YAML files
+    """
+    import yaml
+
+    # Read dbt_project.yml to get profile name
+    project_yml_path = project_dir / "dbt_project.yml"
+    if not project_yml_path.exists():
+        raise FileNotFoundError(f"dbt_project.yml not found in {project_dir}")
+
+    project_yml = yaml.safe_load(project_yml_path.read_text())
+    profile_name = project_yml["profile"]
+
+    # Find profiles.yml (project dir or ~/.dbt/)
+    profiles_path = project_dir / "profiles.yml"
+    if not profiles_path.exists():
+        profiles_path = Path.home() / ".dbt" / "profiles.yml"
+
+    if not profiles_path.exists():
+        raise FileNotFoundError(f"profiles.yml not found in {project_dir} or ~/.dbt/")
+
+    # Read profiles.yml and get adapter type
+    profiles = yaml.safe_load(profiles_path.read_text())
+    profile = profiles[profile_name]
+
+    # Get target (default or first output)
+    target_name = profile.get("target")
+    if target_name is None:
+        target_name = list(profile["outputs"].keys())[0]
+
+    adapter_type = profile["outputs"][target_name]["type"]
+
+    logger.info(f"Detected dbt adapter: {adapter_type}")
+    return str(adapter_type)

dbt_core_mcp/utils/process_check.py

@@ -0,0 +1,130 @@
+"""
+Process checking utilities for detecting running dbt processes.
+
+This module provides utilities to check if dbt is currently running
+in the same project directory, helping prevent concurrent execution issues.
+"""
+
+import logging
+from pathlib import Path
+
+import psutil
+
+logger = logging.getLogger(__name__)
+
+
+def is_dbt_running(project_dir: Path, exclude_pid: int | None = None) -> bool:
+    """
+    Check if dbt is currently running in the specified project directory.
+
+    Args:
+        project_dir: Path to the dbt project directory to check
+        exclude_pid: Optional PID to exclude from detection (e.g., our own daemon process)
+
+    Returns:
+        True if a dbt process is detected running in the project directory
+    """
+    project_dir = project_dir.resolve()  # Normalize path
+    logger.debug(f"Checking for dbt processes in: {project_dir}")
+
+    try:
+        for proc in psutil.process_iter(["pid", "name", "cmdline", "cwd"]):
+            try:
+                # Skip if this is the excluded PID (our own daemon)
+                if exclude_pid and proc.info.get("pid") == exclude_pid:
+                    continue
+
+                # Check if this is a dbt-related process
+                cmdline = proc.info.get("cmdline") or []
+                if not cmdline:
+                    continue
+
+                # Look for 'dbt' command in the command line
+                # We want actual 'dbt' CLI commands, not just processes that import dbt
+                cmdline_str = " ".join(cmdline).lower()
+
+                # Skip if this is our own MCP server or Python imports
+                if "dbt-core-mcp" in cmdline_str or "dbt_core_mcp" in cmdline_str:
+                    continue
+
+                # Skip MCP persistent dbt processes (identifiable by their loop script markers)
+                # These processes run idle waiting for stdin and don't interfere with external dbt commands
+                if '"type": "ready"' in cmdline_str or 'type": "ready' in cmdline_str:
+                    logger.debug(f"Skipping MCP persistent process (PID {proc.info['pid']})")
+                    continue
+
+                # Look for actual dbt CLI usage: 'dbt run', 'dbt parse', 'python -m dbt.cli.main', etc.
+                is_dbt_command = False
+                for arg in cmdline:
+                    arg_lower = arg.lower()
+                    # Check for 'dbt' as a standalone command or module
+                    if arg_lower == "dbt" or arg_lower.endswith("dbt.exe") or arg_lower.endswith("dbt"):
+                        is_dbt_command = True
+                        break
+                    # Check for 'python -m dbt.cli.main'
+                    if "dbt.cli.main" in arg_lower or "dbt/cli/main" in arg_lower:
+                        is_dbt_command = True
+                        break
+
+                if not is_dbt_command:
+                    continue
+
+                # Check if it's running in the same project directory
+                # Compare working directory
+                proc_cwd = proc.info.get("cwd")
+                if proc_cwd:
+                    proc_path = Path(proc_cwd).resolve()
+                    # Only match if:
+                    # 1. Exact match - same directory
+                    # 2. Process is running in a subdirectory of our project
+                    if proc_path == project_dir or proc_path.is_relative_to(project_dir):
+                        logger.info(f"Found running dbt process (PID {proc.info['pid']}): {cmdline}")
+                        return True
+
+                # Also check if project directory is mentioned in command line
+                project_str = str(project_dir)
+                if project_str in cmdline_str or str(project_dir).replace("\\", "/") in cmdline_str:
+                    logger.info(f"Found dbt process with project path (PID {proc.info['pid']}): {cmdline}")
+                    return True
+
+            except (psutil.NoSuchProcess, psutil.AccessDenied, psutil.ZombieProcess):
+                # Process disappeared or we don't have permission - skip it
+                continue
+
+        logger.debug("No running dbt processes detected")
+        return False
+
+    except Exception as e:
+        # If we can't check, assume it's safe to proceed
+        logger.warning(f"Error checking for dbt processes: {e}")
+        return False
+
+
+def wait_for_dbt_completion(project_dir: Path, timeout: float = 10.0, poll_interval: float = 0.2) -> bool:
+    """
+    Wait for any running dbt processes to complete.
+
+    Args:
+        project_dir: Path to the dbt project directory
+        timeout: Maximum time to wait in seconds (default: 10)
+        poll_interval: How often to check in seconds (default: 0.2)
+
+    Returns:
+        True if dbt finished or was not running, False if timeout occurred
+    """
+    import time
+
+    logger.debug(f"Waiting for dbt completion (timeout: {timeout}s)")
+
+    elapsed = 0.0
+    while elapsed < timeout:
+        if not is_dbt_running(project_dir):
+            logger.debug("dbt process check clear")
+            return True
+
+        logger.debug(f"DBT still running, waiting... ({elapsed:.1f}/{timeout}s)")
+        time.sleep(poll_interval)
+        elapsed += poll_interval
+
+    logger.warning(f"Timeout waiting for dbt to complete after {timeout}s")
+    return False