datahub-agent-context 1.3.1.10rc1__tar.gz → 1.4.0rc2__tar.gz

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/PKG-INFO

@@ -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
@@ -29,26 +29,32 @@ Classifier: Topic :: Software Development
 Requires-Python: >=3.9
 Description-Content-Type: text/markdown
 Requires-Dist: jmespath<2.0.0,>=1.0.0
+Requires-Dist: acryl-datahub==1.4.0rc2
 Requires-Dist: cachetools<7.0.0,>=5.0.0
-Requires-Dist: acryl-datahub==1.3.1.10rc1
+Requires-Dist: httpcore<2.0,>=1.0.9
 Requires-Dist: pydantic<3.0.0,>=2.0.0
-Requires-Dist: json-repair<1.0.0,>=0.25.0
 Requires-Dist: h11<1.0,>=0.16
 Requires-Dist: google-re2<2.0,>=1.0
-Requires-Dist: httpcore<2.0,>=1.0.9
+Requires-Dist: json-repair<1.0.0,>=0.25.0
 Provides-Extra: dev
-Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
+Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
+Requires-Dist: click<9.0.0,>=8.0.0; extra == "dev"
 Requires-Dist: types-PyYAML<7.0.0,>=6.0.0; extra == "dev"
-Requires-Dist: mypy==1.17.1; extra == "dev"
 Requires-Dist: types-toml<1.0.0,>=0.10.0; extra == "dev"
+Requires-Dist: langchain-core<2.0.0,>=1.2.7; extra == "dev"
 Requires-Dist: types-cachetools<7.0.0,>=5.0.0; extra == "dev"
-Requires-Dist: pytest<9.0.0,>=8.3.4; extra == "dev"
 Requires-Dist: pytest-cov<7.0.0,>=2.8.0; extra == "dev"
-Requires-Dist: types-requests<3.0.0,>=2.0.0; extra == "dev"
-Requires-Dist: tox<5.0.0,>=4.0.0; extra == "dev"
+Requires-Dist: mypy==1.17.1; extra == "dev"
 Requires-Dist: ruff==0.11.7; extra == "dev"
+Requires-Dist: types-jmespath<2.0.0,>=1.0.0; extra == "dev"
+Requires-Dist: snowflake-connector-python<4.0.0,>=3.0.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"
 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
 

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/README.md

@@ -43,7 +43,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"]},
@@ -51,7 +51,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"]]
     )
@@ -113,6 +113,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
 

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/setup.py

@@ -69,9 +69,16 @@ langchain_requirements = {
     "langchain-core>=1.2.7,<2.0.0",
 }
 
