Tool policies (#263)

dbt Labs' policies dictate that we do not feed row-level data to LLMs in
our systems. There is no such restriction for typical users of dbt-mcp
because they have their own terms with LLM providers. However, when we
use dbt-mcp internally, we need to ensure that we don't send tool
results with row-level data to LLMs. This PR applies policies to each
tool to ensure we are always compliant.

I had to change the get_model_details tool to return raw SQL instead of
compiled SQL in order for it to be compliant with metadata only.

Author: DevonFulcher

.changes/unreleased/Enhancement or New Feature-20250730-202622.yaml

@@ -0,0 +1,3 @@
+kind: Enhancement or New Feature
+body: 'Tool policies '
+time: 2025-07-30T20:26:22.429381-05:00

src/dbt_mcp/discovery/client.py

@@ -50,7 +50,7 @@ class GraphQLQueries:
                             node {
                                 name
                                 uniqueId
-                                compiledCode
+                                rawCode
                                 description
                                 database
                                 schema

src/dbt_mcp/tools/policy.py

@@ -0,0 +1,95 @@
+from enum import Enum
+
+from pydantic.dataclasses import dataclass
+
+from dbt_mcp.tools.tool_names import ToolName
+
+
+class ToolBehavior(Enum):
+    """Behavior of the tool."""
+
+    # The tool can return row-level data.
+    RESULT_SET = "result_set"
+    # The tool only returns metadata.
+    METADATA = "metadata"
+
+
+@dataclass
+class ToolPolicy:
+    """Policy for a tool."""
+
+    name: str
+    behavior: ToolBehavior
+
+
+# Defining tool policies is important for our internal usage of dbt-mcp.
+# Our policies dictate that we do not send row-level data to LLMs.
+tool_policies = {
+    # CLI tools
+    ToolName.SHOW.value: ToolPolicy(
+        name=ToolName.SHOW.value, behavior=ToolBehavior.RESULT_SET
+    ),
+    ToolName.LIST.value: ToolPolicy(
+        name=ToolName.LIST.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.DOCS.value: ToolPolicy(
+        name=ToolName.DOCS.value, behavior=ToolBehavior.METADATA
+    ),
+    # Compile tool can have result_set behavior because of macros like print_table
+    ToolName.COMPILE.value: ToolPolicy(
+        name=ToolName.COMPILE.value, behavior=ToolBehavior.RESULT_SET
+    ),
+    ToolName.TEST.value: ToolPolicy(
+        name=ToolName.TEST.value, behavior=ToolBehavior.METADATA
+    ),
+    # Run tool can have result_set behavior because of macros like print_table
+    ToolName.RUN.value: ToolPolicy(
+        name=ToolName.RUN.value, behavior=ToolBehavior.RESULT_SET
+    ),
+    # Build tool can have result_set behavior because of macros like print_table
+    ToolName.BUILD.value: ToolPolicy(
+        name=ToolName.BUILD.value, behavior=ToolBehavior.RESULT_SET
+    ),
+    ToolName.PARSE.value: ToolPolicy(
+        name=ToolName.PARSE.value, behavior=ToolBehavior.METADATA
+    ),
+    # Semantic Layer tools
+    ToolName.LIST_METRICS.value: ToolPolicy(
+        name=ToolName.LIST_METRICS.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_DIMENSIONS.value: ToolPolicy(
+        name=ToolName.GET_DIMENSIONS.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_ENTITIES.value: ToolPolicy(
+        name=ToolName.GET_ENTITIES.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.QUERY_METRICS.value: ToolPolicy(
+        name=ToolName.QUERY_METRICS.value, behavior=ToolBehavior.RESULT_SET
+    ),
+    ToolName.GET_METRICS_COMPILED_SQL.value: ToolPolicy(
+        name=ToolName.GET_METRICS_COMPILED_SQL.value, behavior=ToolBehavior.METADATA
+    ),
+    # Discovery tools
+    ToolName.GET_MODEL_PARENTS.value: ToolPolicy(
+        name=ToolName.GET_MODEL_PARENTS.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_MODEL_CHILDREN.value: ToolPolicy(
+        name=ToolName.GET_MODEL_CHILDREN.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_MODEL_DETAILS.value: ToolPolicy(
+        name=ToolName.GET_MODEL_DETAILS.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_MART_MODELS.value: ToolPolicy(
+        name=ToolName.GET_MART_MODELS.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.GET_ALL_MODELS.value: ToolPolicy(
+        name=ToolName.GET_ALL_MODELS.value, behavior=ToolBehavior.METADATA
+    ),
+    # SQL tools
+    ToolName.TEXT_TO_SQL.value: ToolPolicy(
+        name=ToolName.TEXT_TO_SQL.value, behavior=ToolBehavior.METADATA
+    ),
+    ToolName.EXECUTE_SQL.value: ToolPolicy(
+        name=ToolName.EXECUTE_SQL.value, behavior=ToolBehavior.RESULT_SET
+    ),
+}

tests/unit/tools/test_tool_names.py

@@ -9,20 +9,18 @@
 @pytest.mark.asyncio
 async def test_tool_names_match_server_tools():
     """Test that the ToolName enum matches the tools registered in the server."""
+    sql_tool_names = {"text_to_sql", "execute_sql"}
+
     with default_env_vars_context():
         config = load_config()
         dbt_mcp = await create_dbt_mcp(config)
 
         # 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 = {
-            n
-            for n in ToolName.get_all_tool_names()
-            if n
-            # Not testing SQL tools for now
-            not in ["text_to_sql", "execute_sql"]
-        }
+        # Manually adding SQL tools here because the server doesn't get them
+        # in this unit test.
+        server_tool_names = {tool.name for tool in server_tools} | sql_tool_names
+        enum_names = {n for n in ToolName.get_all_tool_names()}
 
         # This should not raise any errors if the enum is in sync
         if server_tool_names != enum_names:
@@ -37,3 +35,8 @@ async def test_tool_names_match_server_tools():
             assert isinstance(tool.value, str), (
                 f"Tool {tool.name} value should be a string"
             )
+
+
+def test_tool_names_no_duplicates():
+    """Test that there are no duplicate tool names in the enum."""
+    assert len(ToolName.get_all_tool_names()) == len(set(ToolName.get_all_tool_names()))

tests/unit/tools/test_tool_policies.py

@@ -0,0 +1,44 @@
+from dbt_mcp.config.config import load_config
+from dbt_mcp.mcp.server import create_dbt_mcp
+from dbt_mcp.tools.policy import tool_policies
+from dbt_mcp.tools.tool_names import ToolName
+from tests.env_vars import default_env_vars_context
+
+
+async def test_tool_policies_match_server_tools():
+    """Test that the ToolPolicy enum matches the tools registered in the server."""
+    sql_tool_names = {"text_to_sql", "execute_sql"}
+
+    with default_env_vars_context():
+        config = load_config()
+        dbt_mcp = await create_dbt_mcp(config)
+
+        # Get all tools from the server
+        server_tools = await dbt_mcp.list_tools()
+        # Manually adding SQL tools here because the server doesn't get them
+        # in this unit test.
+        server_tool_names = {tool.name for tool in server_tools} | sql_tool_names
+        policy_names = {policy_name for policy_name in tool_policies}
+
+        if server_tool_names != policy_names:
+            raise ValueError(
+                f"Tool name mismatch:\n"
+                f"In server but not in enum: {server_tool_names - policy_names}\n"
+                f"In enum but not in server: {policy_names - server_tool_names}"
+            )
+
+
+def test_tool_policies_match_tool_names():
+    policy_names = {policy.upper() for policy in tool_policies}
+    tool_names = {tool.name for tool in ToolName}
+    if tool_names != policy_names:
+        raise ValueError(
+            f"Tool name mismatch:\n"
+            f"In tool names but not in policy: {tool_names - policy_names}\n"
+            f"In policy but not in tool names: {policy_names - tool_names}"
+        )
+
+
+def test_tool_policies_no_duplicates():
+    """Test that there are no duplicate tool names in the policy."""
+    assert len(tool_policies) == len(set(tool_policies.keys()))