datahub-agent-context 1.3.1.10rc1__py3-none-any.whl → 1.4.0rc2__py3-none-any.whl

datahub_agent_context/snowflake/udfs/list_schema_fields.py

@@ -0,0 +1,80 @@
+"""LIST_SCHEMA_FIELDS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_list_schema_fields_udf() -> str:
+    """Generate LIST_SCHEMA_FIELDS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.list_schema_fields() to list
+    schema fields for a dataset with optional keyword filtering and pagination.
+
+    Useful when schema fields were truncated in search results and you need to
+    explore specific columns. Supports pagination for large schemas.
+
+    Parameters:
+        dataset_urn (STRING): Dataset URN
+        keywords (STRING): Optional keywords to filter schema fields (OR matching).
+                          Can be a JSON array string or a single keyword.
+                          - Single string: Treated as one keyword for exact field name/phrase
+                          - JSON array: Multiple keywords, matches any (OR logic)
+                          - Empty/null: Returns all fields in priority order
+                          Matches against fieldPath, description, label, tags, and glossary terms.
+        limit (NUMBER): Maximum number of fields to return (default: 100)
+
+    Returns:
+        VARIANT: Dictionary with:
+                - urn: The dataset URN
+                - fields: List of schema fields (paginated)
+                - totalFields: Total number of fields in the schema
+                - returned: Number of fields actually returned
+                - remainingCount: Number of fields not included
+                - matchingCount: Number of fields that matched keywords (if provided)
+
+    Examples:
+        - Single keyword: LIST_SCHEMA_FIELDS(urn, 'user_email', 100)
+        - Multiple keywords: LIST_SCHEMA_FIELDS(urn, '["email", "user"]', 100)
+        - All fields: LIST_SCHEMA_FIELDS(urn, NULL, 100)
+    """
+    function_body = """from datahub_agent_context.mcp_tools import list_schema_fields
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    keyword_list = None
+    if keywords:
+        try:
+            keyword_list = json.loads(keywords) if isinstance(keywords, str) else keywords
+        except json.JSONDecodeError:
+            keyword_list = [keywords]
+
+    with DataHubContext(graph):
+        return list_schema_fields(
+            urn=dataset_urn,
+            keywords=keyword_list,
+            limit=limit
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e),
+        'urn': dataset_urn
+    }"""
+
+    return generate_python_udf_code(
+        function_name="LIST_SCHEMA_FIELDS",
+        parameters=[
+            ("dataset_urn", "STRING"),
+            ("keywords", "STRING"),
+            ("limit", "NUMBER"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/remove_domains.py

@@ -0,0 +1,45 @@
+"""REMOVE_DOMAINS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_remove_domains_udf() -> str:
+    """Generate REMOVE_DOMAINS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.remove_domains() to remove
+    domain assignments from multiple DataHub entities.
+
+    Parameters:
+        entity_urns (STRING): JSON array of entity URNs to remove domain from
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+    """
+    function_body = """from datahub_agent_context.mcp_tools import remove_domains
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+
+    with DataHubContext(graph):
+        return remove_domains(entity_urns=entity_urn_list)
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="REMOVE_DOMAINS",
+        parameters=[("entity_urns", "STRING")],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/remove_glossary_terms.py

@@ -0,0 +1,57 @@
+"""REMOVE_GLOSSARY_TERMS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_remove_glossary_terms_udf() -> str:
+    """Generate REMOVE_GLOSSARY_TERMS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.remove_glossary_terms() to
+    remove glossary terms from multiple DataHub entities or their columns.
+
+    Parameters:
+        term_urns (STRING): JSON array of glossary term URNs to remove
+        entity_urns (STRING): JSON array of entity URNs
+        column_paths (STRING): Optional JSON array of column names (use NULL for entity-level)
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+    """
+    function_body = """from datahub_agent_context.mcp_tools import remove_glossary_terms
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    term_urn_list = json.loads(term_urns) if isinstance(term_urns, str) else term_urns
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+    column_path_list = json.loads(column_paths) if column_paths and isinstance(column_paths, str) else None
+
+    with DataHubContext(graph):
+        return remove_glossary_terms(
+            term_urns=term_urn_list,
+            entity_urns=entity_urn_list,
+            column_paths=column_path_list
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="REMOVE_GLOSSARY_TERMS",
+        parameters=[
+            ("term_urns", "STRING"),
+            ("entity_urns", "STRING"),
+            ("column_paths", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/remove_owners.py

@@ -0,0 +1,56 @@
+"""REMOVE_OWNERS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_remove_owners_udf() -> str:
+    """Generate REMOVE_OWNERS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.remove_owners() to remove
+    owners from multiple DataHub entities.
+
+    Parameters:
+        owner_urns (STRING): JSON array of owner URNs to remove
+        entity_urns (STRING): JSON array of entity URNs to remove ownership from
+        ownership_type_urn (STRING): Optional ownership type URN (use NULL for all types)
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+    """
+    function_body = """from datahub_agent_context.mcp_tools import remove_owners
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    owner_urn_list = json.loads(owner_urns) if isinstance(owner_urns, str) else owner_urns
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+
+    with DataHubContext(graph):
+        return remove_owners(
+            owner_urns=owner_urn_list,
+            entity_urns=entity_urn_list,
+            ownership_type_urn=ownership_type_urn if ownership_type_urn else None
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="REMOVE_OWNERS",
+        parameters=[
+            ("owner_urns", "STRING"),
+            ("entity_urns", "STRING"),
+            ("ownership_type_urn", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/remove_structured_properties.py

@@ -0,0 +1,56 @@
+"""REMOVE_STRUCTURED_PROPERTIES UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_remove_structured_properties_udf() -> str:
+    """Generate REMOVE_STRUCTURED_PROPERTIES UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.remove_structured_properties()
+    to remove structured properties from multiple DataHub entities.
+
+    Parameters:
+        property_urns (STRING): JSON array of structured property URNs to remove
+        entity_urns (STRING): JSON array of entity URNs to remove properties from
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+
+    Example:
+        - Remove: REMOVE_STRUCTURED_PROPERTIES('["urn:li:structuredProperty:retentionTime"]', '["urn:li:dataset:(...)"]')
+    """
+    function_body = """from datahub_agent_context.mcp_tools import remove_structured_properties
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    property_urn_list = json.loads(property_urns) if isinstance(property_urns, str) else property_urns
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+
+    with DataHubContext(graph):
+        return remove_structured_properties(
+            property_urns=property_urn_list,
+            entity_urns=entity_urn_list
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="REMOVE_STRUCTURED_PROPERTIES",
+        parameters=[
+            ("property_urns", "STRING"),
+            ("entity_urns", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/remove_tags.py

@@ -0,0 +1,57 @@
+"""REMOVE_TAGS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_remove_tags_udf() -> str:
+    """Generate REMOVE_TAGS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.remove_tags() to remove tags
+    from multiple DataHub entities or their columns.
+
+    Parameters:
+        tag_urns (STRING): JSON array of tag URNs to remove
+        entity_urns (STRING): JSON array of entity URNs to untag
+        column_paths (STRING): Optional JSON array of column names (use NULL for entity-level)
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+    """
+    function_body = """from datahub_agent_context.mcp_tools import remove_tags
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    tag_urn_list = json.loads(tag_urns) if isinstance(tag_urns, str) else tag_urns
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+    column_path_list = json.loads(column_paths) if column_paths and isinstance(column_paths, str) else None
+
+    with DataHubContext(graph):
+        return remove_tags(
+            tag_urns=tag_urn_list,
+            entity_urns=entity_urn_list,
+            column_paths=column_path_list
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="REMOVE_TAGS",
+        parameters=[
+            ("tag_urns", "STRING"),
+            ("entity_urns", "STRING"),
+            ("column_paths", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/search_datahub.py

@@ -0,0 +1,71 @@
+"""SEARCH_DATAHUB UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_search_datahub_udf() -> str:
+    """Generate SEARCH_DATAHUB UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.search() to enable searching
+    across DataHub entities using structured full-text search from Snowflake.
+
+    Search Syntax:
+        - Structured full-text search - always start queries with /q
+        - Use + operator for AND (handles punctuation better than quotes)
+        - Supports boolean logic: AND (default), OR, NOT, parentheses, field searches
+        - Examples:
+          • /q user+transaction → requires both terms
+          • /q wizard OR pet → entities containing either term
+          • /q revenue* → wildcard matching
+          • /q tag:PII → search by tag name
+
+    Parameters:
+        search_query (STRING): Search query string (use /q prefix for structured queries)
+        entity_type (STRING): Optional entity type filter as JSON array string
+                             (e.g., '["dataset"]', '["dashboard", "chart"]')
+                             Use NULL to search across all entity types
+
+    Returns:
+        VARIANT: Dictionary with search results, facets, and metadata
+    """
+    function_body = """from datahub_agent_context.mcp_tools import search
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    # Parse entity_type filter if provided
+    filters = {}
+    if entity_type:
+        try:
+            entity_type_list = json.loads(entity_type) if isinstance(entity_type, str) else entity_type
+            filters["entity_type"] = entity_type_list
+        except json.JSONDecodeError:
+            # If not valid JSON, treat as single entity type
+            filters["entity_type"] = [entity_type]
+
+    with DataHubContext(graph):
+        return search(
+            query=search_query,
+            filters=filters if filters else None,
+            num_results=10
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e),
+        'query': search_query
+    }"""
+
+    return generate_python_udf_code(
+        function_name="SEARCH_DATAHUB",
+        parameters=[("search_query", "STRING"), ("entity_type", "STRING")],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/search_documents.py

@@ -0,0 +1,58 @@
+"""SEARCH_DOCUMENTS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_search_documents_udf() -> str:
+    """Generate SEARCH_DOCUMENTS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.search_documents() to search for
+    documents stored in DataHub (runbooks, FAQs, knowledge articles from Notion,
+    Confluence, etc.).
+
+    Returns document metadata without content. Use GET_ENTITIES with document URN
+    to retrieve full content.
+
+    Parameters:
+        search_query (STRING): Search query string (use /q prefix for structured queries)
+        num_results (NUMBER): Number of results to return (max 50, default: 10)
+
+    Returns:
+        VARIANT: Dictionary with search results and facets showing document metadata
+
+    Examples:
+        - Find deployment docs: SEARCH_DOCUMENTS('deployment', 10)
+        - Structured search: SEARCH_DOCUMENTS('/q kubernetes OR k8s', 20)
+    """
+    function_body = """from datahub_agent_context.mcp_tools import search_documents
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    with DataHubContext(graph):
+        return search_documents(
+            query=search_query,
+            num_results=int(num_results) if num_results else 10
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e),
+        'query': search_query
+    }"""
+
+    return generate_python_udf_code(
+        function_name="SEARCH_DOCUMENTS",
+        parameters=[
+            ("search_query", "STRING"),
+            ("num_results", "NUMBER"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/set_domains.py

@@ -0,0 +1,55 @@
+"""SET_DOMAINS UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_set_domains_udf() -> str:
+    """Generate SET_DOMAINS UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.set_domains() to assign a domain
+    to multiple DataHub entities.
+
+    Parameters:
+        domain_urn (STRING): Domain URN to assign (e.g., 'urn:li:domain:marketing')
+        entity_urns (STRING): JSON array of entity URNs to assign to the domain
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+
+    Example:
+        - Set domain: SET_DOMAINS('urn:li:domain:marketing', '["urn:li:dataset:(...)"]')
+    """
+    function_body = """from datahub_agent_context.mcp_tools import set_domains
+import json
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    entity_urn_list = json.loads(entity_urns) if isinstance(entity_urns, str) else entity_urns
+
+    with DataHubContext(graph):
+        return set_domains(
+            domain_urn=domain_urn,
+            entity_urns=entity_urn_list
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="SET_DOMAINS",
+        parameters=[
+            ("domain_urn", "STRING"),
+            ("entity_urns", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

datahub_agent_context/snowflake/udfs/update_description.py

@@ -0,0 +1,60 @@
+"""UPDATE_DESCRIPTION UDF generator."""
+
+from datahub_agent_context.snowflake.udfs.base import generate_python_udf_code
+
+
+def generate_update_description_udf() -> str:
+    """Generate UPDATE_DESCRIPTION UDF using datahub-agent-context.
+
+    This UDF wraps datahub_agent_context.mcp_tools.update_description() to update
+    descriptions for DataHub entities or their columns.
+
+    Parameters:
+        entity_urn (STRING): Entity URN to update
+        operation (STRING): Operation type - 'replace', 'append', or 'remove'
+        description (STRING): Description text (not needed for 'remove')
+        column_path (STRING): Optional column name (use NULL for entity-level)
+
+    Returns:
+        VARIANT: Dictionary with success status and message
+
+    Examples:
+        - Replace: UPDATE_DESCRIPTION(urn, 'replace', 'New description', NULL)
+        - Append: UPDATE_DESCRIPTION(urn, 'append', ' (PII)', 'email')
+        - Remove: UPDATE_DESCRIPTION(urn, 'remove', NULL, 'old_field')
+    """
+    function_body = """from datahub_agent_context.mcp_tools import update_description
+try:
+    datahub_url = _snowflake.get_generic_secret_string('datahub_url_secret')
+    datahub_token = _snowflake.get_generic_secret_string('datahub_token_secret')
+    datahub_url = datahub_url.rstrip('/')
+
+    graph = DataHubGraph(
+        config=DatahubClientConfig(server=datahub_url, token=datahub_token)
+    )
+
+    with DataHubContext(graph):
+        return update_description(
+            entity_urn=entity_urn,
+            operation=operation,
+            description=description if description else None,
+            column_path=column_path if column_path else None
+        )
+
+except Exception as e:
+    return {
+        'success': False,
+        'error': str(e)
+    }"""
+
+    return generate_python_udf_code(
+        function_name="UPDATE_DESCRIPTION",
+        parameters=[
+            ("entity_urn", "STRING"),
+            ("operation", "STRING"),
+            ("description", "STRING"),
+            ("column_path", "STRING"),
+        ],
+        return_type="VARIANT",
+        function_body=function_body,
+    )

{datahub_agent_context-1.3.1.10rc1.dist-info → datahub_agent_context-1.4.0rc2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: datahub-agent-context
-Version: 1.3.1.10rc1
+Version: 1.4.0rc2
 Summary: DataHub Agent Context - MCP Tools for AI Agents
 Home-page: https://datahub.io/
 License: Apache License 2.0
@@ -28,27 +28,33 @@ Classifier: Environment :: MacOS X
 Classifier: Topic :: Software Development
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: cachetools<7.0.0,>=5.0.0
-Requires-Dist: json-repair<1.0.0,>=0.25.0
-Requires-Dist: pydantic<3.0.0,>=2.0.0
+Requires-Dist: h11<1.0,>=0.16
 Requires-Dist: google-re2<2.0,>=1.0
-Requires-Dist: jmespath<2.0.0,>=1.0.0
+Requires-Dist: acryl-datahub==1.4.0rc2
+Requires-Dist: pydantic<3.0.0,>=2.0.0
 Requires-Dist: httpcore<2.0,>=1.0.9
-Requires-Dist: h11<1.0,>=0.16
-Requires-Dist: acryl-datahub==1.3.1.10rc1
+Requires-Dist: jmespath<2.0.0,>=1.0.0
+Requires-Dist: json-repair<1.0.0,>=0.25.0
+Requires-Dist: cachetools<7.0.0,>=5.0.0
 Provides-Extra: dev
-Requires-Dist: mypy==1.17.1; extra == "dev"
 Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "dev"
+Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.0; extra == "dev"
+Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
 Requires-Dist: types-cachetools<7.0.0,>=5.0.0; extra == "dev"
-Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
 Requires-Dist: types-toml<1.0.0,>=0.10.0; extra == "dev"
-Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
+Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
+Requires-Dist: click<9.0.0,>=8.0.0; extra == "dev"
+Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
 Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
-Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
 Requires-Dist: pytest<9.0.0,>=8.3.4; extra == "dev"
-Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
+Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
 Provides-Extra: langchain
 Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "langchain"
+Provides-Extra: snowflake
+Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.0; extra == "snowflake"
+Requires-Dist: click<9.0.0,>=8.0.0; extra == "snowflake"
 Dynamic: classifier
 Dynamic: description
 Dynamic: description-content-type
@@ -105,7 +111,7 @@ from datahub_agent_context.mcp_tools.entities import get_entities
 client = DataHubClient.from_env()
 
 # Search for datasets
-with client.graph as graph:
+with client as client:
     results = search(
         query="user_data",
         filters={"entity_type": ["dataset"]},
@@ -113,7 +119,7 @@ with client.graph as graph:
     )
 
 # Get detailed entity information
-with client.graph as graph:
+with client as client:
     entities = get_entities(
         urns=[result["entity"]["urn"] for result in results["searchResults"]]
     )
@@ -175,6 +181,7 @@ agent = create_agent(model, tools=tools, system_prompt="...")
 - `add_owners()`, `remove_owners()` - Manage owners
 - `add_glossary_terms()`, `remove_glossary_terms()` - Manage glossary terms
 - `add_structured_properties()`, `remove_structured_properties()` - Manage structured properties
+- `save_document()` - Save or update a Document.
 
 #### User Tools