Allow passing semantic layer client in context

Author: wiggzz

src/dbt_mcp/dbt_cli/tools.py

@@ -2,7 +2,7 @@
 import subprocess
 from collections.abc import Iterable, Sequence
 
-from mcp.server.fastmcp import Context, FastMCP
+from mcp.server.fastmcp import FastMCP
 from pydantic import Field
 
 from dbt_mcp.config.config import DbtCliConfig

src/dbt_mcp/mcp/server.py

@@ -5,6 +5,7 @@
 from typing import Any
 
 from dbtlabs_vortex.producer import shutdown
+from dbtsl.client.sync import SyncSemanticLayerClient
 from mcp.server.fastmcp import FastMCP
 from mcp.types import ContentBlock, TextContent
 
@@ -109,6 +110,15 @@ def get_context(self) -> DbtMcpContext:
             request_context=request_context,
             fastmcp=self,
             semantic_layer_config=self.config.semantic_layer_config,
+            semantic_layer_client=(
+                SyncSemanticLayerClient(
+                    environment_id=self.config.semantic_layer_config.prod_environment_id,
+                    auth_token=self.config.semantic_layer_config.service_token,
+                    host=self.config.semantic_layer_config.host,
+                )
+                if self.config.semantic_layer_config
+                else None
+            ),
             discovery_config=self.config.discovery_config,
             dbt_cli_config=self.config.dbt_cli_config,
             admin_api_config=self.config.admin_api_config,

src/dbt_mcp/semantic_layer/tools.py

@@ -26,11 +26,7 @@
 
 def get_fetcher(ctx: DbtMcpContext) -> SemanticLayerFetcher:
     sl_config = ctx.get_semantic_layer_config()
-    sl_client = SyncSemanticLayerClient(
-        environment_id=sl_config.prod_environment_id,
-        auth_token=sl_config.service_token,
-        host=sl_config.host,
-    )
+    sl_client = ctx.get_semantic_layer_client()
     return SemanticLayerFetcher(
         sl_client=sl_client,
         config=sl_config,

src/dbt_mcp/tools/config.py

@@ -10,12 +10,14 @@
     SemanticLayerConfig,
 )
 from dbt_mcp.dbt_admin.client import DbtAdminAPIClient
+from dbt_mcp.semantic_layer.client import SemanticLayerClientProtocol
 
 
 class DbtMcpContext(Context[ServerSession, object, Request]):
     """Custom context for the MCP server"""
 
     _semantic_layer_config: SemanticLayerConfig | None = None
+    _semantic_layer_client: SemanticLayerClientProtocol | None = None
     _discovery_config: DiscoveryConfig | None = None
     _dbt_cli_config: DbtCliConfig | None = None
     _admin_api_config: AdminApiConfig | None = None
@@ -26,13 +28,15 @@ def __init__(
         request_context: RequestContext[ServerSession, object, Request] | None = None,
         fastmcp: FastMCP | None = None,
         semantic_layer_config: SemanticLayerConfig | None = None,
+        semantic_layer_client: SemanticLayerClientProtocol | None = None,
         discovery_config: DiscoveryConfig | None = None,
         dbt_cli_config: DbtCliConfig | None = None,
         admin_api_config: AdminApiConfig | None = None,
         admin_api_client: DbtAdminAPIClient | None = None,
     ):
         super().__init__(request_context=request_context, fastmcp=fastmcp)
         self._semantic_layer_config = semantic_layer_config
+        self._semantic_layer_client = semantic_layer_client
         self._discovery_config = discovery_config
         self._dbt_cli_config = dbt_cli_config
         self._admin_api_config = admin_api_config
@@ -43,6 +47,11 @@ def get_semantic_layer_config(self) -> SemanticLayerConfig:
             raise ValueError("Semantic layer config is not set")
         return self._semantic_layer_config
 
+    def get_semantic_layer_client(self) -> SemanticLayerClientProtocol:
+        if self._semantic_layer_client is None:
+            raise ValueError("Semantic layer client is not set")
+        return self._semantic_layer_client
+
     def get_discovery_config(self) -> DiscoveryConfig:
         if self._discovery_config is None:
             raise ValueError("Discovery config is not set")