diff --git a/modelcontextprotocol/middleware.py b/modelcontextprotocol/middleware.py index c90db62a..7861eb30 100644 --- a/modelcontextprotocol/middleware.py +++ b/modelcontextprotocol/middleware.py @@ -100,7 +100,6 @@ async def on_call_tool(self, context: MiddlewareContext, call_next): # Tool is allowed, proceed with execution logger.debug(f"Tool access granted: {tool_name}", tool=tool_name) - return await call_next(context) except ToolError: diff --git a/modelcontextprotocol/settings.py b/modelcontextprotocol/settings.py index 23c22260..af3d0ac0 100644 --- a/modelcontextprotocol/settings.py +++ b/modelcontextprotocol/settings.py @@ -13,6 +13,7 @@ class Settings(BaseSettings): ATLAN_AGENT_ID: str = "NA" ATLAN_AGENT: str = "atlan-mcp" ATLAN_MCP_USER_AGENT: str = f"Atlan MCP Server {MCP_VERSION}" + ATLAN_TOOL_NAME: str = "X-Atlan-Mcp-Tool" MCP_TRANSPORT: str = "stdio" MCP_HOST: str = "0.0.0.0" MCP_PORT: int = 8000 diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index e6abf183..b2e17e75 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -10,6 +10,7 @@ ) from pyatlan.model.assets import Readme, AtlasGlossaryTerm from pyatlan.model.fluent_search import CompoundQuery, FluentSearch +from settings import get_settings # Initialize logging logger = logging.getLogger(__name__) @@ -47,7 +48,6 @@ def update_assets( logger.info( f"Updating {len(updatable_assets)} assets with attribute '{attribute_name}'" ) - # Validate attribute values if len(updatable_assets) != len(attribute_values): error_msg = "Number of asset GUIDs must match number of attribute values" @@ -67,7 +67,8 @@ def update_assets( # Get Atlan client client = get_atlan_client() - + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: "update_assets_tool"}) # Create assets with updated values assets = [] # readme_update_parent_assets: Assets that were updated with readme. diff --git a/modelcontextprotocol/tools/dsl.py b/modelcontextprotocol/tools/dsl.py index 9d3d4185..8e200bdb 100644 --- a/modelcontextprotocol/tools/dsl.py +++ b/modelcontextprotocol/tools/dsl.py @@ -5,6 +5,8 @@ from client import get_atlan_client from pyatlan.model.search import DSL, IndexSearchRequest from utils.search import SearchUtils +from settings import get_settings + # Configure logging logger = logging.getLogger(__name__) @@ -19,6 +21,7 @@ def get_assets_by_dsl(dsl_query: Union[str, Dict[str, Any]]) -> Dict[str, Any]: Dict[str, Any]: A dictionary containing the results and aggregations """ logger.info("Starting DSL-based asset search") + try: # Parse string to dict if needed if isinstance(dsl_query, str): @@ -47,6 +50,8 @@ def get_assets_by_dsl(dsl_query: Union[str, Dict[str, Any]]) -> Dict[str, Any]: logger.info("Executing DSL search request") client = get_atlan_client() + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: "get_assets_by_dsl_tool"}) search_response = client.asset.search(index_request) processed_results = SearchUtils.process_results(search_response) return processed_results diff --git a/modelcontextprotocol/tools/glossary.py b/modelcontextprotocol/tools/glossary.py index ccf8d964..38d8abf8 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -16,11 +16,12 @@ GlossaryCategory, GlossaryTerm, ) +from settings import get_settings logger = logging.getLogger(__name__) -def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]: +def save_assets(assets: List[Asset], tool_name: str) -> List[Dict[str, Any]]: """ Common bulk save and response processing for any asset type. @@ -34,7 +35,10 @@ def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]: Exception: If there's an error saving the assets. """ logger.info("Starting bulk save operation") + client = get_atlan_client() + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: tool_name}) try: response = client.asset.save(assets) except Exception as e: @@ -85,6 +89,7 @@ def create_glossary_assets( """ data = glossaries if isinstance(glossaries, list) else [glossaries] logger.info(f"Creating {len(data)} glossary asset(s)") + logger.debug(f"Glossary specifications: {data}") specs = [Glossary(**item) for item in data] @@ -104,7 +109,7 @@ def create_glossary_assets( logger.debug(f"Set certificate status for {spec.name}: {cs.value}") assets.append(glossary) - return save_assets(assets) + return save_assets(assets, "create_glossaries") def create_glossary_category_assets( @@ -137,6 +142,7 @@ def create_glossary_category_assets( """ data = categories if isinstance(categories, list) else [categories] logger.info(f"Creating {len(data)} glossary category asset(s)") + logger.debug(f"Category specifications: {data}") specs = [GlossaryCategory(**item) for item in data] @@ -166,7 +172,7 @@ def create_glossary_category_assets( assets.append(category) - return save_assets(assets) + return save_assets(assets, "create_glossary_categories") def create_glossary_term_assets( @@ -222,4 +228,4 @@ def create_glossary_term_assets( term.certificate_status = cs.value assets.append(term) - return save_assets(assets) + return save_assets(assets, "create_glossary_terms") diff --git a/modelcontextprotocol/tools/lineage.py b/modelcontextprotocol/tools/lineage.py index e8646c00..53d92d4f 100644 --- a/modelcontextprotocol/tools/lineage.py +++ b/modelcontextprotocol/tools/lineage.py @@ -7,6 +7,7 @@ from pyatlan.model.fields.atlan_fields import AtlanField from utils.search import SearchUtils from utils.constants import DEFAULT_SEARCH_ATTRIBUTES +from settings import get_settings # Configure logging logger = logging.getLogger(__name__) @@ -89,6 +90,8 @@ def traverse_lineage( logger.info("Executing lineage request") client = get_atlan_client() + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: "traverse_lineage_tool"}) response = client.asset.get_lineage_list(request) # Process results using same pattern as search diff --git a/modelcontextprotocol/tools/query.py b/modelcontextprotocol/tools/query.py index 951c9cdc..59a9512f 100644 --- a/modelcontextprotocol/tools/query.py +++ b/modelcontextprotocol/tools/query.py @@ -10,6 +10,7 @@ from client import get_atlan_client from pyatlan.model.query import QueryRequest +from settings import get_settings # Configure logging logger = logging.getLogger(__name__) @@ -78,7 +79,8 @@ def query_asset( # Get Atlan client logger.debug("Getting Atlan client") client = get_atlan_client() - + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: "query_asset_tool"}) # Build query request logger.debug("Building QueryRequest object") query_request = QueryRequest( diff --git a/modelcontextprotocol/tools/search.py b/modelcontextprotocol/tools/search.py index 3a1c3996..ced55481 100644 --- a/modelcontextprotocol/tools/search.py +++ b/modelcontextprotocol/tools/search.py @@ -7,6 +7,7 @@ from pyatlan.model.fields.atlan_fields import AtlanField from utils.search import SearchUtils from utils.constants import DEFAULT_SEARCH_ATTRIBUTES +from settings import get_settings # Configure logging logger = logging.getLogger(__name__) @@ -73,6 +74,8 @@ def search_assets( f"Starting asset search with parameters: asset_type={asset_type}, " f"limit={limit}, include_archived={include_archived}" ) + + # Set tool-specific headers logger.debug( f"Full search parameters: conditions={conditions}, " f"negative_conditions={negative_conditions}, some_conditions={some_conditions}, " @@ -291,6 +294,8 @@ def search_assets( logger.info("Executing search request") client = get_atlan_client() + settings = get_settings() + client.update_headers({settings.ATLAN_TOOL_NAME: "search_assets_tool"}) search_response = client.asset.search(request) processed_results = SearchUtils.process_results(search_response) logger.info(