acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/cli/specific/user_cli.py

@@ -1,14 +1,16 @@
 import logging
 import pathlib
 from pathlib import Path
+from typing import Optional
 
 import click
 from click_default_group import DefaultGroup
 
 from datahub.api.entities.corpuser.corpuser import CorpUser, CorpUserGenerationConfig
 from datahub.cli.specific.file_loader import load_file
-from datahub.ingestion.graph.client import get_default_graph
-from datahub.telemetry import telemetry
+from datahub.configuration.common import OperationalError
+from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
@@ -32,13 +34,12 @@ def user() -> None:
     help="Use this flag to overwrite the information that is set via the UI",
 )
 @upgrade.check_upgrade
-@telemetry.with_telemetry()
 def upsert(file: Path, override_editable: bool) -> None:
     """Create or Update a User in DataHub"""
 
     config_dict = load_file(pathlib.Path(file))
     user_configs = config_dict if isinstance(config_dict, list) else [config_dict]
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         for user_config in user_configs:
             try:
                 datahub_user: CorpUser = CorpUser.parse_obj(user_config)
@@ -56,3 +57,172 @@ def upsert(file: Path, override_editable: bool) -> None:
                     f"Update failed for id {user_config.get('id')}. due to {e}",
                     fg="red",
                 )
+
+
+def validate_user_id_options(
+    user_id: Optional[str], email_as_id: bool, email: str
+) -> str:
+    """
+    Validate user ID options and return the final user ID to use.
+
+    Args:
+        user_id: Optional explicit user ID
+        email_as_id: Whether to use email as the user ID
+        email: User's email address
+
+    Returns:
+        The final user ID to use for the URN
+
+    Raises:
+        ValueError: If validation fails (neither or both options provided)
+    """
+    if not user_id and not email_as_id:
+        raise ValueError("Must specify either --id or --email-as-id flag")
+
+    if user_id and email_as_id:
+        raise ValueError("Cannot specify both --id and --email-as-id flag")
+
+    if email_as_id:
+        return email
+
+    assert user_id is not None
+    return user_id
+
+
+def create_native_user_in_datahub(
+    graph: DataHubGraph,
+    user_id: str,
+    email: str,
+    display_name: str,
+    password: str,
+    role: Optional[str] = None,
+) -> str:
+    """
+    Create a native DataHub user.
+
+    Args:
+        graph: DataHubGraph client
+        user_id: User identifier (used in URN)
+        email: User's email address
+        display_name: User's full display name
+        password: User's password
+        role: Optional role to assign (Admin, Editor, or Reader)
+
+    Returns:
+        The created user's URN
+
+    Raises:
+        ValueError: If user already exists or role is invalid
+        OperationalError: If user creation fails due to API/network errors
+    """
+    user_urn = f"urn:li:corpuser:{user_id}"
+
+    if graph.exists(user_urn):
+        raise ValueError(f"User with ID {user_id} already exists (urn: {user_urn})")
+
+    created_user_urn = graph.create_native_user(
+        user_id=user_id,
+        email=email,
+        display_name=display_name,
+        password=password,
+        role=role,
+    )
+
+    return created_user_urn
+
+
+@user.command(name="add")
+@click.option("--id", "user_id", type=str, help="User identifier (used in URN)")
+@click.option("--email", required=True, type=str, help="User's email address")
+@click.option(
+    "--email-as-id",
+    is_flag=True,
+    default=False,
+    help="Use email address as user ID (alternative to --id)",
+)
+@click.option(
+    "--display-name", required=True, type=str, help="User's full display name"
+)
+@click.option(
+    "--password",
+    is_flag=True,
+    default=False,
+    help="Prompt for password (hidden input)",
+)
+@click.option(
+    "--role",
+    required=False,
+    type=click.Choice(
+        ["Admin", "Editor", "Reader", "admin", "editor", "reader"], case_sensitive=False
+    ),
+    help="Optional role to assign (Admin, Editor, or Reader)",
+)
+@upgrade.check_upgrade
+def add(
+    user_id: str,
+    email: str,
+    email_as_id: bool,
+    display_name: str,
+    password: bool,
+    role: str,
+) -> None:
+    """Create a native DataHub user with email/password authentication"""
+
+    try:
+        final_user_id = validate_user_id_options(user_id, email_as_id, email)
+    except ValueError as e:
+        click.secho(f"Error: {str(e)}", fg="red")
+        raise SystemExit(1) from e
+
+    if not password:
+        click.secho(
+            "Error: --password flag is required to prompt for password input",
+            fg="red",
+        )
+        raise SystemExit(1)
+
+    password_value = click.prompt(
+        "Enter password", hide_input=True, confirmation_prompt=True
+    )
+
+    with get_default_graph(ClientMode.CLI) as graph:
+        try:
+            created_user_urn = create_native_user_in_datahub(
+                graph, final_user_id, email, display_name, password_value, role
+            )
+
+            if role:
+                click.secho(
+                    f"Successfully created user {final_user_id} with role {role.capitalize()} (URN: {created_user_urn})",
+                    fg="green",
+                )
+            else:
+                click.secho(
+                    f"Successfully created user {final_user_id} (URN: {created_user_urn})",
+                    fg="green",
+                )
+        except ValueError as e:
+            click.secho(f"Error: {str(e)}", fg="red")
+            raise SystemExit(1) from e
+        except OperationalError as e:
+            error_msg = e.message if hasattr(e, "message") else str(e.args[0])
+            click.secho(f"Error: {error_msg}", fg="red")
+
+            if hasattr(e, "info") and e.info:
+                logger.debug(f"Error details: {e.info}")
+                if "status_code" in e.info:
+                    click.secho(f"  HTTP Status: {e.info['status_code']}", fg="red")
+                if "response_text" in e.info:
+                    click.secho(
+                        f"  Response: {e.info['response_text'][:200]}", fg="red"
+                    )
+
+            click.secho(
+                "\nTip: Run with DATAHUB_DEBUG=1 environment variable for detailed logs",
+                fg="yellow",
+            )
+            raise SystemExit(1) from e
+        except Exception as e:
+            click.secho(f"Unexpected error: {str(e)}", fg="red")
+            logger.exception("Unexpected error during user creation")
+            raise SystemExit(1) from e

