acryl-datahub 1.3.0.1rc2__py3-none-any.whl → 1.3.0.1rc4__py3-none-any.whl

datahub/cli/quickstart_versioning.py

@@ -1,6 +1,5 @@
 import json
 import logging
-import os
 import os.path
 import re
 from typing import Dict, Optional
@@ -13,10 +12,11 @@ from packaging.version import parse
 from pydantic import BaseModel
 
 from datahub._version import nice_version_name
+from datahub.configuration.env_vars import get_force_local_quickstart_mapping
 
 logger = logging.getLogger(__name__)
 
-LOCAL_QUICKSTART_MAPPING_FILE = os.environ.get("FORCE_LOCAL_QUICKSTART_MAPPING", "")
+LOCAL_QUICKSTART_MAPPING_FILE = get_force_local_quickstart_mapping()
 DEFAULT_LOCAL_CONFIG_PATH = "~/.datahub/quickstart/quickstart_version_mapping.yaml"
 DEFAULT_REMOTE_CONFIG_PATH = "https://raw.githubusercontent.com/datahub-project/datahub/master/docker/quickstart/quickstart_version_mapping.yaml"
 

datahub/cli/specific/dataproduct_cli.py

@@ -1,7 +1,6 @@
 import difflib
 import json
 import logging
-import os
 import pathlib
 import sys
 from pathlib import Path
@@ -14,6 +13,7 @@ from click_default_group import DefaultGroup
 
 from datahub.api.entities.dataproduct.dataproduct import DataProduct
 from datahub.cli.specific.file_loader import load_file
+from datahub.configuration.env_vars import get_dataproduct_external_url
 from datahub.emitter.mce_builder import (
     make_group_urn,
     make_user_urn,
@@ -84,9 +84,7 @@ def mutate(file: Path, validate_assets: bool, external_url: str, upsert: bool) -
     with get_default_graph(ClientMode.CLI) as graph:
         data_product: DataProduct = DataProduct.from_yaml(file, graph)
         external_url_override = (
-            external_url
-            or os.getenv("DATAHUB_DATAPRODUCT_EXTERNAL_URL")
-            or data_product.external_url
+            external_url or get_dataproduct_external_url() or data_product.external_url
         )
         data_product.external_url = external_url_override
         if upsert and not graph.exists(data_product.urn):

datahub/cli/specific/user_cli.py

@@ -1,13 +1,15 @@
 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.upgrade import upgrade
 
@@ -55,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/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/kafka.py

@@ -1,19 +1,21 @@
-import os
-
 from pydantic import Field, validator
 
 from datahub.configuration.common import ConfigModel, ConfigurationError
+from datahub.configuration.env_vars import (
+    get_gms_base_path,
+    get_kafka_schema_registry_url,
+)
 from datahub.configuration.kafka_consumer_config import CallableConsumerConfig
 from datahub.configuration.validate_host_port import validate_host_port
 
 
 def _get_schema_registry_url() -> str:
     """Get schema registry URL with proper base path handling."""
-    explicit_url = os.getenv("KAFKA_SCHEMAREGISTRY_URL")
+    explicit_url = get_kafka_schema_registry_url()
     if explicit_url:
         return explicit_url
 
-    base_path = os.getenv("DATAHUB_GMS_BASE_PATH", "")
+    base_path = get_gms_base_path()
     if base_path in ("/", ""):
         base_path = ""
 

datahub/emitter/mce_builder.py

@@ -3,7 +3,6 @@
 import hashlib
 import json
 import logging
-import os
 import re
 import time
 from datetime import datetime, timezone
@@ -26,6 +25,7 @@ import typing_inspect
 from avrogen.dict_wrapper import DictWrapper
 from typing_extensions import assert_never
 
+from datahub.configuration.env_vars import get_dataset_urn_to_lower
 from datahub.emitter.enum_helpers import get_enum_options
 from datahub.metadata.schema_classes import (
     AssertionKeyClass,
@@ -72,9 +72,7 @@ ALL_ENV_TYPES: Set[str] = set(get_enum_options(FabricTypeClass))
 
 DEFAULT_FLOW_CLUSTER = "prod"
 UNKNOWN_USER = "urn:li:corpuser:unknown"
-DATASET_URN_TO_LOWER: bool = (
-    os.getenv("DATAHUB_DATASET_URN_TO_LOWER", "false") == "true"
-)
+DATASET_URN_TO_LOWER: bool = get_dataset_urn_to_lower() == "true"
 
 if TYPE_CHECKING:
     from datahub.emitter.mcp_builder import DatahubKey