acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/cli/specific/user_cli.py

@@ -1,15 +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.configuration.common import OperationalError
+from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
 from datahub.ingestion.graph.config import ClientMode
-from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
@@ -33,7 +34,6 @@ 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"""
 
@@ -57,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

@@ -6,7 +6,6 @@ from click_default_group import DefaultGroup
 
 from datahub.ingestion.graph.client import get_default_graph
 from datahub.ingestion.graph.config import ClientMode
-from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
 logger = logging.getLogger(__name__)
@@ -22,7 +21,6 @@ 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.

datahub/cli/timeline_cli.py

@@ -10,7 +10,6 @@ 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.ingestion.graph.config import ClientMode
-from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.urns.urn import Urn
 
@@ -130,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
@@ -83,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
@@ -141,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"""
 
@@ -334,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)

datahub/configuration/env_vars.py

@@ -0,0 +1,331 @@
+# ABOUTME: Central registry for all environment variables used in metadata-ingestion.
+# ABOUTME: All environment variable reads should go through this module for discoverability and maintainability.
+
+import os
+from typing import Optional
+
+# ============================================================================
+# Core DataHub Configuration
+# ============================================================================
+
+
+def get_gms_url() -> Optional[str]:
+    """Complete GMS URL (takes precedence over separate host/port)."""
+    return os.getenv("DATAHUB_GMS_URL")
+
+
+def get_gms_host() -> Optional[str]:
+    """GMS host (fallback for URL, deprecated)."""
+    return os.getenv("DATAHUB_GMS_HOST")
+
+
+def get_gms_port() -> Optional[str]:
+    """GMS port number."""
+    return os.getenv("DATAHUB_GMS_PORT")
+
+
+def get_gms_protocol() -> str:
+    """Protocol for GMS connection (http/https)."""
+    return os.getenv("DATAHUB_GMS_PROTOCOL", "http")
+
+
+def get_gms_token() -> Optional[str]:
+    """Authentication token for GMS."""
+    return os.getenv("DATAHUB_GMS_TOKEN")
+
+
+def get_system_client_id() -> Optional[str]:
+    """System client ID for OAuth/auth."""
+    return os.getenv("DATAHUB_SYSTEM_CLIENT_ID")
+
+
+def get_system_client_secret() -> Optional[str]:
+    """System client secret for OAuth/auth."""
+    return os.getenv("DATAHUB_SYSTEM_CLIENT_SECRET")
+
+
+def get_skip_config() -> bool:
+    """Skip loading config file (forces env variables)."""
+    return os.getenv("DATAHUB_SKIP_CONFIG", "").lower() == "true"
+
+
+def get_gms_base_path() -> str:
+    """Base path for GMS API endpoints."""
+    return os.getenv("DATAHUB_GMS_BASE_PATH", "")
+
+
+# ============================================================================
+# REST Emitter Configuration
+# ============================================================================
+
+
+def get_rest_emitter_default_retry_max_times() -> str:
+    """Max retry attempts for failed requests."""
+    return os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
+
+
+def get_rest_emitter_batch_max_payload_bytes() -> int:
+    """Maximum payload size in bytes for batch operations."""
+    return int(
+        os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_BYTES", str(15 * 1024 * 1024))
+    )
+
+
+def get_rest_emitter_batch_max_payload_length() -> int:
+    """Maximum number of MCPs per batch."""
+    return int(os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_LENGTH", "200"))
+
+
+def get_emit_mode() -> Optional[str]:
+    """Emission mode (SYNC_PRIMARY, SYNC_WAIT, ASYNC, ASYNC_WAIT)."""
+    return os.getenv("DATAHUB_EMIT_MODE")
+
+
+def get_rest_emitter_default_endpoint() -> Optional[str]:
+    """REST endpoint type (RESTLI or OPENAPI)."""
+    return os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT")
+
+
+def get_emitter_trace() -> bool:
+    """Enable detailed emitter tracing."""
+    return os.getenv("DATAHUB_EMITTER_TRACE", "").lower() == "true"
+
+
+# ============================================================================
+# REST Sink Configuration
+# ============================================================================
+
+
+def get_rest_sink_default_max_threads() -> int:
+    """Max thread pool size for async operations."""
+    return int(os.getenv("DATAHUB_REST_SINK_DEFAULT_MAX_THREADS", "15"))
+
+
+def get_rest_sink_default_mode() -> Optional[str]:
+    """Sink mode (SYNC, ASYNC, ASYNC_BATCH)."""
+    return os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE")
+
+
+# ============================================================================
+# Telemetry & Monitoring
+# ============================================================================
+
+
+def get_telemetry_timeout() -> str:
+    """Telemetry timeout in seconds."""
+    return os.getenv("DATAHUB_TELEMETRY_TIMEOUT", "10")
+
+
+def get_sentry_dsn() -> Optional[str]:
+    """Sentry error tracking DSN."""
+    return os.getenv("SENTRY_DSN")
+
+
+def get_sentry_environment() -> str:
+    """Sentry environment (dev/prod)."""
+    return os.getenv("SENTRY_ENVIRONMENT", "dev")
+
+
+# ============================================================================
+# Logging & Debug Configuration
+# ============================================================================
+
+
+def get_suppress_logging_manager() -> Optional[str]:
+    """Suppress DataHub logging manager initialization."""
+    return os.getenv("DATAHUB_SUPPRESS_LOGGING_MANAGER")
+
+
+def get_no_color() -> bool:
+    """Disable colored logging output."""
+    return os.getenv("NO_COLOR", "").lower() == "true"
+
+
+def get_test_mode() -> Optional[str]:
+    """Indicates running in test context."""
+    return os.getenv("DATAHUB_TEST_MODE")
+
+
+def get_debug() -> bool:
+    """Enable debug mode."""
+    return os.getenv("DATAHUB_DEBUG", "").lower() == "true"
+
+
+# ============================================================================
+# Data Processing Configuration
+# ============================================================================
+
+
+def get_sql_agg_query_log() -> str:
+    """SQL aggregator query logging level."""
+    return os.getenv("DATAHUB_SQL_AGG_QUERY_LOG", "DISABLED")
+
+
+def get_dataset_urn_to_lower() -> str:
+    """Convert dataset URNs to lowercase."""
+    return os.getenv("DATAHUB_DATASET_URN_TO_LOWER", "false")
+
+
+# ============================================================================
+# Integration-Specific Configuration
+# ============================================================================
+
+
+def get_kafka_schema_registry_url() -> Optional[str]:
+    """Kafka schema registry URL."""
+    return os.getenv("KAFKA_SCHEMAREGISTRY_URL")
+
+
+def get_spark_version() -> Optional[str]:
+    """Spark version (for S3 source)."""
+    return os.getenv("SPARK_VERSION")
+
+
+def get_bigquery_schema_parallelism() -> int:
+    """Parallelism level for BigQuery schema extraction."""
+    return int(os.getenv("DATAHUB_BIGQUERY_SCHEMA_PARALLELISM", "20"))
+
+
+def get_snowflake_schema_parallelism() -> int:
+    """Parallelism level for Snowflake schema extraction."""
+    return int(os.getenv("DATAHUB_SNOWFLAKE_SCHEMA_PARALLELISM", "20"))
+
+
+def get_powerbi_m_query_parse_timeout() -> int:
+    """Timeout for PowerBI M query parsing."""
+    return int(os.getenv("DATAHUB_POWERBI_M_QUERY_PARSE_TIMEOUT", "60"))
+
+
+def get_trace_powerbi_mquery_parser() -> bool:
+    """Enable PowerBI M query parser tracing."""
+    return os.getenv("DATAHUB_TRACE_POWERBI_MQUERY_PARSER", "").lower() == "true"
+
+
+def get_lookml_git_test_ssh_key() -> Optional[str]:
+    """SSH key for LookML Git tests."""
+    return os.getenv("DATAHUB_LOOKML_GIT_TEST_SSH_KEY")
+
+
+# ============================================================================
+# AWS/Cloud Configuration
+# ============================================================================
+
+
+def get_aws_lambda_function_name() -> Optional[str]:
+    """Indicates running in AWS Lambda."""
+    return os.getenv("AWS_LAMBDA_FUNCTION_NAME")
+
+
+def get_aws_execution_env() -> Optional[str]:
+    """AWS execution environment."""
+    return os.getenv("AWS_EXECUTION_ENV")
+
+
+def get_aws_web_identity_token_file() -> Optional[str]:
+    """OIDC token file path."""
+    return os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE")
+
+
+def get_aws_role_arn() -> Optional[str]:
+    """AWS role ARN for OIDC."""
+    return os.getenv("AWS_ROLE_ARN")
+
+
+def get_aws_app_runner_service_id() -> Optional[str]:
+    """AWS App Runner service ID."""
+    return os.getenv("AWS_APP_RUNNER_SERVICE_ID")
+
+
+def get_ecs_container_metadata_uri_v4() -> Optional[str]:
+    """ECS metadata endpoint v4."""
+    return os.getenv("ECS_CONTAINER_METADATA_URI_V4")
+
+
+def get_ecs_container_metadata_uri() -> Optional[str]:
+    """ECS metadata endpoint v3."""
+    return os.getenv("ECS_CONTAINER_METADATA_URI")
+
+
+def get_elastic_beanstalk_environment_name() -> Optional[str]:
+    """Elastic Beanstalk environment."""
+    return os.getenv("ELASTIC_BEANSTALK_ENVIRONMENT_NAME")
+
+
+# ============================================================================
+# Docker & Local Development
+# ============================================================================
+
+
+def get_compose_project_name() -> str:
+    """Docker Compose project name."""
+    return os.getenv("DATAHUB_COMPOSE_PROJECT_NAME", "datahub")
+
+
+def get_docker_compose_base() -> Optional[str]:
+    """Base path for Docker Compose files."""
+    return os.getenv("DOCKER_COMPOSE_BASE")
+
+
+def get_datahub_version() -> Optional[str]:
+    """DataHub version (set during docker init)."""
+    return os.getenv("DATAHUB_VERSION")
+
+
+def get_mapped_mysql_port() -> Optional[str]:
+    """MySQL port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_MYSQL_PORT")
+
+
+def get_mapped_kafka_broker_port() -> Optional[str]:
+    """Kafka broker port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_KAFKA_BROKER_PORT")
+
+
+def get_mapped_elastic_port() -> Optional[str]:
+    """Elasticsearch port mapping (set during docker init)."""
+    return os.getenv("DATAHUB_MAPPED_ELASTIC_PORT")
+
+
+def get_metadata_service_auth_enabled() -> str:
+    """Enable/disable auth in Docker."""
+    return os.getenv("METADATA_SERVICE_AUTH_ENABLED", "false")
+
+
+def get_ui_ingestion_default_cli_version() -> Optional[str]:
+    """CLI version for UI ingestion (set during init)."""
+    return os.getenv("UI_INGESTION_DEFAULT_CLI_VERSION")
+
+
+# ============================================================================
+# Utility & Helper Configuration
+# ============================================================================
+
+
+def get_datahub_component() -> str:
+    """Component name for user agent tracking."""
+    return os.getenv("DATAHUB_COMPONENT", "datahub")
+
+
+def get_force_local_quickstart_mapping() -> str:
+    """Force local quickstart mapping file."""
+    return os.getenv("FORCE_LOCAL_QUICKSTART_MAPPING", "")
+
+
+def get_dataproduct_external_url() -> Optional[str]:
+    """External URL for data products."""
+    return os.getenv("DATAHUB_DATAPRODUCT_EXTERNAL_URL")
+
+
+def get_override_sqlite_version_req() -> str:
+    """Override SQLite version requirement."""
+    return os.getenv("OVERRIDE_SQLITE_VERSION_REQ", "")
+
+
+def get_update_entity_registry() -> str:
+    """Update entity registry during tests."""
+    return os.getenv("UPDATE_ENTITY_REGISTRY", "false")
+
+
+def get_ci() -> Optional[str]:
+    """Indicates running in CI environment."""
+    return os.getenv("CI")

datahub/configuration/import_resolver.py

@@ -1,15 +1,18 @@
-from typing import TypeVar, Union
+from typing import TYPE_CHECKING, Type, TypeVar, Union
 
 import pydantic
 
 from datahub.ingestion.api.registry import import_path
 
-T = TypeVar("T")
+if TYPE_CHECKING:
+    from pydantic.deprecated.class_validators import V1Validator
 
+_T = TypeVar("_T")
 
-def _pydantic_resolver(v: Union[T, str]) -> T:
+
+def _pydantic_resolver(cls: Type, v: Union[str, _T]) -> _T:
     return import_path(v) if isinstance(v, str) else v
 
 
-def pydantic_resolve_key(field: str) -> classmethod:
+def pydantic_resolve_key(field: str) -> "V1Validator":
     return pydantic.validator(field, pre=True, allow_reuse=True)(_pydantic_resolver)