Skip to content

Commit 5f8fbd8

Browse files
committed
updated, addressed comments
1 parent 205f225 commit 5f8fbd8

File tree

3 files changed

+8
-78
lines changed

3 files changed

+8
-78
lines changed

modelcontextprotocol/server.py

Lines changed: 1 addition & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -511,12 +511,7 @@ def update_assets_tool(
511511

512512
@mcp.tool()
513513
def get_asset_history_tool(
514-
guid=None,
515-
qualified_name=None,
516-
type_name=None,
517-
size=10,
518-
sort_order="DESC",
519-
include_attributes=None,
514+
guid=None, qualified_name=None, type_name=None, size=10, sort_order="DESC"
520515
):
521516
"""
522517
Get the audit history of an asset by GUID or qualified name.
@@ -531,8 +526,6 @@ def get_asset_history_tool(
531526
size (int): Number of history entries to return. Defaults to 10.
532527
sort_order (str): Sort order for results. "ASC" for oldest first, "DESC" for newest first.
533528
Defaults to "DESC".
534-
include_attributes (List[Union[str, AtlanField]], optional): List of additional attributes to include in results.
535-
Can be string attribute names or AtlanField objects. These will be added to the default set.
536529
537530
Returns:
538531
Dict[str, Any]: Dictionary containing:
@@ -547,7 +540,6 @@ def get_asset_history_tool(
547540
guid="6fc01478-1263-42ae-b8ca-c4a57da51392",
548541
size=20,
549542
sort_order="DESC",
550-
include_attributes=["owner_users", "description"]
551543
)
552544
553545
# Get history by qualified name
@@ -564,7 +556,6 @@ def get_asset_history_tool(
564556
type_name=type_name,
565557
size=size,
566558
sort_order=sort_order,
567-
include_attributes=include_attributes,
568559
)
569560

570561

modelcontextprotocol/tools/assets.py

Lines changed: 2 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -4,13 +4,11 @@
44
from .models import UpdatableAsset, UpdatableAttribute, CertificateStatus
55
from pyatlan.model.assets import Readme
66
from pyatlan.model.fluent_search import CompoundQuery, FluentSearch
7-
from pyatlan.model.fields.atlan_fields import AtlanField
87
from utils.asset_history import (
98
validate_asset_history_params,
109
create_audit_search_request,
1110
process_audit_result,
1211
create_sort_item,
13-
convert_attributes_to_camel_case,
1412
)
1513

1614
# Initialize logging
@@ -144,7 +142,6 @@ def get_asset_history(
144142
type_name: Optional[str] = None,
145143
size: int = 10,
146144
sort_order: str = "DESC",
147-
include_attributes: Optional[List[Union[str, AtlanField]]] = None,
148145
) -> Dict[str, Any]:
149146
"""
150147
Get the audit history of an asset by GUID or qualified name.
@@ -159,8 +156,6 @@ def get_asset_history(
159156
size (int): Number of history entries to return. Defaults to 10.
160157
sort_order (str): Sort order for results. "ASC" for oldest first, "DESC" for newest first.
161158
Defaults to "DESC".
162-
include_attributes (List[Union[str, AtlanField]], optional): List of additional attributes to include in results.
163-
Can be string attribute names or AtlanField objects. These will be added to the default set.
164159
165160
Returns:
166161
Dict[str, Any]: Dictionary containing:
@@ -196,20 +191,15 @@ def get_asset_history(
196191
# Create sort item
197192
sort_item = create_sort_item(sort_order)
198193

199-
# Convert include_attributes from snake_case to camelCase if needed
200-
if include_attributes:
201-
include_attributes = convert_attributes_to_camel_case(include_attributes)
202-
203194
# Create and execute audit search request
204195
request = create_audit_search_request(
205-
guid, qualified_name, type_name, size, sort_item, include_attributes
196+
guid, qualified_name, type_name, size, sort_item
206197
)
207198
response = client.audit.search(criteria=request, bulk=False)
208199

209200
# Process audit results - use current_page() to respect size parameter
210201
entity_audits = [
211-
process_audit_result(result, include_attributes)
212-
for result in response.current_page()
202+
process_audit_result(result) for result in response.current_page()
213203
]
214204

215205
result_data = {

modelcontextprotocol/utils/asset_history.py

Lines changed: 5 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -5,34 +5,15 @@
55
"""
66

77
import logging
8-
from typing import List, Union, Dict, Any, Optional
8+
from typing import Dict, Any, Optional
99
from pyatlan.client.audit import AuditSearchRequest
1010
from pyatlan.model.search import DSL, Bool, SortItem, Term
1111
from pyatlan.model.enums import SortOrder
12-
from pyatlan.model.fields.atlan_fields import AtlanField
1312

1413
# Initialize logging
1514
logger = logging.getLogger(__name__)
1615

1716

18-
def snake_to_camel(s: str) -> str:
19-
"""
20-
Convert a snake_case string to camelCase.
21-
22-
Args:
23-
s (str): The snake_case string.
24-
25-
Returns:
26-
str: The camelCase string.
27-
"""
28-
parts = s.split("_")
29-
return (
30-
parts[0] + "".join(word.capitalize() for word in parts[1:])
31-
if len(parts) > 1
32-
else s
33-
)
34-
35-
3617
def validate_asset_history_params(
3718
guid: Optional[str],
3819
qualified_name: Optional[str],
@@ -69,7 +50,6 @@ def create_audit_search_request(
6950
type_name: Optional[str],
7051
size: int,
7152
sort_item: SortItem,
72-
include_attributes: Optional[List[Union[str, AtlanField]]],
7353
) -> AuditSearchRequest:
7454
"""
7555
Create an AuditSearchRequest based on the provided parameters.
@@ -80,7 +60,6 @@ def create_audit_search_request(
8060
type_name: Asset type name
8161
size: Number of results to return
8262
sort_item: Sort configuration
83-
include_attributes: Attributes to include in results
8463
8564
Returns:
8665
Configured AuditSearchRequest
@@ -107,11 +86,7 @@ def create_audit_search_request(
10786
f"Created audit search request by qualified name: {qualified_name}"
10887
)
10988

110-
# Only pass attributes if they are provided
111-
if include_attributes:
112-
return AuditSearchRequest(dsl=dsl, attributes=include_attributes)
113-
else:
114-
return AuditSearchRequest(dsl=dsl)
89+
return AuditSearchRequest(dsl=dsl)
11590

11691

11792
def extract_basic_audit_info(result) -> Dict[str, Any]:
@@ -135,29 +110,21 @@ def extract_basic_audit_info(result) -> Dict[str, Any]:
135110
}
136111

137112

138-
def process_audit_result(
139-
result, include_attributes: Optional[List[Union[str, AtlanField]]]
140-
) -> Dict[str, Any]:
113+
def process_audit_result(result) -> Dict[str, Any]:
141114
"""
142115
Process a single audit result into a formatted dictionary.
143116
144117
Args:
145118
result: Audit result object
146-
include_attributes: List of attributes to include
147119
148120
Returns:
149121
Formatted audit entry dictionary
150122
"""
151123
try:
152124
# Extract basic audit information
153125
audit_entry = extract_basic_audit_info(result)
154-
155-
# Add requested attributes if specified
156-
if include_attributes:
157-
requested_attrs = result.detail.attributes.dict(
158-
exclude_unset=True, by_alias=True
159-
)
160-
audit_entry.update(requested_attrs)
126+
updates = result.detail.dict(exclude_unset=True)
127+
audit_entry.update(updates)
161128

162129
return audit_entry
163130

@@ -184,21 +151,3 @@ def create_sort_item(sort_order: str) -> SortItem:
184151
"created",
185152
order=SortOrder.DESCENDING if sort_order == "DESC" else SortOrder.ASCENDING,
186153
)
187-
188-
189-
def convert_attributes_to_camel_case(
190-
include_attributes: List[Union[str, AtlanField]],
191-
) -> List[Union[str, AtlanField]]:
192-
"""
193-
Convert string attributes from snake_case to camelCase.
194-
195-
Args:
196-
include_attributes: List of attributes to convert
197-
198-
Returns:
199-
List with string attributes converted to camelCase
200-
"""
201-
return [
202-
snake_to_camel(attr) if isinstance(attr, str) else attr
203-
for attr in include_attributes
204-
]

0 commit comments

Comments
 (0)