datahub-agent-context 1.3.1.8__py3-none-any.whl

datahub_agent_context/mcp_tools/search.py

@@ -0,0 +1,239 @@
+"""Search tools for DataHub."""
+
+import json
+import logging
+import pathlib
+from typing import Any, Dict, Literal, Optional
+
+from datahub.sdk.search_client import compile_filters
+from datahub.sdk.search_filters import Filter, load_filters
+from datahub_agent_context.context import get_graph
+from datahub_agent_context.mcp_tools.base import (
+    clean_gql_response,
+    execute_graphql,
+    fetch_global_default_view,
+)
+
+logger = logging.getLogger(__name__)
+
+# Load GraphQL queries
+_gql_dir = pathlib.Path(__file__).parent / "gql"
+search_gql = (_gql_dir / "search.gql").read_text()
+
+
+def _convert_custom_filter_format(filters_obj: Any) -> Any:
+    """
+    Convert chatbot's intuitive {"custom": {...}} format to the format expected by _CustomCondition.
+
+    Transforms:
+    {"custom": {"field": "urn", "condition": "EQUAL", "values": [...]}}
+
+    Into:
+    {"field": "urn", "condition": "EQUAL", "values": [...]}
+
+    This allows the discriminator to correctly identify it as _custom.
+    """
+    if isinstance(filters_obj, dict):
+        # Check if this is a "custom" or "custom_condition" wrapper that needs unwrapping
+        if len(filters_obj) == 1 and (
+            "custom" in filters_obj or "custom_condition" in filters_obj
+        ):
+            wrapper_key = "custom" if "custom" in filters_obj else "custom_condition"
+            custom_content = filters_obj[wrapper_key]
+            # Ensure it has the expected structure for _CustomCondition
+            if isinstance(custom_content, dict) and "field" in custom_content:
+                return custom_content
+
+        # Recursively process nested filters (for "and", "or", etc.)
+        result = {}
+        for key, value in filters_obj.items():
+            if isinstance(value, (list, dict)):
+                result[key] = _convert_custom_filter_format(value)
+            else:
+                result[key] = value
+        return result
+    elif isinstance(filters_obj, list):
+        # Process list of filters
+        return [_convert_custom_filter_format(item) for item in filters_obj]
+    else:
+        # Return primitive values unchanged
+        return filters_obj
+
+
+def search(
+    query: str = "*",
+    filters: Optional[Filter | str] = None,
+    num_results: int = 10,
+    sort_by: Optional[str] = None,
+    sort_order: Optional[Literal["asc", "desc"]] = "desc",
+    offset: int = 0,
+) -> dict:
+    """Search across DataHub entities using structured full-text search.
+    Results are ordered by relevance and importance - examine top results first.
+
+    SEARCH SYNTAX:
+    - Structured full-text search - **always start queries with /q**
+    - **Recommended: Use + operator for AND** (handles punctuation better than quotes)
+    - Supports full boolean logic: AND (default), OR, NOT, parentheses, field searches
+    - Examples:
+      • /q user+transaction → requires both terms (better for field names with _ or punctuation)
+      • /q point+sale+app → requires all terms (works with point_of_sale_app_usage)
+      • /q wizard OR pet → entities containing either term
+      • /q revenue* → wildcard matching (revenue_2023, revenue_2024, revenue_monthly, etc.)
+      • /q tag:PII → search by tag name
+      • /q "exact table name" → exact phrase matching (use sparingly)
+      • /q (sales OR revenue) AND quarterly → complex boolean combinations
+    - Fast and precise for exact matching, technical terms, and complex queries
+    - Best for: entity names, identifiers, column names, or any search needing boolean logic
+
+    PAGINATION:
+    - num_results: Number of results to return per page (max: 50)
+    - offset: Starting position in results (default: 0)
+    - Examples:
+      • First page: offset=0, num_results=10
+      • Second page: offset=10, num_results=10
+      • Third page: offset=20, num_results=10
+
+    FACET EXPLORATION - Discover metadata without returning results:
+    - Set num_results=0 to get ONLY facets (no search results)
+    - Facets show ALL tags, glossaryTerms, platforms, domains used in the catalog
+    - Example: search(query="*", filters={"entity_type": ["DATASET"]}, num_results=0)
+      → Returns facets showing all tags/glossaryTerms applied to datasets
+    - Use this to discover what metadata exists before doing filtered searches
+
+    TYPICAL WORKFLOW:
+    1. Facet exploration: search(query="*", filters={"entity_type": ["DATASET"]}, num_results=0)
+       → Examine tags/glossaryTerms facets to see what metadata exists
+    2. Filtered search: search(query="*", filters={"tag": ["urn:li:tag:pii"]}, num_results=30)
+       → Get entities with specific tag using URN from step 1
+    3. Get details: Use get_entities() on specific results
+
+    Here are some example filters:
+    - All Looker assets
+    ```
+    {"platform": ["looker"]}
+    ```
+    - Production environment warehouse assets
+    ```
+    {
+      "and": [
+        {"env": ["PROD"]},
+        {"platform": ["snowflake", "bigquery", "redshift"]}
+      ]
+    }
+    ```
+    - Filter by domain (MUST use full URN format)
+    ```
+    {"domain": ["urn:li:domain:marketing"]}
+    {"domain": ["urn:li:domain:9f8e7d6c-5b4a-3928-1765-432109876543", "urn:li:domain:7c6b5a49-3827-1654-9032-8f7e6d5c4b3a"]}
+    ```
+    IMPORTANT: Domain filters require full URN format starting with "urn:li:domain:",
+    NOT short names like "marketing" or "customer". Domain URNs can be readable names
+    or GUIDs. Always search with {"entity_type": ["domain"]}
+    to find valid domain URNs first, then use the exact URN from the results.
+
+    SUPPORTED FILTER TYPES (only these will work):
+    - entity_type: ["dataset"], ["dashboard", "chart"], ["corp_user"], ["corp_group"]
+    - entity_subtype: ["Table"], ["View", "Model"]
+    - platform: ["snowflake"], ["looker", "tableau"]
+    - domain: ["urn:li:domain:marketing"] (full URN required)
+    - container: ["urn:li:container:..."] (full URN required)
+    - tag: ["urn:li:tag:PII"] (full tag URN required)
+    - glossary_term: ["urn:li:glossaryTerm:uuid"] (full term URN required)
+    - owner: ["urn:li:corpuser:alice", "urn:li:corpGroup:marketing"] (full user or group URN required)
+    - custom: {"field": "fieldName", "condition": "EQUAL", "values": [...]}
+    - status: ["NOT_SOFT_DELETED"] (for non-deleted entities)
+    - env: ["PROD"], ["DEV", "STAGING"] (Should not use unless explicitly requested)
+    - and: [filter1, filter2] (combines multiple filters)
+    - or: [filter1, filter2] (matches any filter)
+    - not: {"entity_type": ["dataset"]} (excludes matches)
+
+    CRITICAL: Use only ONE discriminator key per filter object. Never mix
+    entity_type with custom, domain, etc. at the same level. Use "and" or "or" to combine.
+
+    SEARCH STRATEGY EXAMPLES:
+    - /q customer+behavior → finds tables with both terms (works with customer_behavior fields)
+    - /q customer OR user → finds tables with either term
+    - /q (financial OR revenue) AND metrics → complex boolean logic
+
+    SORTING - Order results by specific fields:
+    - sort_by: Field name to sort by (optional)
+    - sort_order: "desc" (default) or "asc"
+
+    Note: If sort_by is not provided, search results use default ranking by relevance and
+    importance. When using sort_by, results are strictly ordered by that field.
+
+    Args:
+        query: Search query string (use /q prefix for structured queries)
+        filters: Optional filter object or JSON string
+        num_results: Number of results to return (max 50)
+        sort_by: Optional field name to sort by
+        sort_order: Sort order ("asc" or "desc")
+        offset: Starting position for pagination
+
+    Returns:
+        Dictionary with search results, facets, and metadata
+
+    Example:
+        from datahub_agent_context.context import DataHubContext
+
+        with DataHubContext(client.graph):
+            result = search(query="/q users", filters={"entity_type": ["dataset"]})
+    """
+    graph = get_graph()
+    # Cap num_results at 50 to prevent excessive requests
+    num_results = min(num_results, 50)
+
+    # Handle stringified JSON filters or dict filters
+    if isinstance(filters, str):
+        # Parse JSON first to allow preprocessing
+        filters_dict = json.loads(filters)
+
+        # Convert "custom" wrapper to direct _custom format for compatibility
+        filters_dict = _convert_custom_filter_format(filters_dict)
+
+        filters = load_filters(filters_dict)
+    elif isinstance(filters, dict):
+        # Convert dict to Filter object
+        filters_dict = _convert_custom_filter_format(filters)
+        filters = load_filters(filters_dict)
+
+    types, compiled_filters = compile_filters(filters)
+
+    # Fetch and apply default view (returns None if disabled or not configured)
+    view_urn = fetch_global_default_view(graph)
+    if view_urn:
+        logger.debug(f"Applying default view: {view_urn}")
+    else:
+        logger.debug("No default view to apply")
+
+    variables: Dict[str, Any] = {
+        "query": query,
+        "types": types,
+        "orFilters": compiled_filters,
+        "count": max(num_results, 1),  # 0 is not a valid value for count.
+        "start": offset,
+        "viewUrn": view_urn,  # Will be None if disabled or not set
+    }
+
+    # Add sorting if requested
+    if sort_by is not None:
+        sort_order_enum = "ASCENDING" if sort_order == "asc" else "DESCENDING"
+        variables["sortInput"] = {
+            "sortCriteria": [{"field": sort_by, "sortOrder": sort_order_enum}]
+        }
+
+    # Use keyword search
+    response = execute_graphql(
+        graph,
+        query=search_gql,
+        variables=variables,
+        operation_name="search",
+    )["searchAcrossEntities"]
+
+    if num_results == 0 and isinstance(response, dict):
+        # Hack to support num_results=0 without support for it in the backend.
+        response.pop("searchResults", None)
+        response.pop("count", None)
+
+    return clean_gql_response(response)

