Feature/add dq tools

modelcontextprotocol/README.md

@@ -146,6 +146,7 @@ Open `Cursor > Settings > Tools & Integrations > New MCP Server` to include the
 | `create_glossaries` | Create glossaries                                                 |
 | `create_glossary_categories` | Create glossary categories                               |
 | `create_glossary_terms` | Create glossary terms                                         |
+| `create_dq_rules`   | Create data quality rules (column-level, table-level, custom SQL) |
 | `query_asset`       | Execute SQL queries on table/view assets                          |
 
 ## Tool Access Control
@@ -212,13 +213,14 @@ You can restrict any of the following tools:
 - `create_glossaries` - Glossary creation
 - `create_glossary_categories` - Category creation
 - `create_glossary_terms` - Term creation
+- `create_dq_rules_tool` - Data quality rule creation
 
 ### Common Use Cases
 
 #### Read-Only Access
 Restrict all write operations:
 ```
-RESTRICTED_TOOLS=update_assets_tool,create_glossaries,create_glossary_categories,create_glossary_terms
+RESTRICTED_TOOLS=update_assets_tool,create_glossaries,create_glossary_categories,create_glossary_terms,create_dq_rules_tool
 ```
 
 #### Disable DSL Queries
@@ -230,7 +232,7 @@ RESTRICTED_TOOLS=get_assets_by_dsl_tool
 #### Minimal Access
 Allow only basic search:
 ```
-RESTRICTED_TOOLS=get_assets_by_dsl_tool,update_assets_tool,traverse_lineage_tool,create_glossaries,create_glossary_categories,create_glossary_terms
+RESTRICTED_TOOLS=get_assets_by_dsl_tool,update_assets_tool,traverse_lineage_tool,create_glossaries,create_glossary_categories,create_glossary_terms,create_dq_rules_tool
 ```
 
 ### How It Works

modelcontextprotocol/server.py

@@ -12,6 +12,7 @@
     create_glossary_category_assets,
     create_glossary_assets,
     create_glossary_term_assets,
+    create_dq_rules,
     UpdatableAttribute,
     CertificateStatus,
     UpdatableAsset,
@@ -905,6 +906,126 @@ def create_glossary_categories(categories) -> List[Dict[str, Any]]:
     return create_glossary_category_assets(categories)
 
 