+snowflake_requirements = {
+    "click>=8.0.0,<9.0.0",
+    "snowflake-connector-python>=3.0.0,<4.0.0",
+}
+
 dev_requirements = {
     *lint_requirements,
     *mypy_stubs,
+    *snowflake_requirements,
+    *langchain_requirements,
     "pytest>=8.3.4,<9.0.0",
     "pytest-cov>=2.8.0,<7.0.0",
     "tox>=4.0.0,<5.0.0",
@@ -123,5 +130,6 @@ setuptools.setup(
     extras_require={
         "dev": list(dev_requirements),
         "langchain": list(langchain_requirements),
+        "snowflake": list(snowflake_requirements),
     },
 )

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/__init__.py

@@ -17,9 +17,17 @@
 from datahub_agent_context._version import __version__
 from datahub_agent_context.context import (
     DataHubContext,
+    get_datahub_client,
     get_graph,
-    reset_graph,
-    set_graph,
+    reset_client,
+    set_client,
 )
 
-__all__ = ["__version__", "DataHubContext", "get_graph", "set_graph", "reset_graph"]
+__all__ = [
+    "__version__",
+    "DataHubContext",
+    "get_datahub_client",
+    "get_graph",
+    "set_client",
+    "reset_client",
+]

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/_version.py

@@ -13,4 +13,4 @@
 # limitations under the License.
 
 __package_name__ = "datahub-agent-context"
-__version__ = "1.3.1.10rc1"
+__version__ = "1.4.0rc2"

datahub_agent_context-1.4.0rc2/src/datahub_agent_context/cli.py

@@ -0,0 +1,152 @@
+import logging
+
+import click
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def agent() -> None:
+    """Helper commands for Creating and managing Agent on DataHub."""
+    pass
+
+
+@agent.group()
+def create() -> None:
+    """Create an agent on DataHub."""
+    pass
+
+
+@create.command(name="snowflake")
+@click.option(
+    "--sf-account",
+    default=None,
+    help="Snowflake account identifier (auto-detected if not provided)",
+)
+@click.option(
+    "--sf-user",
+    default=None,
+    help="Snowflake user name (auto-detected if not provided)",
+)
+@click.option(
+    "--sf-role",
+    default=None,
+    help="Snowflake role (auto-detected if not provided)",
+)
+@click.option(
+    "--sf-warehouse",
+    default=None,
+    help="Snowflake warehouse name (auto-detected if not provided)",
+)
+@click.option(
+    "--sf-database",
+    default=None,
+    help="Snowflake database name (auto-detected if not provided)",
+)
+@click.option(
+    "--sf-schema",
+    default=None,
+    help="Snowflake schema name (auto-detected if not provided)",
+)
+@click.option("--datahub-url", required=True, help="DataHub instance URL")
+@click.option("--datahub-token", required=True, help="DataHub Personal Access Token")
+@click.option(
+    "--agent-name", default="DATAHUB_SQL_AGENT", help="Agent name in Snowflake"
+)
+@click.option(
+    "--agent-display-name",
+    default="DataHub SQL Assistant",
+    help="Agent display name in Snowflake UI",
+)
+@click.option(
+    "--agent-color",
+    default="blue",
+    help="Agent color in Snowflake UI",
+    type=click.Choice(["blue", "red", "green", "yellow", "purple", "orange"]),
+)
+@click.option(
+    "--output-dir",
+    default="./snowflake_setup",
+    help="Output directory for generated SQL files",
+)
+@click.option(
+    "--execute",
+    is_flag=True,
+    default=False,
+    help="Connect to Snowflake and execute the SQL scripts directly",
+)
+@click.option(
+    "--sf-password",
+    help="Snowflake password (required if --execute is used with password authentication)",
+)
+@click.option(
+    "--sf-authenticator",
+    default="snowflake",
+    type=click.Choice(["snowflake", "externalbrowser", "oauth"], case_sensitive=False),
+    help="Authentication method: 'snowflake' (password), 'externalbrowser' (SSO), or 'oauth' (token-based). Default: snowflake",
+)
+@click.option(
+    "--enable-mutations/--no-enable-mutations",
+    default=True,
+    help="Include mutation/write tools (tags, descriptions, owners, etc.). Default: enabled",
+)
+def create_snowflake(
+    sf_account: str | None,
+    sf_user: str | None,
+    sf_role: str | None,
+    sf_warehouse: str | None,
+    sf_database: str | None,
+    sf_schema: str | None,
+    datahub_url: str,
+    datahub_token: str,
+    agent_name: str,
+    agent_display_name: str,
+    agent_color: str,
+    output_dir: str,
+    execute: bool,
+    sf_password: str | None,
+    sf_authenticator: str,
+    enable_mutations: bool,
+) -> None:
+    """Create a Snowflake agent on DataHub.
+
+    Snowflake connection parameters (account, user, role, warehouse, database, schema)
+    will be auto-detected from your current session if not provided explicitly.
+
+    Authentication methods:
+    - snowflake (default): Standard password authentication
+    - externalbrowser: SSO authentication via browser (no password required)
+    - oauth: OAuth token-based authentication
+
+    Examples:
+        # Generate SQL files only (default)
+        $ datahub agent create snowflake --datahub-url=... --datahub-token=...
+
+        # Execute with password authentication
+        $ datahub agent create snowflake --execute --sf-password=secret ...
+
+        # Execute with SSO authentication (browser-based)
+        $ datahub agent create snowflake --execute --sf-authenticator=externalbrowser ...
+    """
+    from datahub_agent_context.snowflake.snowflake import create_snowflake_agent
+
+    ctx = click.get_current_context()
+    ctx.invoke(
+        create_snowflake_agent,
+        sf_account=sf_account,
+        sf_user=sf_user,
+        sf_role=sf_role,
+        sf_warehouse=sf_warehouse,
+        sf_database=sf_database,
+        sf_schema=sf_schema,
+        datahub_url=datahub_url,
+        datahub_token=datahub_token,
+        agent_name=agent_name,
+        agent_display_name=agent_display_name,
+        agent_color=agent_color,
+        output_dir=output_dir,
+        execute=execute,
+        sf_password=sf_password,
+        sf_authenticator=sf_authenticator,
+        enable_mutations=enable_mutations,
+    )

datahub_agent_context-1.4.0rc2/src/datahub_agent_context/context.py

@@ -0,0 +1,110 @@
+"""Context management for DataHub tools.
+
+This module provides a context manager pattern for managing DataHubClient instances
+across tool calls without explicit parameter passing.
+"""
+
+import contextvars
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from datahub.ingestion.graph.client import DataHubGraph
+    from datahub.sdk.main_client import DataHubClient
+
+# Context variable to store the current DataHubClient instance
+_client_context: contextvars.ContextVar[Optional["DataHubClient"]] = (
+    contextvars.ContextVar("datahub_client", default=None)
+)
+
+
+def get_datahub_client() -> "DataHubClient":
+    """Get the current DataHubClient from context.
+
+    Returns:
+        DataHubClient instance from context
+
+    Raises:
+        RuntimeError: If no client is set in context
+    """
+    client = _client_context.get()
+    if client is None:
+        raise RuntimeError(
+            "No DataHubClient in context. "
+            "Make sure to use DataHubContext context manager or set_client() before calling tools."
+        )
+    return client
+
+
+def get_graph() -> "DataHubGraph":
+    """Get the current DataHubGraph from context (convenience method).
+
+    Returns:
+        DataHubGraph instance from the client in context
+
+    Raises:
+        RuntimeError: If no client is set in context
+    """
+    return get_datahub_client()._graph
+
+
+def set_client(client: "DataHubClient") -> contextvars.Token:
+    """Set the DataHubClient in context.
+
+    Args:
+        client: DataHubClient instance to set
+
+    Returns:
+        Token that can be used to reset the context
+    """
+    return _client_context.set(client)
+
+
+def reset_client(token: contextvars.Token) -> None:
+    """Reset the DataHubClient context to its previous value.
+
+    Args:
+        token: Token returned by set_client()
+    """
+    _client_context.reset(token)
+
+
+class DataHubContext:
+    """Context manager for DataHub tool execution.
+
+    This context manager sets the DataHubClient in context for the duration
+    of the with block, allowing tools to access it without explicit parameter passing.
+
+    Example:
+        from datahub.sdk.main_client import DataHubClient
+        from datahub_agent_context.context import DataHubContext
+        from datahub_agent_context.mcp_tools import search
+
+        client = DataHubClient(...)
+
+        with DataHubContext(client):
+            results = search(query="users")  # No client parameter needed!
+    """
+
+    def __init__(self, client: "DataHubClient"):
+        """Initialize the context manager.
+
+        Args:
+            client: DataHubClient instance to use in this context
+        """
+        self.client = client
+        self._token: Optional[contextvars.Token] = None
+
+    def __enter__(self) -> "DataHubClient":
+        """Enter the context and set the client.
+
+        Returns:
+            The DataHubClient instance
+        """
+        self._token = set_client(self.client)
+        return self.client
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Exit the context and reset the client."""
+        if self._token is not None:
+            reset_client(self._token)
+            self._token = None

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/langchain_tools/builder.py

@@ -3,7 +3,7 @@
 import functools
 from typing import TYPE_CHECKING, Callable
 
-from datahub_agent_context.context import set_graph
+from datahub_agent_context.context import set_client
 from datahub_agent_context.mcp_tools import get_me
 from datahub_agent_context.mcp_tools.documents import grep_documents, search_documents
 from datahub_agent_context.mcp_tools.domains import remove_domains, set_domains
@@ -32,6 +32,7 @@ from datahub_agent_context.mcp_tools.lineage import (
 )
 from datahub_agent_context.mcp_tools.owners import add_owners, remove_owners
 from datahub_agent_context.mcp_tools.queries import get_dataset_queries
+from datahub_agent_context.mcp_tools.save_document import save_document
 from datahub_agent_context.mcp_tools.search import search
 from datahub_agent_context.mcp_tools.tags import add_tags, remove_tags
 from datahub_agent_context.mcp_tools.terms import (
@@ -57,14 +58,14 @@ def create_context_wrapper(func: Callable, client: "DataHubClient") -> Callable:
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
         # Set graph in context for this function call
-        token = set_graph(client._graph)
+        token = set_client(client)
         try:
             return func(*args, **kwargs)
         finally:
             # Always reset context, even if function raises
-            from datahub_agent_context.context import reset_graph
+            from datahub_agent_context.context import reset_client
 
-            reset_graph(token)
+            reset_client(token)
 
     return wrapper
 
@@ -123,5 +124,6 @@ def build_langchain_tools(
         tools.append(tool(create_context_wrapper(remove_tags, client)))
         tools.append(tool(create_context_wrapper(add_glossary_terms, client)))
         tools.append(tool(create_context_wrapper(remove_glossary_terms, client)))
+        tools.append(tool(create_context_wrapper(save_document, client)))
 
     return tools

{datahub_agent_context-1.3.1.10rc1 → datahub_agent_context-1.4.0rc2}/src/datahub_agent_context/mcp_tools/base.py

@@ -152,9 +152,12 @@ def _is_datahub_cloud(graph: DataHubGraph) -> bool:
         )
         return False
 
-    is_cloud = hasattr(graph, "frontend_base_url") and graph.frontend_base_url
-    logger.debug(f"Cloud detection: {is_cloud}")
-    return bool(is_cloud)
+    try:
+        is_cloud = hasattr(graph, "frontend_base_url") and graph.frontend_base_url
+        logger.debug(f"Cloud detection: {is_cloud}")
+        return bool(is_cloud)
+    except ValueError:
+        return False
 
 
 def _is_field_validation_error(error_msg: str) -> bool: