parent d122674

author Bala <[email protected]> 1757908811 +0530
committer Bala <[email protected]> 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

Author: bladata1990

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)

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)

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):