diff --git a/modelcontextprotocol/server.py b/modelcontextprotocol/server.py index 515adfd..1aef447 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 f9d18f6..08092dd 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 e6abf18..fb2e7d6 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,8 +161,19 @@ 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_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) diff --git a/modelcontextprotocol/tools/models.py b/modelcontextprotocol/tools/models.py index 95bd048..de7ffd7 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