Create list of tool names

Author: DevonFulcher

src/dbt_mcp/config/config.py

@@ -5,6 +5,8 @@
 import yaml
 from dotenv import load_dotenv
 
+from dbt_mcp.tools.tool_names import ToolName
+
 
 @dataclass
 class TrackingConfig:
@@ -56,7 +58,7 @@ class Config:
     dbt_cli_config: DbtCliConfig | None
     discovery_config: DiscoveryConfig | None
     semantic_layer_config: SemanticLayerConfig | None
-    disable_tools: list[str]
+    disable_tools: list[ToolName]
 
 
 def load_config() -> Config:
@@ -77,7 +79,10 @@ def load_config() -> Config:
     disable_remote = os.environ.get("DISABLE_REMOTE", "true") == "true"
     multicell_account_prefix = os.environ.get("MULTICELL_ACCOUNT_PREFIX", None)
     dbt_cli_timeout = int(os.environ.get("DBT_CLI_TIMEOUT", 10))
-    disable_tools = os.environ.get("DISABLE_TOOLS", "").split(",")
+    disable_tools = [
+        ToolName(tool_name)
+        for tool_name in os.environ.get("DISABLE_TOOLS", "").split(",")
+    ]
 
     # set default warn error options if not provided
     if os.environ.get("DBT_WARN_ERROR_OPTIONS") is None:

src/dbt_mcp/discovery/tools.py

@@ -1,13 +1,14 @@
 import logging
-
 from collections.abc import Sequence
+
 from mcp.server.fastmcp import FastMCP
 
 from dbt_mcp.config.config import DiscoveryConfig
 from dbt_mcp.discovery.client import MetadataAPIClient, ModelsFetcher
 from dbt_mcp.prompts.prompts import get_prompt
 from dbt_mcp.tools.definitions import ToolDefinition
 from dbt_mcp.tools.register import register_tools
+from dbt_mcp.tools.tool_names import ToolName
 
 logger = logging.getLogger(__name__)
 