datahub_agent_context/mcp_tools/structured_properties.py

@@ -0,0 +1,447 @@
+"""Structured property management tools for DataHub MCP server."""
+
+import logging
+from datetime import datetime
+from typing import Dict, List, Union
+
+from datahub.utilities.urns._urn_base import Urn
+from datahub_agent_context.context import get_graph
+from datahub_agent_context.mcp_tools.base import execute_graphql
+
+logger = logging.getLogger(__name__)
+
+
+def _validate_and_fetch_structured_property(property_urn: str) -> Dict:
+    """
+    Validate that the structured property exists and fetch its definition.
+
+    Returns:
+        Dictionary with property definition including valueType and entityTypes
+
+    Raises:
+        ValueError: If the property URN does not exist or is invalid
+    """
+    graph = get_graph()
+    query = """
+        query getStructuredProperty($urn: String!) {
+            entity(urn: $urn) {
+                urn
+                type
+                ... on StructuredPropertyEntity {
+                    definition {
+                        qualifiedName
+                        entityTypes {
+                            urn
+                            type
+                            info {
+                                type
+                            }
+                        }
+                        valueType {
+                            urn
+                            info {
+                                qualifiedName
+                            }
+                        }
+                        cardinality
+                    }
+                }
+            }
+        }
+    """
+
+    try:
+        result = execute_graphql(
+            graph,
+            query=query,
+            variables={"urn": property_urn},
+            operation_name="getStructuredProperty",
+        )
+
+        entity = result.get("entity")
+
+        if entity is None:
+            raise ValueError(
+                f"Structured property URN does not exist in DataHub: {property_urn}. "
+                f"Please use the search tool to find existing structured properties, "
+                f"or create the property first before assigning it."
+            )
+
+        if entity.get("type") != "STRUCTURED_PROPERTY":
+            raise ValueError(
+                f"The URN is not a structured property entity: {property_urn} (type: {entity.get('type')})"
+            )
+
+        return entity.get("definition", {})
+
+    except Exception as e:
+        if isinstance(e, ValueError):
+            raise
+        raise ValueError(f"Failed to validate structured property URN: {str(e)}") from e
+
+
+def _validate_property_value(
+    property_definition: Dict, value: Union[str, float, int]
+) -> Dict:
+    """
+    Validate and convert a property value to the appropriate GraphQL format.
+
+    Supports 5 data types:
+    - datahub.string: Plain text strings
+    - datahub.number: Numeric values (int, float, double, long)
+    - datahub.urn: DataHub URN references
+    - datahub.date: ISO 8601 date strings
+    - datahub.rich_text: Rich text/markdown content
+
+    Args:
+        property_definition: The property definition containing valueType info
+        value: The value to validate and convert
+
+    Returns:
+        Dictionary with either stringValue or numberValue key
+
+    Raises:
+        ValueError: If the value type doesn't match the property's valueType
+    """
+    value_type_info = property_definition.get("valueType", {}).get("info", {})
+    qualified_name = value_type_info.get("qualifiedName", "").lower()
+
+    # Determine the data type
+    is_numeric_type = any(
+        numeric_type in qualified_name
+        for numeric_type in ["number", "int", "float", "double", "long"]
+    )
+    is_urn_type = "urn" in qualified_name and "datahub.urn" in qualified_name
+    is_date_type = "date" in qualified_name
+    is_rich_text_type = "rich_text" in qualified_name or "richtext" in qualified_name
+
+    if is_numeric_type:
+        # Value should be numeric
+        if isinstance(value, (int, float)):
+            return {"numberValue": float(value)}
+        elif isinstance(value, str):
+            try:
+                return {"numberValue": float(value)}
+            except ValueError as e:
+                raise ValueError(
+                    f"Property expects numeric type ({qualified_name}), but got non-numeric string: {value}"
+                ) from e
+        else:
+            raise ValueError(
+                f"Property expects numeric type ({qualified_name}), got {type(value).__name__}"
+            )
+
+    elif is_urn_type:
+        # Value should be a valid DataHub URN
+        if not isinstance(value, str):
+            value = str(value)
+
+        try:
+            # Validate URN format
+            Urn.from_string(value)
+            return {"stringValue": value}
+        except Exception as e:
+            raise ValueError(
+                f"Property expects URN type ({qualified_name}), but got invalid URN: {value}. "
+                f"URNs must be in format 'urn:li:entityType:...' Error: {str(e)}"
+            ) from e
+
+    elif is_date_type:
+        # Value should be an ISO 8601 date string
+        if not isinstance(value, str):
+            value = str(value)
+
+        # Try to parse as ISO 8601 date
+        try:
+            # Support various ISO 8601 formats
+            # Examples: 2024-12-22, 2024-12-22T10:30:00, 2024-12-22T10:30:00Z, 2024-12-22T10:30:00+00:00
+            datetime.fromisoformat(value.replace("Z", "+00:00"))
+            return {"stringValue": value}
+        except ValueError as e:
+            raise ValueError(
+                f"Property expects date type ({qualified_name}), but got invalid date format: {value}. "
+                f"Dates must be in ISO 8601 format (e.g., '2024-12-22', '2024-12-22T10:30:00Z')"
+            ) from e
+
+    elif is_rich_text_type:
+        # Value should be string (can contain markdown/HTML)
+        if isinstance(value, str):
+            return {"stringValue": value}
+        else:
+            # Convert to string for non-string types
+            return {"stringValue": str(value)}
+
+    else:
+        # Default to string type (datahub.string or unknown types)
+        if isinstance(value, str):
+            return {"stringValue": value}
+        else:
+            # Convert to string for non-string types
+            return {"stringValue": str(value)}
+
+    raise ValueError(
+        f"Value type mismatch: property expects {qualified_name}, got {type(value).__name__}"
+    )
+
+
+def add_structured_properties(
+    property_values: Dict[str, List[Union[str, float, int]]],
+    entity_urns: List[str],
+) -> dict:
+    """Add structured properties with values to multiple DataHub entities.
+
+    This tool allows you to assign structured properties to multiple entities in a single operation.
+    Structured properties are schema-defined metadata fields that can store typed values (strings, numbers, etc.).
+
+    Args:
+        property_values: Dictionary mapping structured property URNs to lists of values.
+                        Example: {
+                            "urn:li:structuredProperty:io.acryl.privacy.retentionTime": ["90"],
+                            "urn:li:structuredProperty:io.acryl.common.businessCriticality": ["HIGH"]
+                        }
+        entity_urns: List of entity URNs to assign properties to (e.g., dataset URNs, dashboard URNs)
+
+    Examples:
+        # Add retention time and criticality to datasets
+        add_structured_properties(
+            property_values={
+                "urn:li:structuredProperty:io.acryl.privacy.retentionTime": ["90"],
+                "urn:li:structuredProperty:io.acryl.common.businessCriticality": ["HIGH"]
+            },
+            entity_urns=[
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.users,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.customers,PROD)"
+            ]
+        )
+
+        # Add numeric property
+        add_structured_properties(
+            property_values={
+                "urn:li:structuredProperty:io.acryl.dataQuality.scoreThreshold": [0.95]
+            },
+            entity_urns=[
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.verified_data,PROD)"
+            ]
+        )
+
+        # Add multiple values for a multi-valued property
+        add_structured_properties(
+            property_values={
+                "urn:li:structuredProperty:io.acryl.common.dataClassification": ["PII", "SENSITIVE"]
+            },
+            entity_urns=[
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.users,PROD)"
+            ]
+        )
+
+    Example:
+        from datahub_agent_context.context import DataHubContext
+
+        with DataHubContext(client.graph):
+            result = add_structured_properties(
+                property_values={"urn:li:structuredProperty:...": ["value"]},
+                entity_urns=["urn:li:dataset:(...)"]
+            )
+    """
+    graph = get_graph()
+    if not property_values:
+        raise ValueError("property_values cannot be empty")
+    if not entity_urns:
+        raise ValueError("entity_urns cannot be empty")
+
+    # Validate all structured properties and fetch their definitions
+    property_definitions = {}
+    for property_urn in property_values:
+        property_definitions[property_urn] = _validate_and_fetch_structured_property(
+            property_urn
+        )
+
+    # Build structured property input params with type validation
+    structured_property_params = []
+    for property_urn, values in property_values.items():
+        property_def = property_definitions[property_urn]
+
+        # Validate and convert each value
+        converted_values = []
+        for value in values:
+            try:
+                converted_value = _validate_property_value(property_def, value)
+                converted_values.append(converted_value)
+            except ValueError as e:
+                raise ValueError(
+                    f"Value validation failed for {property_urn}: {str(e)}"
+                ) from e
+
+        structured_property_params.append(
+            {"structuredPropertyUrn": property_urn, "values": converted_values}
+        )
+
+    # Execute upsert for each entity
+    mutation = """
+        mutation upsertStructuredProperties($input: UpsertStructuredPropertiesInput!) {
+            upsertStructuredProperties(input: $input) {
+                properties {
+                    structuredProperty {
+                        urn
+                    }
+                }
+            }
+        }
+    """
+
+    success_count = 0
+    failed_urns = []
+    error_messages = []
+
+    for entity_urn in entity_urns:
+        variables = {
+            "input": {
+                "assetUrn": entity_urn,
+                "structuredPropertyInputParams": structured_property_params,
+            }
+        }
+
+        try:
+            result = execute_graphql(
+                graph,
+                query=mutation,
+                variables=variables,
+                operation_name="upsertStructuredProperties",
+            )
+
+            if result.get("upsertStructuredProperties"):
+                success_count += 1
+            else:
+                failed_urns.append(entity_urn)
+                error_messages.append(
+                    f"{entity_urn}: operation returned false or empty result"
+                )
+
+        except Exception as e:
+            failed_urns.append(entity_urn)
+            error_messages.append(f"{entity_urn}: {str(e)}")
+
+    if failed_urns:
+        error_details = "; ".join(error_messages[:3])
+        if len(error_messages) > 3:
+            error_details += f"; and {len(error_messages) - 3} more error(s)"
+        raise RuntimeError(
+            f"Failed to add structured properties to {len(failed_urns)} entit(ies). Errors: {error_details}"
+        )
+
+    return {
+        "success": True,
+        "message": f"Successfully added {len(property_values)} structured propert(ies) to {success_count} entit(ies)",
+    }
+
+
+def remove_structured_properties(
+    property_urns: List[str],
+    entity_urns: List[str],
+) -> dict:
+    """Remove structured properties from multiple DataHub entities.
+
+    This tool allows you to remove structured property assignments from multiple entities in a single operation.
+
+    Args:
+        property_urns: List of structured property URNs to remove
+                      Example: ["urn:li:structuredProperty:io.acryl.privacy.retentionTime"]
+        entity_urns: List of entity URNs to remove properties from (e.g., dataset URNs, dashboard URNs)
+
+    Examples:
+        # Remove retention time property from datasets
+        remove_structured_properties(
+            property_urns=["urn:li:structuredProperty:io.acryl.privacy.retentionTime"],
+            entity_urns=[
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.old_data,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.archived,PROD)"
+            ]
+        )
+
+        # Remove multiple properties at once
+        remove_structured_properties(
+            property_urns=[
+                "urn:li:structuredProperty:io.acryl.privacy.retentionTime",
+                "urn:li:structuredProperty:io.acryl.common.businessCriticality"
+            ],
+            entity_urns=[
+                "urn:li:dataset:(urn:li:dataPlatform:snowflake,db.schema.temp_table,PROD)"
+            ]
+        )
+
+    Example:
+        from datahub_agent_context.context import DataHubContext
+
+        with DataHubContext(client.graph):
+            result = remove_structured_properties(
+                property_urns=["urn:li:structuredProperty:..."],
+                entity_urns=["urn:li:dataset:(...)"]
+            )
+    """
+    graph = get_graph()
+    if not property_urns:
+        raise ValueError("property_urns cannot be empty")
+    if not entity_urns:
+        raise ValueError("entity_urns cannot be empty")
+
+    # Validate all structured properties exist
+    for property_urn in property_urns:
+        _validate_and_fetch_structured_property(property_urn)
+
+    # Execute remove for each entity
+    mutation = """
+        mutation removeStructuredProperties($input: RemoveStructuredPropertiesInput!) {
+            removeStructuredProperties(input: $input) {
+                properties {
+                    structuredProperty {
+                        urn
+                    }
+                }
+            }
+        }
+    """
+
+    success_count = 0
+    failed_urns = []
+    error_messages = []
+
+    for entity_urn in entity_urns:
+        variables = {
+            "input": {
+                "assetUrn": entity_urn,
+                "structuredPropertyUrns": property_urns,
+            }
+        }
+
+        try:
+            result = execute_graphql(
+                graph,
+                query=mutation,
+                variables=variables,
+                operation_name="removeStructuredProperties",
+            )
+
+            if result.get("removeStructuredProperties"):
+                success_count += 1
+            else:
+                failed_urns.append(entity_urn)
+                error_messages.append(
+                    f"{entity_urn}: operation returned false or empty result"
+                )
+
+        except Exception as e:
+            failed_urns.append(entity_urn)
+            error_messages.append(f"{entity_urn}: {str(e)}")
+
+    if failed_urns:
+        error_details = "; ".join(error_messages[:3])
+        if len(error_messages) > 3:
+            error_details += f"; and {len(error_messages) - 3} more error(s)"
+        raise RuntimeError(
+            f"Failed to remove structured properties from {len(failed_urns)} entit(ies). Errors: {error_details}"
+        )
+
+    return {
+        "success": True,
+        "message": f"Successfully removed {len(property_urns)} structured propert(ies) from {success_count} entit(ies)",
+    }