From a759e05a155bb0dd67182ca3c03cf44cb216127f Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Wed, 17 Sep 2025 14:13:50 +0530 Subject: [PATCH 1/7] package name header update for tool name in logs --- modelcontextprotocol/middleware.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/modelcontextprotocol/middleware.py b/modelcontextprotocol/middleware.py index c90db62a..13e7141d 100644 --- a/modelcontextprotocol/middleware.py +++ b/modelcontextprotocol/middleware.py @@ -9,6 +9,7 @@ from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.exceptions import ToolError import logging +from client import get_atlan_client logger = logging.getLogger(__name__) @@ -100,6 +101,15 @@ 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) + try: + client = get_atlan_client() + client.update_headers({"x-atlan-package-name": tool_name}) + except Exception: + logger.warning( + "Could not set x-atlan-package-name header", + tool=tool_name, + exc_info=True, + ) return await call_next(context) From 363dd123d0fc46d096be6a5a32661c7854b96fb7 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Thu, 18 Sep 2025 14:59:01 +0530 Subject: [PATCH 2/7] update header at tool level --- modelcontextprotocol/middleware.py | 11 ----------- modelcontextprotocol/settings.py | 1 + modelcontextprotocol/tools/assets.py | 4 ++++ modelcontextprotocol/tools/dsl.py | 5 +++++ modelcontextprotocol/tools/glossary.py | 14 ++++++++++++++ modelcontextprotocol/tools/lineage.py | 4 ++++ modelcontextprotocol/tools/query.py | 3 +++ modelcontextprotocol/tools/search.py | 4 ++++ modelcontextprotocol/utils/__init__.py | 2 ++ modelcontextprotocol/utils/headers.py | 18 ++++++++++++++++++ 10 files changed, 55 insertions(+), 11 deletions(-) create mode 100644 modelcontextprotocol/utils/headers.py diff --git a/modelcontextprotocol/middleware.py b/modelcontextprotocol/middleware.py index 13e7141d..7861eb30 100644 --- a/modelcontextprotocol/middleware.py +++ b/modelcontextprotocol/middleware.py @@ -9,7 +9,6 @@ from fastmcp.server.middleware import Middleware, MiddlewareContext from fastmcp.exceptions import ToolError import logging -from client import get_atlan_client logger = logging.getLogger(__name__) @@ -101,16 +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) - try: - client = get_atlan_client() - client.update_headers({"x-atlan-package-name": tool_name}) - except Exception: - logger.warning( - "Could not set x-atlan-package-name header", - tool=tool_name, - exc_info=True, - ) - return await call_next(context) except ToolError: diff --git a/modelcontextprotocol/settings.py b/modelcontextprotocol/settings.py index ed2dcead..7c002988 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_PACKAGE_NAME: str = "x-atlan-package-name" @property def headers(self) -> dict: diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index e6abf183..24e86847 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 utils.headers import set_tool_headers # Initialize logging logger = logging.getLogger(__name__) @@ -47,6 +48,9 @@ def update_assets( logger.info( f"Updating {len(updatable_assets)} assets with attribute '{attribute_name}'" ) + + # Set tool-specific headers + set_tool_headers("update_assets_tool") # Validate attribute values if len(updatable_assets) != len(attribute_values): diff --git a/modelcontextprotocol/tools/dsl.py b/modelcontextprotocol/tools/dsl.py index 9d3d4185..48e76b26 100644 --- a/modelcontextprotocol/tools/dsl.py +++ b/modelcontextprotocol/tools/dsl.py @@ -5,6 +5,7 @@ from client import get_atlan_client from pyatlan.model.search import DSL, IndexSearchRequest from utils.search import SearchUtils +from utils.headers import set_tool_headers # Configure logging logger = logging.getLogger(__name__) @@ -19,6 +20,10 @@ 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") + + # Set tool-specific headers + set_tool_headers("get_assets_by_dsl_tool") + try: # Parse string to dict if needed if isinstance(dsl_query, str): diff --git a/modelcontextprotocol/tools/glossary.py b/modelcontextprotocol/tools/glossary.py index ccf8d964..84489275 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -10,6 +10,7 @@ ) from utils.parameters import parse_list_parameter from client import get_atlan_client +from utils.headers import set_tool_headers from .models import ( CertificateStatus, Glossary, @@ -34,6 +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") + + # Set tool-specific headers + set_tool_headers("glossary_save_operation") + client = get_atlan_client() try: response = client.asset.save(assets) @@ -85,6 +90,9 @@ def create_glossary_assets( """ data = glossaries if isinstance(glossaries, list) else [glossaries] logger.info(f"Creating {len(data)} glossary asset(s)") + + # Set tool-specific headers + set_tool_headers("create_glossaries") logger.debug(f"Glossary specifications: {data}") specs = [Glossary(**item) for item in data] @@ -137,6 +145,9 @@ def create_glossary_category_assets( """ data = categories if isinstance(categories, list) else [categories] logger.info(f"Creating {len(data)} glossary category asset(s)") + + # Set tool-specific headers + set_tool_headers("create_glossary_categories") logger.debug(f"Category specifications: {data}") specs = [GlossaryCategory(**item) for item in data] @@ -200,6 +211,9 @@ def create_glossary_term_assets( """ data = terms if isinstance(terms, list) else [terms] logger.info(f"Creating {len(data)} glossary term asset(s)") + + # Set tool-specific headers + set_tool_headers("create_glossary_terms") logger.debug(f"Term specifications: {data}") specs = [GlossaryTerm(**item) for item in data] diff --git a/modelcontextprotocol/tools/lineage.py b/modelcontextprotocol/tools/lineage.py index e8646c00..68a6e16c 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 utils.headers import set_tool_headers # Configure logging logger = logging.getLogger(__name__) @@ -50,6 +51,9 @@ def traverse_lineage( ) logger.debug(f"Include attributes parameter: {include_attributes}") + # Set tool-specific headers + set_tool_headers("traverse_lineage_tool") + try: # Initialize base request logger.debug("Initializing FluentLineage object") diff --git a/modelcontextprotocol/tools/query.py b/modelcontextprotocol/tools/query.py index 951c9cdc..3f95df8f 100644 --- a/modelcontextprotocol/tools/query.py +++ b/modelcontextprotocol/tools/query.py @@ -50,6 +50,9 @@ def query_asset( logger.info( f"Starting SQL query execution on connection: {connection_qualified_name}" ) + + # Set tool-specific headers + set_tool_headers("query_asset_tool") logger.debug(f"SQL query: {sql}") logger.debug(f"Parameters - default_schema: {default_schema}") diff --git a/modelcontextprotocol/tools/search.py b/modelcontextprotocol/tools/search.py index 3a1c3996..1641d7a9 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 utils.headers import set_tool_headers # Configure logging logger = logging.getLogger(__name__) @@ -73,6 +74,9 @@ def search_assets( f"Starting asset search with parameters: asset_type={asset_type}, " f"limit={limit}, include_archived={include_archived}" ) + + # Set tool-specific headers + set_tool_headers("search_assets_tool") logger.debug( f"Full search parameters: conditions={conditions}, " f"negative_conditions={negative_conditions}, some_conditions={some_conditions}, " diff --git a/modelcontextprotocol/utils/__init__.py b/modelcontextprotocol/utils/__init__.py index 6de6b6d6..a871d399 100644 --- a/modelcontextprotocol/utils/__init__.py +++ b/modelcontextprotocol/utils/__init__.py @@ -10,10 +10,12 @@ parse_json_parameter, parse_list_parameter, ) +from .headers import set_tool_headers __all__ = [ "DEFAULT_SEARCH_ATTRIBUTES", "SearchUtils", "parse_json_parameter", "parse_list_parameter", + "set_tool_headers", ] diff --git a/modelcontextprotocol/utils/headers.py b/modelcontextprotocol/utils/headers.py new file mode 100644 index 00000000..8631b521 --- /dev/null +++ b/modelcontextprotocol/utils/headers.py @@ -0,0 +1,18 @@ +"""Header utilities for Atlan MCP tools.""" + +import logging +from client import get_atlan_client +from settings import get_settings + +logger = logging.getLogger(__name__) + + +def set_tool_headers(tool_name: str) -> None: + """Set tool-specific headers for the Atlan client.""" + try: + settings = get_settings() + get_atlan_client().update_headers({settings.ATLAN_PACKAGE_NAME: tool_name}) + except Exception as e: + logger.warning(f"Could not set headers for {tool_name}: {e}") + + From b83c2ae0a924db01d00f9b8922155224ff4eb99b Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Thu, 18 Sep 2025 16:12:10 +0530 Subject: [PATCH 3/7] add tool name header at tool level --- modelcontextprotocol/settings.py | 2 +- modelcontextprotocol/tools/assets.py | 2 +- modelcontextprotocol/tools/dsl.py | 4 ++-- modelcontextprotocol/tools/glossary.py | 12 ++++++------ modelcontextprotocol/tools/query.py | 2 +- modelcontextprotocol/tools/search.py | 2 +- modelcontextprotocol/utils/headers.py | 2 -- 7 files changed, 12 insertions(+), 14 deletions(-) diff --git a/modelcontextprotocol/settings.py b/modelcontextprotocol/settings.py index 7c002988..c4d7dcaa 100644 --- a/modelcontextprotocol/settings.py +++ b/modelcontextprotocol/settings.py @@ -13,7 +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_PACKAGE_NAME: str = "x-atlan-package-name" + ATLAN_PACKAGE_NAME: str = "x-atlan-mcp-tool" @property def headers(self) -> dict: diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index 24e86847..4e12781f 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -48,7 +48,7 @@ def update_assets( logger.info( f"Updating {len(updatable_assets)} assets with attribute '{attribute_name}'" ) - + # Set tool-specific headers set_tool_headers("update_assets_tool") diff --git a/modelcontextprotocol/tools/dsl.py b/modelcontextprotocol/tools/dsl.py index 48e76b26..341b2b88 100644 --- a/modelcontextprotocol/tools/dsl.py +++ b/modelcontextprotocol/tools/dsl.py @@ -20,10 +20,10 @@ 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") - + # Set tool-specific headers set_tool_headers("get_assets_by_dsl_tool") - + try: # Parse string to dict if needed if isinstance(dsl_query, str): diff --git a/modelcontextprotocol/tools/glossary.py b/modelcontextprotocol/tools/glossary.py index 84489275..e81f5d10 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -35,10 +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") - - # Set tool-specific headers + + # Set tool-specific headers set_tool_headers("glossary_save_operation") - + client = get_atlan_client() try: response = client.asset.save(assets) @@ -90,7 +90,7 @@ def create_glossary_assets( """ data = glossaries if isinstance(glossaries, list) else [glossaries] logger.info(f"Creating {len(data)} glossary asset(s)") - + # Set tool-specific headers set_tool_headers("create_glossaries") logger.debug(f"Glossary specifications: {data}") @@ -145,7 +145,7 @@ def create_glossary_category_assets( """ data = categories if isinstance(categories, list) else [categories] logger.info(f"Creating {len(data)} glossary category asset(s)") - + # Set tool-specific headers set_tool_headers("create_glossary_categories") logger.debug(f"Category specifications: {data}") @@ -211,7 +211,7 @@ def create_glossary_term_assets( """ data = terms if isinstance(terms, list) else [terms] logger.info(f"Creating {len(data)} glossary term asset(s)") - + # Set tool-specific headers set_tool_headers("create_glossary_terms") logger.debug(f"Term specifications: {data}") diff --git a/modelcontextprotocol/tools/query.py b/modelcontextprotocol/tools/query.py index 3f95df8f..b11a2230 100644 --- a/modelcontextprotocol/tools/query.py +++ b/modelcontextprotocol/tools/query.py @@ -50,7 +50,7 @@ def query_asset( logger.info( f"Starting SQL query execution on connection: {connection_qualified_name}" ) - + # Set tool-specific headers set_tool_headers("query_asset_tool") logger.debug(f"SQL query: {sql}") diff --git a/modelcontextprotocol/tools/search.py b/modelcontextprotocol/tools/search.py index 1641d7a9..aad1fc8d 100644 --- a/modelcontextprotocol/tools/search.py +++ b/modelcontextprotocol/tools/search.py @@ -74,7 +74,7 @@ def search_assets( f"Starting asset search with parameters: asset_type={asset_type}, " f"limit={limit}, include_archived={include_archived}" ) - + # Set tool-specific headers set_tool_headers("search_assets_tool") logger.debug( diff --git a/modelcontextprotocol/utils/headers.py b/modelcontextprotocol/utils/headers.py index 8631b521..98fb9f41 100644 --- a/modelcontextprotocol/utils/headers.py +++ b/modelcontextprotocol/utils/headers.py @@ -14,5 +14,3 @@ def set_tool_headers(tool_name: str) -> None: get_atlan_client().update_headers({settings.ATLAN_PACKAGE_NAME: tool_name}) except Exception as e: logger.warning(f"Could not set headers for {tool_name}: {e}") - - From ea71011f5de2a1970990d8e8ec246979dc4ea323 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Fri, 19 Sep 2025 11:06:15 +0530 Subject: [PATCH 4/7] tool level headers --- modelcontextprotocol/server.py | 6 +++--- modelcontextprotocol/settings.py | 2 +- modelcontextprotocol/tools/assets.py | 9 +++------ modelcontextprotocol/tools/dsl.py | 8 ++++---- modelcontextprotocol/tools/glossary.py | 25 ++++++++++--------------- modelcontextprotocol/tools/lineage.py | 7 +++---- modelcontextprotocol/tools/query.py | 7 +++---- modelcontextprotocol/tools/search.py | 5 +++-- modelcontextprotocol/utils/__init__.py | 2 -- modelcontextprotocol/utils/headers.py | 16 ---------------- 10 files changed, 30 insertions(+), 57 deletions(-) delete mode 100644 modelcontextprotocol/utils/headers.py diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 515adfdb..dfde4587 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -766,7 +766,7 @@ def create_glossaries(glossaries) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for glossaries parameter: {str(e)}"} - return create_glossary_assets(glossaries) + return create_glossary_assets(glossaries, "create_glossaries") @mcp.tool() @@ -831,7 +831,7 @@ def create_glossary_terms(terms) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for terms parameter: {str(e)}"} - return create_glossary_term_assets(terms) + return create_glossary_term_assets(terms, "create_glossary_terms") @mcp.tool() @@ -900,7 +900,7 @@ def create_glossary_categories(categories) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for categories parameter: {str(e)}"} - return create_glossary_category_assets(categories) + return create_glossary_category_assets(categories, "create_glossary_categories") def main(): diff --git a/modelcontextprotocol/settings.py b/modelcontextprotocol/settings.py index c4d7dcaa..b40a1b13 100644 --- a/modelcontextprotocol/settings.py +++ b/modelcontextprotocol/settings.py @@ -13,7 +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_PACKAGE_NAME: str = "x-atlan-mcp-tool" + ATLAN_TOOL_NAME: str = "x-atlan-mcp-tool" @property def headers(self) -> dict: diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index 4e12781f..b2e17e75 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -10,7 +10,7 @@ ) from pyatlan.model.assets import Readme, AtlasGlossaryTerm from pyatlan.model.fluent_search import CompoundQuery, FluentSearch -from utils.headers import set_tool_headers +from settings import get_settings # Initialize logging logger = logging.getLogger(__name__) @@ -48,10 +48,6 @@ def update_assets( logger.info( f"Updating {len(updatable_assets)} assets with attribute '{attribute_name}'" ) - - # Set tool-specific headers - set_tool_headers("update_assets_tool") - # Validate attribute values if len(updatable_assets) != len(attribute_values): error_msg = "Number of asset GUIDs must match number of attribute values" @@ -71,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 341b2b88..8e200bdb 100644 --- a/modelcontextprotocol/tools/dsl.py +++ b/modelcontextprotocol/tools/dsl.py @@ -5,7 +5,8 @@ from client import get_atlan_client from pyatlan.model.search import DSL, IndexSearchRequest from utils.search import SearchUtils -from utils.headers import set_tool_headers +from settings import get_settings + # Configure logging logger = logging.getLogger(__name__) @@ -21,9 +22,6 @@ def get_assets_by_dsl(dsl_query: Union[str, Dict[str, Any]]) -> Dict[str, Any]: """ logger.info("Starting DSL-based asset search") - # Set tool-specific headers - set_tool_headers("get_assets_by_dsl_tool") - try: # Parse string to dict if needed if isinstance(dsl_query, str): @@ -52,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 e81f5d10..ceeea8d3 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -10,18 +10,18 @@ ) from utils.parameters import parse_list_parameter from client import get_atlan_client -from utils.headers import set_tool_headers from .models import ( CertificateStatus, Glossary, 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. @@ -36,10 +36,9 @@ def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]: """ logger.info("Starting bulk save operation") - # Set tool-specific headers - set_tool_headers("glossary_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: @@ -65,6 +64,7 @@ def save_assets(assets: List[Asset]) -> List[Dict[str, Any]]: def create_glossary_assets( glossaries: Union[Dict[str, Any], List[Dict[str, Any]]], + tool_name: str, ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossary assets in Atlan. @@ -91,8 +91,6 @@ def create_glossary_assets( data = glossaries if isinstance(glossaries, list) else [glossaries] logger.info(f"Creating {len(data)} glossary asset(s)") - # Set tool-specific headers - set_tool_headers("create_glossaries") logger.debug(f"Glossary specifications: {data}") specs = [Glossary(**item) for item in data] @@ -112,11 +110,12 @@ 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, tool_name) def create_glossary_category_assets( categories: Union[Dict[str, Any], List[Dict[str, Any]]], + tool_name: str, ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossaryCategory assets in Atlan. @@ -146,8 +145,6 @@ def create_glossary_category_assets( data = categories if isinstance(categories, list) else [categories] logger.info(f"Creating {len(data)} glossary category asset(s)") - # Set tool-specific headers - set_tool_headers("create_glossary_categories") logger.debug(f"Category specifications: {data}") specs = [GlossaryCategory(**item) for item in data] @@ -177,11 +174,12 @@ def create_glossary_category_assets( assets.append(category) - return save_assets(assets) + return save_assets(assets, tool_name) def create_glossary_term_assets( terms: Union[Dict[str, Any], List[Dict[str, Any]]], + tool_name: str, ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossaryTerm assets in Atlan. @@ -211,9 +209,6 @@ def create_glossary_term_assets( """ data = terms if isinstance(terms, list) else [terms] logger.info(f"Creating {len(data)} glossary term asset(s)") - - # Set tool-specific headers - set_tool_headers("create_glossary_terms") logger.debug(f"Term specifications: {data}") specs = [GlossaryTerm(**item) for item in data] @@ -236,4 +231,4 @@ def create_glossary_term_assets( term.certificate_status = cs.value assets.append(term) - return save_assets(assets) + return save_assets(assets, tool_name) diff --git a/modelcontextprotocol/tools/lineage.py b/modelcontextprotocol/tools/lineage.py index 68a6e16c..53d92d4f 100644 --- a/modelcontextprotocol/tools/lineage.py +++ b/modelcontextprotocol/tools/lineage.py @@ -7,7 +7,7 @@ from pyatlan.model.fields.atlan_fields import AtlanField from utils.search import SearchUtils from utils.constants import DEFAULT_SEARCH_ATTRIBUTES -from utils.headers import set_tool_headers +from settings import get_settings # Configure logging logger = logging.getLogger(__name__) @@ -51,9 +51,6 @@ def traverse_lineage( ) logger.debug(f"Include attributes parameter: {include_attributes}") - # Set tool-specific headers - set_tool_headers("traverse_lineage_tool") - try: # Initialize base request logger.debug("Initializing FluentLineage object") @@ -93,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 b11a2230..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__) @@ -50,9 +51,6 @@ def query_asset( logger.info( f"Starting SQL query execution on connection: {connection_qualified_name}" ) - - # Set tool-specific headers - set_tool_headers("query_asset_tool") logger.debug(f"SQL query: {sql}") logger.debug(f"Parameters - default_schema: {default_schema}") @@ -81,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 aad1fc8d..ced55481 100644 --- a/modelcontextprotocol/tools/search.py +++ b/modelcontextprotocol/tools/search.py @@ -7,7 +7,7 @@ from pyatlan.model.fields.atlan_fields import AtlanField from utils.search import SearchUtils from utils.constants import DEFAULT_SEARCH_ATTRIBUTES -from utils.headers import set_tool_headers +from settings import get_settings # Configure logging logger = logging.getLogger(__name__) @@ -76,7 +76,6 @@ def search_assets( ) # Set tool-specific headers - set_tool_headers("search_assets_tool") logger.debug( f"Full search parameters: conditions={conditions}, " f"negative_conditions={negative_conditions}, some_conditions={some_conditions}, " @@ -295,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( diff --git a/modelcontextprotocol/utils/__init__.py b/modelcontextprotocol/utils/__init__.py index a871d399..6de6b6d6 100644 --- a/modelcontextprotocol/utils/__init__.py +++ b/modelcontextprotocol/utils/__init__.py @@ -10,12 +10,10 @@ parse_json_parameter, parse_list_parameter, ) -from .headers import set_tool_headers __all__ = [ "DEFAULT_SEARCH_ATTRIBUTES", "SearchUtils", "parse_json_parameter", "parse_list_parameter", - "set_tool_headers", ] diff --git a/modelcontextprotocol/utils/headers.py b/modelcontextprotocol/utils/headers.py deleted file mode 100644 index 98fb9f41..00000000 --- a/modelcontextprotocol/utils/headers.py +++ /dev/null @@ -1,16 +0,0 @@ -"""Header utilities for Atlan MCP tools.""" - -import logging -from client import get_atlan_client -from settings import get_settings - -logger = logging.getLogger(__name__) - - -def set_tool_headers(tool_name: str) -> None: - """Set tool-specific headers for the Atlan client.""" - try: - settings = get_settings() - get_atlan_client().update_headers({settings.ATLAN_PACKAGE_NAME: tool_name}) - except Exception as e: - logger.warning(f"Could not set headers for {tool_name}: {e}") From 86e11791ca46875a586608f02bc9a14772313940 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Fri, 19 Sep 2025 12:31:54 +0530 Subject: [PATCH 5/7] pass tool name directly --- modelcontextprotocol/tools/glossary.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/modelcontextprotocol/tools/glossary.py b/modelcontextprotocol/tools/glossary.py index ceeea8d3..01842bcd 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -63,8 +63,7 @@ def save_assets(assets: List[Asset], tool_name: str) -> List[Dict[str, Any]]: def create_glossary_assets( - glossaries: Union[Dict[str, Any], List[Dict[str, Any]]], - tool_name: str, + glossaries: Union[Dict[str, Any], List[Dict[str, Any]]] ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossary assets in Atlan. @@ -110,12 +109,11 @@ def create_glossary_assets( logger.debug(f"Set certificate status for {spec.name}: {cs.value}") assets.append(glossary) - return save_assets(assets, tool_name) + return save_assets(assets,"create_glossaries") def create_glossary_category_assets( categories: Union[Dict[str, Any], List[Dict[str, Any]]], - tool_name: str, ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossaryCategory assets in Atlan. @@ -174,12 +172,11 @@ def create_glossary_category_assets( assets.append(category) - return save_assets(assets, tool_name) + return save_assets(assets, "create_glossary_categories") def create_glossary_term_assets( terms: Union[Dict[str, Any], List[Dict[str, Any]]], - tool_name: str, ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossaryTerm assets in Atlan. @@ -231,4 +228,4 @@ def create_glossary_term_assets( term.certificate_status = cs.value assets.append(term) - return save_assets(assets, tool_name) + return save_assets(assets, "create_glossary_terms") From 1cf29c27dc109f01fe69c7fcf2883108964b7dd7 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Fri, 19 Sep 2025 12:33:22 +0530 Subject: [PATCH 6/7] header case change --- modelcontextprotocol/settings.py | 2 +- modelcontextprotocol/tools/glossary.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/modelcontextprotocol/settings.py b/modelcontextprotocol/settings.py index b40a1b13..efb9fa6b 100644 --- a/modelcontextprotocol/settings.py +++ b/modelcontextprotocol/settings.py @@ -13,7 +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" + ATLAN_TOOL_NAME: str = "X-Atlan-Mcp-Tool" @property def headers(self) -> dict: diff --git a/modelcontextprotocol/tools/glossary.py b/modelcontextprotocol/tools/glossary.py index 01842bcd..38d8abf8 100644 --- a/modelcontextprotocol/tools/glossary.py +++ b/modelcontextprotocol/tools/glossary.py @@ -63,7 +63,7 @@ def save_assets(assets: List[Asset], tool_name: str) -> List[Dict[str, Any]]: def create_glossary_assets( - glossaries: Union[Dict[str, Any], List[Dict[str, Any]]] + glossaries: Union[Dict[str, Any], List[Dict[str, Any]]], ) -> List[Dict[str, Any]]: """ Create one or multiple AtlasGlossary assets in Atlan. @@ -109,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,"create_glossaries") + return save_assets(assets, "create_glossaries") def create_glossary_category_assets( From 46bb6719ed5ec52af3f704f51d72f793fb23a521 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Thu, 25 Sep 2025 11:06:51 +0530 Subject: [PATCH 7/7] server-fix --- modelcontextprotocol/server.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index dfde4587..515adfdb 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -766,7 +766,7 @@ def create_glossaries(glossaries) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for glossaries parameter: {str(e)}"} - return create_glossary_assets(glossaries, "create_glossaries") + return create_glossary_assets(glossaries) @mcp.tool() @@ -831,7 +831,7 @@ def create_glossary_terms(terms) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for terms parameter: {str(e)}"} - return create_glossary_term_assets(terms, "create_glossary_terms") + return create_glossary_term_assets(terms) @mcp.tool() @@ -900,7 +900,7 @@ def create_glossary_categories(categories) -> List[Dict[str, Any]]: except json.JSONDecodeError as e: return {"error": f"Invalid JSON format for categories parameter: {str(e)}"} - return create_glossary_category_assets(categories, "create_glossary_categories") + return create_glossary_category_assets(categories) def main():