From 666993a875261647139f083189828ed6b5b6113b Mon Sep 17 00:00:00 2001 From: mustafahasankhan Date: Tue, 16 Sep 2025 16:41:29 +0530 Subject: [PATCH 1/4] feat: add announcements --- modelcontextprotocol/server.py | 39 ++++++++++++ modelcontextprotocol/tools/__init__.py | 4 ++ modelcontextprotocol/tools/assets.py | 87 ++++++++++++++++++++++++++ modelcontextprotocol/tools/models.py | 17 +++++ 4 files changed, 147 insertions(+) diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 515adfdb..1aef447d 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -600,6 +600,30 @@ def update_assets_tool( "term_guids": ["term-guid-to-remove"] }] ) + + # Add warning announcement to an asset + update_assets_tool( + assets={ + "guid": "abc-123", + "name": "sales_data", + "type_name": "Table", + "qualified_name": "default/snowflake/db/schema/sales_data" + }, + attribute_name="announcement", + attribute_values=[{ + "announcement_type": "WARNING", + "announcement_title": "Data Quality Issue", + "announcement_message": "Missing records for Q4 2024. ETL team investigating." + }] + ) + + # Remove announcement + update_assets_tool( + assets={"guid": "abc-123", ...}, + attribute_name="announcement", + attribute_values=[None] # or [{}] + ) + """ try: # Parse JSON parameters @@ -627,6 +651,21 @@ def update_assets_tool( CertificateStatus(val) for val in parsed_attribute_values ] + elif attr_enum == UpdatableAttribute.ANNOUNCEMENT: + # Validate announcement structure + for val in parsed_attribute_values: + if val and not isinstance(val, dict): + raise ValueError(f"Announcement must be a dict, got {type(val)}") + if val and not all( + k in val + for k in [ + "announcement_type", + "announcement_title", + "announcement_message", + ] + ): + raise ValueError("Announcement must have type, title, and message") + # Convert assets to UpdatableAsset objects if isinstance(parsed_assets, dict): updatable_assets = [UpdatableAsset(**parsed_assets)] diff --git a/modelcontextprotocol/tools/__init__.py b/modelcontextprotocol/tools/__init__.py index f9d18f67..08092dd8 100644 --- a/modelcontextprotocol/tools/__init__.py +++ b/modelcontextprotocol/tools/__init__.py @@ -16,6 +16,8 @@ Glossary, GlossaryCategory, GlossaryTerm, + Announcement, + AnnouncementType, ) __all__ = [ @@ -34,4 +36,6 @@ "Glossary", "GlossaryCategory", "GlossaryTerm", + "AnnouncementType", + "Announcement", ] diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index e6abf183..95f1ce05 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -9,6 +9,8 @@ TermOperations, ) from pyatlan.model.assets import Readme, AtlasGlossaryTerm +from pyatlan.model.core import Announcement as AtlanAnnouncement +from pyatlan.model.enums import AnnouncementType as AtlanAnnouncementType from pyatlan.model.fluent_search import CompoundQuery, FluentSearch # Initialize logging @@ -32,6 +34,10 @@ def update_assets( 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 TermOperations object with operation and term_guids. + For announcement, each value should be a dict with: + - announcement_type: "INFORMATION", "WARNING", or "ISSUE" + - announcement_title: Title of the announcement + - announcement_message: Message content Returns: Dict[str, Any]: Dictionary containing: @@ -155,6 +161,87 @@ 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.ANNOUNCEMENT: + announcement_results = [] + + for index, updatable_asset in enumerate(updatable_assets): + try: + announcement_data = attribute_values[index] + + # Handle announcement removal (None or empty dict) + if not announcement_data: + # Remove announcement + asset_cls = getattr( + __import__( + "pyatlan.model.assets", + fromlist=[updatable_asset.type_name], + ), + updatable_asset.type_name, + ) + + updated_asset = client.asset.remove_announcement( + asset_type=asset_cls, + qualified_name=updatable_asset.qualified_name, + name=updatable_asset.name, + ) + + if updated_asset: + announcement_results.append( + { + "asset": updatable_asset.qualified_name, + "action": "removed", + } + ) + continue + + # Create or update announcement + announcement = AtlanAnnouncement( + announcement_type=AtlanAnnouncementType[ + announcement_data["announcement_type"] + ], + announcement_title=announcement_data["announcement_title"], + announcement_message=announcement_data[ + "announcement_message" + ], + ) + + # Get the asset class + asset_cls = getattr( + __import__( + "pyatlan.model.assets", + fromlist=[updatable_asset.type_name], + ), + updatable_asset.type_name, + ) + + # Update announcement + updated_asset = client.asset.update_announcement( + asset_type=asset_cls, + qualified_name=updatable_asset.qualified_name, + name=updatable_asset.name, + announcement=announcement, + ) + + if updated_asset: + announcement_results.append( + { + "asset": updatable_asset.qualified_name, + "type": announcement_data["announcement_type"], + "title": announcement_data["announcement_title"], + } + ) + + except Exception as e: + logger.error( + f"Error updating announcement for {updatable_asset.qualified_name}: {e}" + ) + result["errors"].append( + f"Failed to update {updatable_asset.qualified_name}: {str(e)}" + ) + + result["updated_count"] = len(announcement_results) + result["announcements"] = announcement_results + else: # Regular attribute update flow setattr(asset, attribute_name.value, attribute_values[index]) diff --git a/modelcontextprotocol/tools/models.py b/modelcontextprotocol/tools/models.py index 95bd048a..de7ffd76 100644 --- a/modelcontextprotocol/tools/models.py +++ b/modelcontextprotocol/tools/models.py @@ -19,6 +19,7 @@ class UpdatableAttribute(str, Enum): CERTIFICATE_STATUS = "certificate_status" README = "readme" TERM = "term" + ANNOUNCEMENT = "announcement" class TermOperation(str, Enum): @@ -73,3 +74,19 @@ class GlossaryTerm(BaseModel): user_description: Optional[str] = None certificate_status: Optional[CertificateStatus] = None category_guids: Optional[List[str]] = None + + +class AnnouncementType(str, Enum): + """Enum for announcement types.""" + + INFORMATION = "INFORMATION" + WARNING = "WARNING" + ISSUE = "ISSUE" + + +class Announcement(BaseModel): + """Class representing an announcement.""" + + announcement_type: AnnouncementType + announcement_title: str + announcement_message: str From 4804dd1d4f89903bccd36b5d98a244cb7ff01466 Mon Sep 17 00:00:00 2001 From: mustafahasankhan Date: Tue, 16 Sep 2025 19:14:34 +0530 Subject: [PATCH 2/4] chore: address review --- modelcontextprotocol/tools/assets.py | 92 ++++------------------------ 1 file changed, 11 insertions(+), 81 deletions(-) diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index 95f1ce05..fb2e7d65 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -162,88 +162,18 @@ def update_assets( logger.error(error_msg) result["errors"].append(error_msg) elif attribute_name == UpdatableAttribute.ANNOUNCEMENT: - announcement_results = [] - - for index, updatable_asset in enumerate(updatable_assets): - try: - announcement_data = attribute_values[index] - - # Handle announcement removal (None or empty dict) - if not announcement_data: - # Remove announcement - asset_cls = getattr( - __import__( - "pyatlan.model.assets", - fromlist=[updatable_asset.type_name], - ), - updatable_asset.type_name, - ) - - updated_asset = client.asset.remove_announcement( - asset_type=asset_cls, - qualified_name=updatable_asset.qualified_name, - name=updatable_asset.name, - ) - - if updated_asset: - announcement_results.append( - { - "asset": updatable_asset.qualified_name, - "action": "removed", - } - ) - continue - - # Create or update announcement - announcement = AtlanAnnouncement( - announcement_type=AtlanAnnouncementType[ - announcement_data["announcement_type"] - ], - announcement_title=announcement_data["announcement_title"], - announcement_message=announcement_data[ - "announcement_message" - ], - ) - - # Get the asset class - asset_cls = getattr( - __import__( - "pyatlan.model.assets", - fromlist=[updatable_asset.type_name], - ), - updatable_asset.type_name, - ) - - # Update announcement - updated_asset = client.asset.update_announcement( - asset_type=asset_cls, - qualified_name=updatable_asset.qualified_name, - name=updatable_asset.name, - announcement=announcement, - ) - - if updated_asset: - announcement_results.append( - { - "asset": updatable_asset.qualified_name, - "type": announcement_data["announcement_type"], - "title": announcement_data["announcement_title"], - } - ) - - except Exception as e: - logger.error( - f"Error updating announcement for {updatable_asset.qualified_name}: {e}" - ) - result["errors"].append( - f"Failed to update {updatable_asset.qualified_name}: {str(e)}" - ) - - result["updated_count"] = len(announcement_results) - result["announcements"] = announcement_results - + announcement_data = attribute_values[index] + if not announcement_data: # None or empty dict + asset.remove_announcement() + else: + announcement = AtlanAnnouncement( + announcement_type=AtlanAnnouncementType[announcement_data["announcement_type"]], + announcement_title=announcement_data["announcement_title"], + announcement_message=announcement_data["announcement_message"] + ) + asset.set_announcement(announcement) + assets.append(asset) else: - # Regular attribute update flow setattr(asset, attribute_name.value, attribute_values[index]) assets.append(asset) From 6f959360eaa1531dec4fdd280d1aaac8ad18e652 Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Sat, 8 Nov 2025 13:35:09 +0530 Subject: [PATCH 3/4] fix : direct update --- modelcontextprotocol/server.py | 2 +- modelcontextprotocol/tools/assets.py | 12 ++++-------- 2 files changed, 5 insertions(+), 9 deletions(-) diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 6d13472e..67d418b6 100644 --- a/modelcontextprotocol/server.py +++ b/modelcontextprotocol/server.py @@ -480,7 +480,7 @@ def update_assets_tool( Can be a single UpdatableAsset or a list of UpdatableAsset objects. For asset of type_name=AtlasGlossaryTerm or type_name=AtlasGlossaryCategory, each asset dictionary MUST include a "glossary_guid" key which is the GUID of the glossary that the term belongs to. attribute_name (str): Name of the attribute to update. - Supports "user_description", "certificate_status", "readme", and "term". + Supports "user_description", "certificate_status", "readme", "term", and "announcement". 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. diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index a85251d0..aa6aafaa 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -9,7 +9,6 @@ TermOperations, ) from pyatlan.model.assets import Readme, AtlasGlossaryTerm, AtlasGlossaryCategory -from pyatlan.model.core import Announcement as AtlanAnnouncement from pyatlan.model.enums import AnnouncementType as AtlanAnnouncementType from pyatlan.model.fluent_search import CompoundQuery, FluentSearch @@ -30,7 +29,7 @@ def update_assets( Can be a single UpdatableAsset or a list of UpdatableAssets. For asset of type_name=AtlasGlossaryTerm or type_name=AtlasGlossaryCategory, each asset dictionary MUST include a "glossary_guid" key which is the GUID of the glossary that the term belongs to. attribute_name (UpdatableAttribute): Name of the attribute to update. - Supports userDescription, certificateStatus, readme, and term. + Supports userDescription, certificateStatus, readme, term, and announcement. 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. @@ -179,12 +178,9 @@ def update_assets( if not announcement_data: # None or empty dict asset.remove_announcement() else: - announcement = AtlanAnnouncement( - announcement_type=AtlanAnnouncementType[announcement_data["announcement_type"]], - announcement_title=announcement_data["announcement_title"], - announcement_message=announcement_data["announcement_message"] - ) - asset.set_announcement(announcement) + asset.announcement_type = AtlanAnnouncementType[announcement_data["announcement_type"]] + asset.announcement_title = announcement_data["announcement_title"] + asset.announcement_message = announcement_data["announcement_message"] assets.append(asset) else: setattr(asset, attribute_name.value, attribute_values[index]) From a38b435f02ee799a8537e5404fc46521c6919a6d Mon Sep 17 00:00:00 2001 From: Abhinav Mathur Date: Mon, 10 Nov 2025 11:11:27 +0530 Subject: [PATCH 4/4] chore: fix formatting --- modelcontextprotocol/tools/assets.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/modelcontextprotocol/tools/assets.py b/modelcontextprotocol/tools/assets.py index aa6aafaa..b70cacad 100644 --- a/modelcontextprotocol/tools/assets.py +++ b/modelcontextprotocol/tools/assets.py @@ -178,9 +178,13 @@ def update_assets( if not announcement_data: # None or empty dict asset.remove_announcement() else: - asset.announcement_type = AtlanAnnouncementType[announcement_data["announcement_type"]] + asset.announcement_type = AtlanAnnouncementType[ + announcement_data["announcement_type"] + ] asset.announcement_title = announcement_data["announcement_title"] - asset.announcement_message = announcement_data["announcement_message"] + asset.announcement_message = announcement_data[ + "announcement_message" + ] assets.append(asset) else: setattr(asset, attribute_name.value, attribute_values[index])