@@ -85,7 +86,7 @@ def get_model_children(
 def register_discovery_tools(
     dbt_mcp: FastMCP,
     config: DiscoveryConfig,
-    exclude_tools: Sequence[str] = [],
+    exclude_tools: Sequence[ToolName] = [],
 ) -> None:
     register_tools(
         dbt_mcp,

src/dbt_mcp/mcp/server.py

@@ -99,7 +99,8 @@ async def create_dbt_mcp():
 
     if config.dbt_cli_config:
         logger.info("Registering dbt cli tools")
-        register_dbt_cli_tools(dbt_mcp, config.dbt_cli_config, config.disable_tools)
+        # TODO: allow for disabling CLI tools
+        register_dbt_cli_tools(dbt_mcp, config.dbt_cli_config)
 
     if config.remote_config:
         logger.info("Registering remote tools")

src/dbt_mcp/remote/tools.py

@@ -26,6 +26,7 @@
 from pydantic_core import PydanticUndefined
 
 from dbt_mcp.config.config import RemoteConfig
+from dbt_mcp.tools.tool_names import ToolName
 
 logger = logging.getLogger(__name__)
 
@@ -72,7 +73,7 @@ def _get_remote_tools(base_url: str, headers: dict[str, str]) -> list[RemoteTool
 async def register_remote_tools(
     dbt_mcp: FastMCP,
     config: RemoteConfig,
-    exclude_tools: Sequence[str] = [],
+    exclude_tools: Sequence[ToolName] = [],
 ) -> None:
     is_local = config.host and config.host.startswith("localhost")
     path = "/mcp" if is_local else "/api/ai/mcp"
@@ -92,7 +93,7 @@ async def register_remote_tools(
         f"Loaded remote tools: {', '.join([tool.name for tool in remote_tools])}",
     )
     for tool in remote_tools:
-        if tool.name in exclude_tools:
+        if tool.name.lower() in [tool.value.lower() for tool in exclude_tools]:
             continue
 
         # Create a new function using a factory to avoid closure issues

src/dbt_mcp/semantic_layer/tools.py

@@ -1,5 +1,5 @@
-from collections.abc import Sequence
 import logging
+from collections.abc import Sequence
 
 from dbtsl.api.shared.query_params import GroupByParam
 from dbtsl.client.sync import SyncSemanticLayerClient
@@ -20,6 +20,7 @@
 )
 from dbt_mcp.tools.definitions import ToolDefinition
 from dbt_mcp.tools.register import register_tools
+from dbt_mcp.tools.tool_names import ToolName
 
 logger = logging.getLogger(__name__)
 
@@ -95,7 +96,7 @@ def query_metrics(
 def register_sl_tools(
     dbt_mcp: FastMCP,
     config: SemanticLayerConfig,
-    exclude_tools: Sequence[str] = [],
+    exclude_tools: Sequence[ToolName] = [],
 ) -> None:
     register_tools(
         dbt_mcp,

src/dbt_mcp/tools/register.py

@@ -1,16 +1,20 @@
 from collections.abc import Sequence
+
 from mcp.server.fastmcp import FastMCP
 
 from dbt_mcp.tools.definitions import ToolDefinition
+from dbt_mcp.tools.tool_names import ToolName
 
 
 def register_tools(
     dbt_mcp: FastMCP,
     tool_definitions: list[ToolDefinition],
-    exclude_tools: Sequence[str] = [],
+    exclude_tools: Sequence[ToolName] = [],
 ) -> None:
     for tool_definition in tool_definitions:
-        if tool_definition.get_name() in exclude_tools:
+        if tool_definition.get_name().lower() in [
+            tool.value.lower() for tool in exclude_tools
+        ]:
             continue
         dbt_mcp.tool(
             name=tool_definition.get_name(),

src/dbt_mcp/tools/tool_names.py

@@ -0,0 +1,42 @@
+from enum import Enum
+
+
+class ToolName(Enum):
+    """Tool names available in the FastMCP server.
+
+    This enum provides type safety and autocompletion for tool names.
+    The validate_server_tools() function should be used to ensure
+    this enum stays in sync with the actual server tools.
+    """
+
+    # dbt CLI tools
+    BUILD = "build"
+    COMPILE = "compile"
+    DOCS = "docs"
+    LIST = "list"
+    PARSE = "parse"
+    RUN = "run"
+    TEST = "test"
+    SHOW = "show"
+
+    # Semantic Layer tools
+    LIST_METRICS = "list_metrics"
+    GET_DIMENSIONS = "get_dimensions"
+    GET_ENTITIES = "get_entities"
+    QUERY_METRICS = "query_metrics"
+
+    # Discovery tools
+    GET_MART_MODELS = "get_mart_models"
+    GET_ALL_MODELS = "get_all_models"
+    GET_MODEL_DETAILS = "get_model_details"
+    GET_MODEL_PARENTS = "get_model_parents"
+    GET_MODEL_CHILDREN = "get_model_children"
+
+    # Remote tools
+    TEXT_TO_SQL = "text_to_sql"
+    EXECUTE_SQL = "execute_sql"
+
+    @classmethod
+    def get_all_tool_names(cls) -> set[str]:
+        """Returns a set of all tool names as strings."""
+        return {member.value for member in cls}

tests/unit/tools/test_tool_names.py

@@ -0,0 +1,28 @@
+import pytest
+
+from dbt_mcp.mcp.server import create_dbt_mcp
+from dbt_mcp.tools.tool_names import ToolName
+
+
+@pytest.mark.asyncio
+async def test_tool_names_match_server_tools():
+    """Test that the ToolName enum matches the tools registered in the server."""
+    # Create the dbt_mcp server instance
+    dbt_mcp = await create_dbt_mcp()
+
+    # Get all tools from the server
+    server_tools = await dbt_mcp.list_tools()
+    server_tool_names = {tool.name for tool in server_tools}
+    enum_names = ToolName.get_all_tool_names()
+
+    # This should not raise any errors if the enum is in sync
+    if server_tool_names != enum_names:
+        raise ValueError(
+            f"Tool name mismatch:\n"
+            f"In server but not in enum: {server_tool_names - enum_names}\n"
+            f"In enum but not in server: {enum_names - server_tool_names}"
+        )
+
+    # Double check that all enum values are strings
+    for tool in ToolName:
+        assert isinstance(tool.value, str), f"Tool {tool.name} value should be a string"