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

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/_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/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/context.py

@@ -1,6 +1,6 @@
 """Context management for DataHub tools.
 
-This module provides a context manager pattern for managing DataHubGraph instances
+This module provides a context manager pattern for managing DataHubClient instances
 across tool calls without explicit parameter passing.
 """
 
@@ -9,56 +9,69 @@ 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 DataHubGraph instance
-_graph_context: contextvars.ContextVar[Optional["DataHubGraph"]] = (
-    contextvars.ContextVar("datahub_graph", default=None)
+# Context variable to store the current DataHubClient instance
+_client_context: contextvars.ContextVar[Optional["DataHubClient"]] = (
+    contextvars.ContextVar("datahub_client", default=None)
 )
 
 
-def get_graph() -> "DataHubGraph":
-    """Get the current DataHubGraph from context.
+def get_datahub_client() -> "DataHubClient":
+    """Get the current DataHubClient from context.
 
     Returns:
-        DataHubGraph instance from context
+        DataHubClient instance from context
 
     Raises:
-        RuntimeError: If no graph is set in context
+        RuntimeError: If no client is set in context
     """
-    graph = _graph_context.get()
-    if graph is None:
+    client = _client_context.get()
+    if client is None:
         raise RuntimeError(
-            "No DataHubGraph in context. "
-            "Make sure to use DataHubContext context manager or set_graph() before calling tools."
+            "No DataHubClient in context. "
+            "Make sure to use DataHubContext context manager or set_client() before calling tools."
         )
-    return graph
+    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_graph(graph: "DataHubGraph") -> contextvars.Token:
-    """Set the DataHubGraph in context.
+def set_client(client: "DataHubClient") -> contextvars.Token:
+    """Set the DataHubClient in context.
 
     Args:
-        graph: DataHubGraph instance to set
+        client: DataHubClient instance to set
 
     Returns:
         Token that can be used to reset the context
     """
-    return _graph_context.set(graph)
+    return _client_context.set(client)
 
 
-def reset_graph(token: contextvars.Token) -> None:
-    """Reset the DataHubGraph context to its previous value.
+def reset_client(token: contextvars.Token) -> None:
+    """Reset the DataHubClient context to its previous value.
 
     Args:
-        token: Token returned by set_graph()
+        token: Token returned by set_client()
     """
-    _graph_context.reset(token)
+    _client_context.reset(token)
 
 
 class DataHubContext:
     """Context manager for DataHub tool execution.
 
-    This context manager sets the DataHubGraph in context for the duration
+    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:
@@ -68,30 +81,30 @@ class DataHubContext:
 
         client = DataHubClient(...)
 
-        with DataHubContext(client.graph):
-            results = search(query="users")  # No graph parameter needed!
+        with DataHubContext(client):
+            results = search(query="users")  # No client parameter needed!
     """
 
-    def __init__(self, graph: "DataHubGraph"):
+    def __init__(self, client: "DataHubClient"):
         """Initialize the context manager.
 
         Args:
-            graph: DataHubGraph instance to use in this context
+            client: DataHubClient instance to use in this context
         """
-        self.graph = graph
+        self.client = client
         self._token: Optional[contextvars.Token] = None
 
-    def __enter__(self) -> "DataHubGraph":
-        """Enter the context and set the graph.
+    def __enter__(self) -> "DataHubClient":
+        """Enter the context and set the client.
 
         Returns:
-            The DataHubGraph instance
+            The DataHubClient instance
         """
-        self._token = set_graph(self.graph)
-        return self.graph
+        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 graph."""
+        """Exit the context and reset the client."""
         if self._token is not None:
-            reset_graph(self._token)
+            reset_client(self._token)
             self._token = None

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/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: