From 3d3d4b7c0cc2e872060c1c88f8807ebd10467047 Mon Sep 17 00:00:00 2001 From: Bala Date: Mon, 15 Sep 2025 09:30:11 +0530 Subject: [PATCH 1/2] Add support to update_Assets_tool with update classifications d122674afe61599ada0cf5bfe6e0408f56490826 author Bala 1757908811 +0530 committer Bala 1758000441 +0530 Add support to update_Assets_tool with update classifications update Args instructions remove unused imports updated review comments Reverted all changes to uv lock --- modelcontextprotocol/server.py | 16 +++++++++++++++- modelcontextprotocol/tools/assets.py | 18 ++++++++++++++++-- modelcontextprotocol/tools/models.py | 3 +++ modelcontextprotocol/uv.lock | 2 +- 4 files changed, 35 insertions(+), 4 deletions(-) diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 515adfd..5e89071 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -478,11 +478,13 @@ def update_assets_tool( assets (Union[Dict[str, Any], List[Dict[str, Any]]]): Asset(s) to update. Can be a single UpdatableAsset or a list of UpdatableAsset objects. attribute_name (str): Name of the attribute to update. - Supports "user_description", "certificate_status", "readme", and "term". + Supports "user_description", "certificate_status", "readme", "classifications" and "term". attribute_values (List[Union[str, Dict[str, Any]]]): List of values to set for the attribute. For certificateStatus, only "VERIFIED", "DRAFT", or "DEPRECATED" are allowed. For readme, the value must be a valid Markdown string. For term, the value must be a dict with "operation" and "term_guids" keys. + For classifications, the value must already exist in Atlan. + Returns: Dict[str, Any]: Dictionary containing: @@ -502,6 +504,17 @@ def update_assets_tool( attribute_name="certificate_status", attribute_values=["VERIFIED"] ) + # Update classification for a single asset + result = update_assets_tool( + assets={ + "guid": "asset-guid-here", + "name": "Asset Name", + "type_name": "Asset Type Name", + "qualified_name": "Asset Qualified Name" + }, + attribute_name="classifications", + attribute_values=["P0"] + ) # Update user description for multiple assets result = update_assets_tool( @@ -602,6 +615,7 @@ def update_assets_tool( ) """ try: + # Parse JSON parameters parsed_assets = parse_json_parameter(assets) parsed_attribute_values = parse_list_parameter(attribute_values) diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index e6abf18..7a8f863 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -27,7 +27,7 @@ def update_assets( updatable_assets (Union[UpdatableAsset, List[UpdatableAsset]]): Asset(s) to update. Can be a single UpdatableAsset or a list of UpdatableAssets. attribute_name (UpdatableAttribute): Name of the attribute to update. - Supports userDescription, certificateStatus, readme, and term. + Supports userDescription, certificateStatus, readme, classifications and term. attribute_values (List[Union[str, CertificateStatus, TermOperations]]): List of values to set for the attribute. For certificateStatus, only VERIFIED, DRAFT, or DEPRECATED are allowed. For readme, the value must be a valid Markdown string. @@ -155,6 +155,21 @@ def update_assets( error_msg = f"Error updating terms on asset {updatable_asset.qualified_name}: {str(e)}" logger.error(error_msg) result["errors"].append(error_msg) + elif attribute_name == UpdatableAttribute.CLASSIFICATIONS: + tag_names = attribute_values[index] if isinstance(attribute_values[index], list) else [ + attribute_values[index]] + + try: + client.asset.add_atlan_tags( + asset_type=asset_cls, + qualified_name=asset.qualified_name, + atlan_tag_names=tag_names, + ) + result["updated_count"] += 1 + except Exception as e: + error_msg = f"Error updating tags on asset {updatable_asset.qualified_name}: {str(e)}" + logger.error(error_msg) + result["errors"].append(error_msg) else: # Regular attribute update flow setattr(asset, attribute_name.value, attribute_values[index]) @@ -172,7 +187,6 @@ def update_assets( f"Successfully updated {result['readme_updated']} readme assets: {result['updated_readme_assets']}" ) - # Proces response if len(assets) > 0: response = client.asset.save(assets) result["updated_count"] = len(response.guid_assignments) diff --git a/modelcontextprotocol/tools/models.py b/modelcontextprotocol/tools/models.py index 95bd048..0857178 100644 --- a/modelcontextprotocol/tools/models.py +++ b/modelcontextprotocol/tools/models.py @@ -19,6 +19,8 @@ class UpdatableAttribute(str, Enum): CERTIFICATE_STATUS = "certificate_status" README = "readme" TERM = "term" + CLASSIFICATIONS ="classifications" + class TermOperation(str, Enum): @@ -45,6 +47,7 @@ class UpdatableAsset(BaseModel): type_name: str user_description: Optional[str] = None certificate_status: Optional[CertificateStatus] = None + classifications: Optional[List[str]] = None class Glossary(BaseModel): diff --git a/modelcontextprotocol/uv.lock b/modelcontextprotocol/uv.lock index 2893b99..39637c9 100644 --- a/modelcontextprotocol/uv.lock +++ b/modelcontextprotocol/uv.lock @@ -770,4 +770,4 @@ dependencies = [ sdist = { url = "https://files.pythonhosted.org/packages/4b/4d/938bd85e5bf2edeec766267a5015ad969730bb91e31b44021dfe8b22df6c/uvicorn-0.34.0.tar.gz", hash = "sha256:404051050cd7e905de2c9a7e61790943440b3416f49cb409f965d9dcd0fa73e9", size = 76568, upload-time = "2024-12-15T13:33:30.42Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/61/14/33a3a1352cfa71812a3a21e8c9bfb83f60b0011f5e36f2b1399d51928209/uvicorn-0.34.0-py3-none-any.whl", hash = "sha256:023dc038422502fa28a09c7a30bf2b6991512da7dcdb8fd35fe57cfc154126f4", size = 62315, upload-time = "2024-12-15T13:33:27.467Z" }, -] +] \ No newline at end of file From 0abe56f6cbf39a59ff8da3bc3a2dfa08d42cebe5 Mon Sep 17 00:00:00 2001 From: Bala Date: Tue, 16 Sep 2025 12:14:00 +0530 Subject: [PATCH 2/2] updated pre commit amends --- modelcontextprotocol/server.py | 1 - modelcontextprotocol/tools/assets.py | 7 +++++-- modelcontextprotocol/tools/models.py | 3 +-- modelcontextprotocol/uv.lock | 2 +- 4 files changed, 7 insertions(+), 6 deletions(-) diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 5e89071..5c7c9e6 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -615,7 +615,6 @@ def update_assets_tool( ) """ try: - # Parse JSON parameters parsed_assets = parse_json_parameter(assets) parsed_attribute_values = parse_list_parameter(attribute_values) diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index 7a8f863..663d011 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -156,8 +156,11 @@ def update_assets( logger.error(error_msg) result["errors"].append(error_msg) elif attribute_name == UpdatableAttribute.CLASSIFICATIONS: - tag_names = attribute_values[index] if isinstance(attribute_values[index], list) else [ - attribute_values[index]] + tag_names = ( + attribute_values[index] + if isinstance(attribute_values[index], list) + else [attribute_values[index]] + ) try: client.asset.add_atlan_tags( diff --git a/modelcontextprotocol/tools/models.py b/modelcontextprotocol/tools/models.py index 0857178..3e18cfd 100644 --- a/modelcontextprotocol/tools/models.py +++ b/modelcontextprotocol/tools/models.py @@ -19,8 +19,7 @@ class UpdatableAttribute(str, Enum): CERTIFICATE_STATUS = "certificate_status" README = "readme" TERM = "term" - CLASSIFICATIONS ="classifications" - + CLASSIFICATIONS = "classifications" class TermOperation(str, Enum): diff --git a/modelcontextprotocol/uv.lock b/modelcontextprotocol/uv.lock index 39637c9..2893b99 100644 --- a/modelcontextprotocol/uv.lock +++ b/modelcontextprotocol/uv.lock @@ -770,4 +770,4 @@ dependencies = [ sdist = { url = "https://files.pythonhosted.org/packages/4b/4d/938bd85e5bf2edeec766267a5015ad969730bb91e31b44021dfe8b22df6c/uvicorn-0.34.0.tar.gz", hash = "sha256:404051050cd7e905de2c9a7e61790943440b3416f49cb409f965d9dcd0fa73e9", size = 76568, upload-time = "2024-12-15T13:33:30.42Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/61/14/33a3a1352cfa71812a3a21e8c9bfb83f60b0011f5e36f2b1399d51928209/uvicorn-0.34.0-py3-none-any.whl", hash = "sha256:023dc038422502fa28a09c7a30bf2b6991512da7dcdb8fd35fe57cfc154126f4", size = 62315, upload-time = "2024-12-15T13:33:27.467Z" }, -] \ No newline at end of file +]