Skip to content
Open
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 0 additions & 1 deletion modelcontextprotocol/middleware.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
6 changes: 3 additions & 3 deletions modelcontextprotocol/server.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down Expand Up @@ -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")
Comment thread
abhinavmathur-atlan marked this conversation as resolved.
Outdated


@mcp.tool()
Expand Down Expand Up @@ -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():
Expand Down
1 change: 1 addition & 0 deletions modelcontextprotocol/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Comment thread
abhinavmathur-atlan marked this conversation as resolved.
Outdated

@property
def headers(self) -> dict:
Expand Down
5 changes: 3 additions & 2 deletions modelcontextprotocol/tools/assets.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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"
Expand All @@ -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.
Expand Down
5 changes: 5 additions & 0 deletions modelcontextprotocol/tools/dsl.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand All @@ -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):
Expand Down Expand Up @@ -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
Expand Down
17 changes: 13 additions & 4 deletions modelcontextprotocol/tools/glossary.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand All @@ -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:
Expand All @@ -60,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,
Comment thread
abhinavmathur-atlan marked this conversation as resolved.
Outdated
) -> List[Dict[str, Any]]:
"""
Create one or multiple AtlasGlossary assets in Atlan.
Expand All @@ -85,6 +90,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]
Expand All @@ -104,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)
Comment thread
abhinavmathur-atlan marked this conversation as resolved.
Outdated


def create_glossary_category_assets(
categories: Union[Dict[str, Any], List[Dict[str, Any]]],
tool_name: str,
Comment thread
abhinavmathur-atlan marked this conversation as resolved.
Outdated
) -> List[Dict[str, Any]]:
"""
Create one or multiple AtlasGlossaryCategory assets in Atlan.
Expand Down Expand Up @@ -137,6 +144,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]
Expand Down Expand Up @@ -166,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.
Expand Down Expand Up @@ -222,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)
3 changes: 3 additions & 0 deletions modelcontextprotocol/tools/lineage.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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
Expand Down
4 changes: 3 additions & 1 deletion modelcontextprotocol/tools/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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(
Expand Down
5 changes: 5 additions & 0 deletions modelcontextprotocol/tools/search.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)
Expand Down Expand Up @@ -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}, "
Expand Down Expand Up @@ -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(
Expand Down