datahub/cli/state_cli.py

@@ -5,7 +5,7 @@ import click
 from click_default_group import DefaultGroup
 
 from datahub.ingestion.graph.client import get_default_graph
-from datahub.telemetry import telemetry
+from datahub.ingestion.graph.config import ClientMode
 from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
@@ -21,14 +21,13 @@ def state() -> None:
 @click.option("--pipeline-name", required=True, type=str)
 @click.option("--platform", required=True, type=str)
 @upgrade.check_upgrade
-@telemetry.with_telemetry()
 def inspect(pipeline_name: str, platform: str) -> None:
     """
     Get the latest stateful ingestion state for a given pipeline.
     Only works for state entity removal for now.
     """
 
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.CLI)
     checkpoint = datahub_graph.get_latest_pipeline_checkpoint(pipeline_name, platform)
     if not checkpoint:
         click.secho("No ingestion state found.", fg="red")

datahub/cli/timeline_cli.py

@@ -9,7 +9,7 @@ from requests import Response
 
 from datahub.emitter.mce_builder import dataset_urn_to_key, schema_field_urn_to_key
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
-from datahub.telemetry import telemetry
+from datahub.ingestion.graph.config import ClientMode
 from datahub.upgrade import upgrade
 from datahub.utilities.urns.urn import Urn
 
@@ -63,7 +63,7 @@ def get_timeline(
     diff: bool,
     graph: Optional[DataHubGraph] = None,
 ) -> Any:
-    client = graph if graph else get_default_graph()
+    client = graph if graph else get_default_graph(ClientMode.CLI)
     session = client._session
     host = client.config.server
     if urn.startswith("urn%3A"):