+@mcp.tool()
+def create_dq_rules_tool(rules):
+    """
+    Create one or multiple data quality rules in Atlan.
+
+    Supports all rule types: column-level, table-level, and custom SQL rules.
+    Rules can be created individually or in bulk for efficient setup.
+
+    Args:
+        rules (Union[Dict[str, Any], List[Dict[str, Any]]]): Either a single rule
+            specification or a list of rule specifications. Each specification
+            must include:
+            - rule_type (str): Type of rule (see Supported Rule Types)
+            - asset_qualified_name (str): Qualified name of the table/view
+            - Additional fields based on rule type (see examples)
+
+    Returns:
+        Dict[str, Any]: Dictionary containing:
+            - created_count: Number of rules successfully created
+            - created_rules: List of created rules with guid, qualified_name, rule_type
+            - errors: List of any errors encountered
+
+    Examples:
+        # Column-level rules (Null Count, Min/Max Value, Unique/Duplicate Count, etc.)
+        rule = create_dq_rules_tool({
+            "rule_type": "Null Count",  # or "Min Value", "Max Value", "Unique Count", etc.
+            "asset_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE",
+            "column_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE/EMAIL",
+            "threshold_compare_operator": "LESS_THAN_EQUAL",  # EQUAL, GREATER_THAN, etc.
+            "threshold_value": 5,
+            "alert_priority": "URGENT",  # LOW, NORMAL, URGENT
+            "row_scope_filtering_enabled": True,
+            "description": "Email column should have minimal nulls"
+        })
+
+        # Conditional rules (String Length, Regex, Valid Values)
+        rule = create_dq_rules_tool({
+            "rule_type": "String Length",  # or "Regex", "Valid Values"
+            "asset_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE",
+            "column_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE/PHONE",
+            "threshold_value": 10,
+            "alert_priority": "URGENT",
+            "rule_conditions": [{
+                "type": "STRING_LENGTH_BETWEEN",  # See Rule Condition Types below
+                "min_value": 10,
+                "max_value": 15
+            }],
+            # For Regex: {"type": "REGEX_NOT_MATCH", "value": "pattern"}
+            # For Valid Values: {"type": "IN_LIST", "value": ["ACTIVE", "INACTIVE"]}
+            "row_scope_filtering_enabled": True
+        })
+
+        # Table-level (Row Count) and Time-based (Freshness)
+        rule = create_dq_rules_tool({
+            "rule_type": "Row Count",  # No column_qualified_name needed
+            "asset_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE",
+            "threshold_compare_operator": "GREATER_THAN_EQUAL",
+            "threshold_value": 1000,
+            "alert_priority": "URGENT"
+        })
+        # For Freshness: Add "column_qualified_name" + "threshold_unit": "DAYS"/"HOURS"/"MINUTES"
+
+        # Custom SQL rule
+        rule = create_dq_rules_tool({
+            "rule_type": "Custom SQL",
+            "asset_qualified_name": "default/snowflake/123/DB/SCHEMA/TABLE",
+            "rule_name": "Revenue Consistency Check",
+            "custom_sql": "SELECT COUNT(*) FROM TABLE WHERE revenue < 0 OR revenue > 1000000",
+            "threshold_compare_operator": "EQUAL",
+            "threshold_value": 0,
+            "alert_priority": "URGENT",
+            "dimension": "CONSISTENCY",  # See Data Quality Dimensions below
+            "description": "Ensure revenue values are within expected range"
+        })
+
+        # Bulk creation - Pass array instead of single dict
+        rules = create_dq_rules_tool([
+            {"rule_type": "Null Count", "column_qualified_name": "...EMAIL", ...},
+            {"rule_type": "Duplicate Count", "column_qualified_name": "...USER_ID", ...},
+            {"rule_type": "Row Count", "asset_qualified_name": "...", ...}
+        ])
+
+    Supported Rule Types:
+        Completeness: "Null Count", "Null Percentage", "Blank Count", "Blank Percentage"
+        Statistical: "Min Value", "Max Value", "Average", "Standard Deviation"
+        Uniqueness: "Unique Count", "Duplicate Count"
+        Validity: "Regex", "String Length", "Valid Values"
+        Timeliness: "Freshness"
+        Volume: "Row Count"
+        Custom: "Custom SQL"
+
+    Valid Alert Priority Levels:
+        "LOW", "NORMAL" (default), "URGENT"
+
+    Threshold Operators:
+        "EQUAL", "GREATER_THAN", "GREATER_THAN_EQUAL", "LESS_THAN", "LESS_THAN_EQUAL", "BETWEEN"
+
+    Threshold Units (Freshness only):
+        "DAYS", "HOURS", "MINUTES"
+
+    Data Quality Dimensions (Custom SQL only):
+        "COMPLETENESS", "VALIDITY", "UNIQUENESS", "TIMELINESS", "VOLUME", "ACCURACY", "CONSISTENCY"
+
+    Rule Condition Types:
+        String Length: "STRING_LENGTH_EQUALS", "STRING_LENGTH_BETWEEN",
+                      "STRING_LENGTH_GREATER_THAN", "STRING_LENGTH_LESS_THAN"
+        Regex: "REGEX_MATCH", "REGEX_NOT_MATCH"
+        Valid Values: "IN_LIST", "NOT_IN_LIST"
+    """
+    try:
+        parsed_rules = parse_json_parameter(rules)
+        return create_dq_rules(parsed_rules)
+    except (json.JSONDecodeError, ValueError) as e:
+        return {
+            "created_count": 0,
+            "created_rules": [],
+            "errors": [f"Parameter parsing error: {str(e)}"],
+        }
+
+
 def main():
     """Main entry point for the Atlan MCP Server."""
 

modelcontextprotocol/tools/__init__.py

@@ -3,6 +3,7 @@
 from .lineage import traverse_lineage
 from .assets import update_assets
 from .query import query_asset
+from .dq_rules import create_dq_rules
 from .glossary import (
     create_glossary_category_assets,
     create_glossary_assets,
@@ -16,6 +17,8 @@
     Glossary,
     GlossaryCategory,
     GlossaryTerm,
+    DQRuleType,
+    DQRuleSpecification,
 )
 
 __all__ = [
@@ -34,4 +37,7 @@
     "Glossary",
     "GlossaryCategory",
     "GlossaryTerm",
+    "create_dq_rules",
+    "DQRuleType",
+    "DQRuleSpecification",
 ]