@@ -129,7 +129,6 @@ def get_timeline(
 @click.option("--raw", type=bool, is_flag=True, help="Show the raw diff")
 @click.pass_context
 @upgrade.check_upgrade
-@telemetry.with_telemetry()
 def timeline(
     ctx: Any,
     urn: str,

datahub/configuration/common.py

@@ -1,20 +1,25 @@
+import dataclasses
 import re
 import unittest.mock
 from abc import ABC, abstractmethod
 from enum import auto
 from typing import (
     IO,
+    TYPE_CHECKING,
+    Annotated,
     Any,
     ClassVar,
     Dict,
     List,
     Optional,
     Type,
+    TypeVar,
     Union,
     runtime_checkable,
 )
 
 import pydantic
+import pydantic_core
 from cached_property import cached_property
 from pydantic import BaseModel, Extra, ValidationError
 from pydantic.fields import Field
@@ -33,10 +38,15 @@ REDACT_KEYS = {
 }
 REDACT_SUFFIXES = {
     "_password",
+    "-password",
     "_secret",
+    "-secret",
     "_token",
+    "-token",
     "_key",
+    "-key",
     "_key_id",
+    "-key-id",
 }
 
 
@@ -78,6 +88,29 @@ def redact_raw_config(obj: Any) -> Any:
         return obj
 
 
+if TYPE_CHECKING:
+    AnyType = TypeVar("AnyType")
+    HiddenFromDocs = Annotated[AnyType, ...]
+else:
+    HiddenFromDocs = pydantic.json_schema.SkipJsonSchema
+
+LaxStr = Annotated[str, pydantic.BeforeValidator(lambda v: str(v))]
+
+
+@dataclasses.dataclass(frozen=True)
+class SupportedSources:
+    sources: List[str]
+
+    def __get_pydantic_json_schema__(
+        self,
+        core_schema: pydantic_core.core_schema.CoreSchema,
+        handler: pydantic.GetJsonSchemaHandler,
+    ) -> pydantic.json_schema.JsonSchemaValue:
+        json_schema = handler(core_schema)
+        json_schema.setdefault("schema_extra", {})["supported_sources"] = self.sources
+        return json_schema
+
+
 class ConfigModel(BaseModel):
     class Config:
         @staticmethod
@@ -136,6 +169,17 @@ class PermissiveConfigModel(ConfigModel):
             extra = Extra.allow
 
 
+class ConnectionModel(BaseModel):
+    """Represents the config associated with a connection"""
+
+    class Config:
+        if PYDANTIC_VERSION_2:
+            extra = "allow"
+        else:
+            extra = Extra.allow
+            underscore_attrs_are_private = True
+
+
 class TransformerSemantics(ConfigEnum):
     """Describes semantics for aspect changes"""
 
@@ -317,7 +361,7 @@ class KeyValuePattern(ConfigModel):
         return KeyValuePattern()
 
     def value(self, string: str) -> List[str]:
-        matching_keys = [key for key in self.rules.keys() if re.match(key, string)]
+        matching_keys = [key for key in self.rules if re.match(key, string)]
         if not matching_keys:
             return []
         elif self.first_match_only:
@@ -329,4 +373,4 @@ class KeyValuePattern(ConfigModel):
 
 
 class VersionedConfig(ConfigModel):
-    version: str = "1"
+    version: LaxStr = "1"

datahub/configuration/connection_resolver.py

@@ -1,13 +1,16 @@
-from typing import Type
+from typing import TYPE_CHECKING, Type
 
 import pydantic
 
 from datahub.ingestion.api.global_context import get_graph_context
 
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1RootValidator
+
 
 def auto_connection_resolver(
     connection_field: str = "connection",
-) -> classmethod:
+) -> "V1RootValidator":
     def _resolve_connection(cls: Type, values: dict) -> dict:
         if connection_field in values:
             connection_urn = values.pop(connection_field)