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

datahub/emitter/rest_emitter.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 import functools
 import json
 import logging
-import os
 import re
 import time
 from collections import defaultdict
@@ -33,7 +32,6 @@ from typing_extensions import deprecated
 from datahub._version import nice_version_name
 from datahub.cli import config_utils
 from datahub.cli.cli_utils import ensure_has_system_metadata, fixup_gms_url, get_or_else
-from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.configuration.common import (
     ConfigEnum,
     ConfigModel,
@@ -42,6 +40,14 @@ from datahub.configuration.common import (
     TraceTimeoutError,
     TraceValidationError,
 )
+from datahub.configuration.env_vars import (
+    get_emit_mode,
+    get_emitter_trace,
+    get_rest_emitter_batch_max_payload_bytes,
+    get_rest_emitter_batch_max_payload_length,
+    get_rest_emitter_default_endpoint,
+    get_rest_emitter_default_retry_max_times,
+)
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import OpenApiRequest, make_curl_command
@@ -82,11 +88,9 @@ _DEFAULT_RETRY_STATUS_CODES = [  # Additional status codes to retry on
     504,
 ]
 _DEFAULT_RETRY_METHODS = ["HEAD", "GET", "POST", "PUT", "DELETE", "OPTIONS", "TRACE"]
-_DEFAULT_RETRY_MAX_TIMES = int(
-    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_RETRY_MAX_TIMES", "4")
-)
+_DEFAULT_RETRY_MAX_TIMES = int(get_rest_emitter_default_retry_max_times())
 
-_DATAHUB_EMITTER_TRACE = get_boolean_env_variable("DATAHUB_EMITTER_TRACE", False)
+_DATAHUB_EMITTER_TRACE = get_emitter_trace()
 
 _DEFAULT_CLIENT_MODE: ClientMode = ClientMode.SDK
 
@@ -98,17 +102,13 @@ TRACE_BACKOFF_FACTOR = 2.0  # Double the wait time each attempt
 # The limit is 16,000,000 bytes. We will use a max of 15mb to have some space
 # for overhead like request headers.
 # This applies to pretty much all calls to GMS.
-INGEST_MAX_PAYLOAD_BYTES = int(
-    os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_BYTES", 15 * 1024 * 1024)
-)
+INGEST_MAX_PAYLOAD_BYTES = get_rest_emitter_batch_max_payload_bytes()
 
 # This limit is somewhat arbitrary. All GMS endpoints will timeout
 # and return a 500 if processing takes too long. To avoid sending
 # too much to the backend and hitting a timeout, we try to limit
 # the number of MCPs we send in a batch.
-BATCH_INGEST_MAX_PAYLOAD_LENGTH = int(
-    os.getenv("DATAHUB_REST_EMITTER_BATCH_MAX_PAYLOAD_LENGTH", 200)
-)
+BATCH_INGEST_MAX_PAYLOAD_LENGTH = get_rest_emitter_batch_max_payload_length()
 
 
 def preserve_unicode_escapes(obj: Any) -> Any:
@@ -147,7 +147,7 @@ class EmitMode(ConfigEnum):
 
 _DEFAULT_EMIT_MODE = pydantic.parse_obj_as(
     EmitMode,
-    os.getenv("DATAHUB_EMIT_MODE", EmitMode.SYNC_PRIMARY),
+    get_emit_mode() or EmitMode.SYNC_PRIMARY,
 )
 
 
@@ -158,7 +158,7 @@ class RestSinkEndpoint(ConfigEnum):
 
 DEFAULT_REST_EMITTER_ENDPOINT = pydantic.parse_obj_as(
     RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
+    get_rest_emitter_default_endpoint() or RestSinkEndpoint.RESTLI,
 )
 
 
@@ -478,7 +478,7 @@ class DataHubRestEmitter(Closeable, Emitter):
         if self._openapi_ingestion is None:
             # No constructor parameter
             if (
-                not os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT")
+                not get_rest_emitter_default_endpoint()
                 and self._session_config.client_mode == ClientMode.SDK
                 and self._server_config.supports_feature(ServiceFeature.OPEN_API_SDK)
             ):

datahub/entrypoints.py

@@ -22,6 +22,7 @@ from datahub.cli.docker_cli import docker
 from datahub.cli.env_utils import get_boolean_env_variable
 from datahub.cli.exists_cli import exists
 from datahub.cli.get_cli import get
+from datahub.cli.graphql_cli import graphql
 from datahub.cli.ingest_cli import ingest
 from datahub.cli.migrate import migrate
 from datahub.cli.put_cli import put
@@ -169,6 +170,7 @@ datahub.add_command(ingest)
 datahub.add_command(delete)
 datahub.add_command(exists)
 datahub.add_command(get)
+datahub.add_command(graphql)
 datahub.add_command(put)
 datahub.add_command(state)
 datahub.add_command(telemetry_cli)

datahub/ingestion/api/auto_work_units/auto_validate_input_fields.py

@@ -0,0 +1,87 @@
+import logging
+from typing import TYPE_CHECKING, Iterable, List
+
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.metadata.schema_classes import InputFieldClass, InputFieldsClass
+
+if TYPE_CHECKING:
+    from datahub.ingestion.api.source import SourceReport
+
+logger = logging.getLogger(__name__)
+
+
+class ValidateInputFieldsProcessor:
+    def __init__(self, report: "SourceReport"):
+        self.report = report
+
+    def validate_input_fields(
+        self,
+        stream: Iterable[MetadataWorkUnit],
+    ) -> Iterable[MetadataWorkUnit]:
+        """
+        Validate input fields and filter out invalid ones.
+
+        Invalid input fields have empty or missing fieldPath values, which would cause
+        URN generation to fail when sent to the server. This processor filters them out
+        and reports them as warnings.
+        """
+        for wu in stream:
+            input_fields_aspect = wu.get_aspect_of_type(InputFieldsClass)
+            if input_fields_aspect and input_fields_aspect.fields:
+                valid_fields: List[InputFieldClass] = []
+                invalid_count = 0
+
+                for input_field in input_fields_aspect.fields:
+                    if (
+                        input_field.schemaField
+                        and input_field.schemaField.fieldPath
+                        and input_field.schemaField.fieldPath.strip()
+                    ):
+                        valid_fields.append(input_field)
+                    else:
+                        invalid_count += 1
+
+                if invalid_count > 0:
+                    logger.debug(
+                        f"Filtered {invalid_count} invalid input field(s) with empty fieldPath for {wu.get_urn()}"
+                    )
+                    self.report.num_input_fields_filtered += invalid_count
+                    self.report.warning(
+                        title="Invalid input fields filtered",
+                        message="Input fields with empty fieldPath values were filtered out to prevent ingestion errors",
+                        context=f"Filtered {invalid_count} invalid input field(s) for {wu.get_urn()}",
+                    )
+
+                    # Update the aspect with only valid fields
+                    if valid_fields:
+                        input_fields_aspect.fields = valid_fields
+                    else:
+                        # If no valid fields remain, skip this workunit entirely
+                        logger.debug(
+                            f"All input fields were invalid for {wu.get_urn()}, skipping InputFieldsClass workunit"
+                        )
+                        # Don't yield this workunit
+                        continue
+
+            yield wu
+
+    def _remove_input_fields_aspect(self, wu: MetadataWorkUnit) -> MetadataWorkUnit:
+        """Remove InputFieldsClass aspect from a workunit."""
+        # For MCPs, we can simply not yield the aspect
+        # For MCEs, we need to remove it from the snapshot
+        if hasattr(wu.metadata, "aspect") and isinstance(
+            wu.metadata.aspect, InputFieldsClass
+        ):
+            # This is an MCP with InputFieldsClass, skip it
+            return wu
+
+        if hasattr(wu.metadata, "proposedSnapshot"):
+            snapshot = wu.metadata.proposedSnapshot
+            if hasattr(snapshot, "aspects"):
+                snapshot.aspects = [
+                    aspect
+                    for aspect in snapshot.aspects
+                    if not isinstance(aspect, InputFieldsClass)
+                ]
+
+        return wu

datahub/ingestion/api/source.py

@@ -31,6 +31,9 @@ from datahub.ingestion.api.auto_work_units.auto_dataset_properties_aspect import
 from datahub.ingestion.api.auto_work_units.auto_ensure_aspect_size import (
     EnsureAspectSizeProcessor,
 )
+from datahub.ingestion.api.auto_work_units.auto_validate_input_fields import (
+    ValidateInputFieldsProcessor,
+)
 from datahub.ingestion.api.closeable import Closeable
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope, WorkUnit
 from datahub.ingestion.api.report import ExamplesReport, Report
@@ -215,6 +218,7 @@ class SourceReport(ExamplesReport, IngestionStageReport):
     event_not_produced_warn: bool = True
     events_produced: int = 0
     events_produced_per_sec: int = 0
+    num_input_fields_filtered: int = 0
 
     _structured_logs: StructuredLogs = field(default_factory=StructuredLogs)
 
@@ -543,6 +547,7 @@ class Source(Closeable, metaclass=ABCMeta):
             browse_path_processor,
             partial(auto_workunit_reporter, self.get_report()),
             auto_patch_last_modified,
+            ValidateInputFieldsProcessor(self.get_report()).validate_input_fields,
             EnsureAspectSizeProcessor(self.get_report()).ensure_aspect_size,
         ]
 

datahub/ingestion/graph/client.py

@@ -30,6 +30,7 @@ from typing_extensions import deprecated
 
 from datahub._codegen.aspect import _Aspect
 from datahub.cli import config_utils
+from datahub.cli.cli_utils import guess_frontend_url_from_gms_url
 from datahub.configuration.common import ConfigModel, GraphError, OperationalError
 from datahub.emitter.aspect import TIMESERIES_ASPECT_MAP
 from datahub.emitter.mce_builder import DEFAULT_ENV, Aspect
@@ -2071,6 +2072,202 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
 
         return res["reportAssertionResult"]
 
+    def _get_invite_token(self) -> str:
+        """
+        Retrieve an invite token for user creation.
+
+        Returns:
+            Invite token string
+
+        Raises:
+            OperationalError: If invite token retrieval fails
+        """
+        get_invite_token_query = """
+            query getInviteToken($input: GetInviteTokenInput!) {
+                getInviteToken(input: $input) {
+                    inviteToken
+                }
+            }
+        """
+
+        try:
+            invite_token_response = self.execute_graphql(
+                query=get_invite_token_query,
+                variables={"input": {}},
+            )
+            invite_token = invite_token_response.get("getInviteToken", {}).get(
+                "inviteToken"
+            )
+            if not invite_token:
+                raise OperationalError(
+                    "Failed to retrieve invite token. Ensure you have admin permissions.",
+                    {},
+                )
+            return invite_token
+        except Exception as e:
+            raise OperationalError(
+                f"Failed to retrieve invite token: {str(e)}", {}
+            ) from e
+
+    def _create_user_with_token(
+        self,
+        user_urn: str,
+        email: str,
+        display_name: str,
+        password: str,
+        invite_token: str,
+    ) -> None:
+        """
+        Create a user using the signup endpoint.
+
+        Args:
+            user_urn: User URN (urn:li:corpuser:{user_id})
+            email: User's email address
+            display_name: Full display name for the user
+            password: User's password
+            invite_token: Invite token for user creation
+
+        Raises:
+            OperationalError: If user creation fails
+        """
+        frontend_url = guess_frontend_url_from_gms_url(self._gms_server)
+        signup_url = f"{frontend_url}/signUp"
+        signup_payload = {
+            "userUrn": user_urn,
+            "email": email,
+            "fullName": display_name,
+            "password": password,
+            "title": "Other",
+            "inviteToken": invite_token,
+        }
+
+        logger.debug(
+            f"Creating user with URN={user_urn}, email={email} at URL: {signup_url}"
+        )
+        logger.debug(
+            f"Signup payload: {json.dumps({**signup_payload, 'password': '***'})}"
+        )
+
+        try:
+            response = self._session.post(signup_url, json=signup_payload)
+            logger.debug(f"Response status code: {response.status_code}")
+            logger.debug(f"Response headers: {dict(response.headers)}")
+            logger.debug(f"Response content length: {len(response.text)}")
+
+            response.raise_for_status()
+
+            # The /signUp endpoint returns 200 with empty body on success
+            logger.debug("User created successfully")
+
+        except HTTPError as http_err:
+            error_details = {
+                "url": signup_url,
+                "status_code": response.status_code,
+                "response_text": response.text[:500],
+            }
+            try:
+                error_json = response.json()
+                error_details["error_response"] = error_json
+                error_msg = error_json.get("message", str(http_err))
+            except JSONDecodeError:
+                error_msg = f"HTTP {response.status_code}: {response.text[:200]}"
+
+            raise OperationalError(
+                f"Failed to create user: {error_msg}",
+                error_details,
+            ) from http_err
+        except Exception as e:
+            raise OperationalError(
+                f"Failed to create user: {str(e)}",
+                {"url": signup_url, "error_type": type(e).__name__},
+            ) from e
+
+    def _assign_role_to_user(self, user_urn: str, role: str) -> None:
+        """
+        Assign a role to a user.
+
+        Args:
+            user_urn: User URN
+            role: Role to assign (Admin, Editor, or Reader)
+
+        Raises:
+            ValueError: If role is invalid
+        """
+        normalized_role = role.capitalize()
+        valid_roles = ["Admin", "Editor", "Reader"]
+        if normalized_role not in valid_roles:
+            raise ValueError(
+                f"Invalid role '{role}'. Must be one of: {', '.join(valid_roles)}"
+            )
+
+        role_urn = f"urn:li:dataHubRole:{normalized_role}"
+
+        batch_assign_role_mutation = """
+            mutation batchAssignRole($input: BatchAssignRoleInput!) {
+                batchAssignRole(input: $input)
+            }
+        """
+
+        try:
+            self.execute_graphql(
+                query=batch_assign_role_mutation,
+                variables={"input": {"roleUrn": role_urn, "actors": [user_urn]}},
+            )
+        except Exception as e:
+            logger.warning(f"Role assignment failed for user {user_urn}: {str(e)}")
+            raise
+
+    def create_native_user(
+        self,
+        user_id: str,
+        email: str,
+        display_name: str,
+        password: str,
+        role: Optional[str] = None,
+    ) -> str:
+        """
+        Create a native DataHub user with email/password authentication.
+
+        Args:
+            user_id: User identifier (will be used in the URN)
+            email: User's email address
+            display_name: Full display name for the user
+            password: User's password
+            role: Optional role to assign (Admin, Editor, or Reader)
+
+        Returns:
+            User URN of the created user (urn:li:corpuser:{user_id})
+
+        Raises:
+            OperationalError: If user creation fails
+            ValueError: If role is invalid
+        """
+        # Validate role before creating user
+        if role:
+            normalized_role = role.capitalize()
+            valid_roles = ["Admin", "Editor", "Reader"]
+            if normalized_role not in valid_roles:
+                raise ValueError(
+                    f"Invalid role '{role}'. Must be one of: {', '.join(valid_roles)}"
+                )
+
+        user_urn = f"urn:li:corpuser:{user_id}"
+
+        invite_token = self._get_invite_token()
+        self._create_user_with_token(
+            user_urn, email, display_name, password, invite_token
+        )
+
+        if role:
+            try:
+                self._assign_role_to_user(user_urn, role)
+            except Exception as e:
+                logger.warning(
+                    f"User {email} created successfully, but role assignment failed: {str(e)}"
+                )
+
+        return user_urn
+
     def close(self) -> None:
         self._make_schema_resolver.cache_clear()
         super().close()

datahub/ingestion/graph/config.py

@@ -1,8 +1,8 @@
-import os
 from enum import Enum, auto
 from typing import Dict, List, Optional
 
 from datahub.configuration.common import ConfigModel
+from datahub.configuration.env_vars import get_datahub_component
 
 
 class ClientMode(Enum):
@@ -11,7 +11,7 @@ class ClientMode(Enum):
     SDK = auto()
 
 
-DATAHUB_COMPONENT_ENV: str = os.getenv("DATAHUB_COMPONENT", "datahub").lower()
+DATAHUB_COMPONENT_ENV: str = get_datahub_component().lower()
 
 
 class DatahubClientConfig(ConfigModel):

datahub/ingestion/sink/datahub_rest.py

@@ -3,7 +3,6 @@ import contextlib
 import dataclasses
 import functools
 import logging
-import os
 import threading
 import uuid
 from enum import auto
@@ -16,6 +15,10 @@ from datahub.configuration.common import (
     ConfigurationError,
     OperationalError,
 )
+from datahub.configuration.env_vars import (
+    get_rest_sink_default_max_threads,
+    get_rest_sink_default_mode,
+)
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.emitter.rest_emitter import (
@@ -47,9 +50,7 @@ from datahub.utilities.server_config_util import set_gms_config
 
 logger = logging.getLogger(__name__)
 
-_DEFAULT_REST_SINK_MAX_THREADS = int(
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_MAX_THREADS", 15)
-)
+_DEFAULT_REST_SINK_MAX_THREADS = get_rest_sink_default_max_threads()
 
 
 class RestSinkMode(ConfigEnum):
@@ -63,7 +64,7 @@ class RestSinkMode(ConfigEnum):
 
 
 _DEFAULT_REST_SINK_MODE = pydantic.parse_obj_as(
-    RestSinkMode, os.getenv("DATAHUB_REST_SINK_DEFAULT_MODE", RestSinkMode.ASYNC_BATCH)
+    RestSinkMode, get_rest_sink_default_mode() or RestSinkMode.ASYNC_BATCH
 )
 
 

datahub/ingestion/source/aws/aws_common.py

@@ -1,5 +1,4 @@
 import logging
-import os
 from datetime import datetime, timedelta, timezone
 from enum import Enum
 from http import HTTPStatus
@@ -17,6 +16,16 @@ from datahub.configuration.common import (
     ConfigModel,
     PermissiveConfigModel,
 )
+from datahub.configuration.env_vars import (
+    get_aws_app_runner_service_id,
+    get_aws_execution_env,
+    get_aws_lambda_function_name,
+    get_aws_role_arn,
+    get_aws_web_identity_token_file,
+    get_ecs_container_metadata_uri,
+    get_ecs_container_metadata_uri_v4,
+    get_elastic_beanstalk_environment_name,
+)
 from datahub.configuration.source_common import EnvConfigMixin
 
 logger = logging.getLogger(__name__)
@@ -100,27 +109,25 @@ def detect_aws_environment() -> AwsEnvironment:
     Order matters as some environments may have multiple indicators.
     """
     # Check Lambda first as it's most specific
-    if os.getenv("AWS_LAMBDA_FUNCTION_NAME"):
-        if os.getenv("AWS_EXECUTION_ENV", "").startswith("CloudFormation"):
+    if get_aws_lambda_function_name():
+        if (get_aws_execution_env() or "").startswith("CloudFormation"):
             return AwsEnvironment.CLOUD_FORMATION
         return AwsEnvironment.LAMBDA
 
     # Check EKS (IRSA)
-    if os.getenv("AWS_WEB_IDENTITY_TOKEN_FILE") and os.getenv("AWS_ROLE_ARN"):
+    if get_aws_web_identity_token_file() and get_aws_role_arn():
         return AwsEnvironment.EKS
 
     # Check App Runner
-    if os.getenv("AWS_APP_RUNNER_SERVICE_ID"):
+    if get_aws_app_runner_service_id():
         return AwsEnvironment.APP_RUNNER
 
     # Check ECS
-    if os.getenv("ECS_CONTAINER_METADATA_URI_V4") or os.getenv(
-        "ECS_CONTAINER_METADATA_URI"
-    ):
+    if get_ecs_container_metadata_uri_v4() or get_ecs_container_metadata_uri():
         return AwsEnvironment.ECS
 
     # Check Elastic Beanstalk
-    if os.getenv("ELASTIC_BEANSTALK_ENVIRONMENT_NAME"):
+    if get_elastic_beanstalk_environment_name():
         return AwsEnvironment.BEANSTALK
 
     if is_running_on_ec2():
@@ -155,7 +162,7 @@ def get_instance_role_arn() -> Optional[str]:
 def get_lambda_role_arn() -> Optional[str]:
     """Get the Lambda function's role ARN"""
     try:
-        function_name = os.getenv("AWS_LAMBDA_FUNCTION_NAME")
+        function_name = get_aws_lambda_function_name()
         if not function_name:
             return None
 
@@ -181,7 +188,7 @@ def get_current_identity() -> Tuple[Optional[str], Optional[str]]:
         return role_arn, AwsServicePrincipal.LAMBDA.value
 
     elif env == AwsEnvironment.EKS:
-        role_arn = os.getenv("AWS_ROLE_ARN")
+        role_arn = get_aws_role_arn()
         return role_arn, AwsServicePrincipal.EKS.value
 
     elif env == AwsEnvironment.APP_RUNNER:
@@ -194,8 +201,8 @@ def get_current_identity() -> Tuple[Optional[str], Optional[str]]:
 
     elif env == AwsEnvironment.ECS:
         try:
-            metadata_uri = os.getenv("ECS_CONTAINER_METADATA_URI_V4") or os.getenv(
-                "ECS_CONTAINER_METADATA_URI"
+            metadata_uri = (
+                get_ecs_container_metadata_uri_v4() or get_ecs_container_metadata_uri()
             )
             if metadata_uri:
                 response = requests.get(f"{metadata_uri}/task", timeout=1)

datahub/ingestion/source/bigquery_v2/bigquery_config.py

@@ -1,5 +1,4 @@
 import logging
-import os
 import re
 from copy import deepcopy
 from datetime import timedelta
@@ -8,6 +7,7 @@ from typing import Dict, List, Optional, Union
 from pydantic import Field, PositiveInt, PrivateAttr, root_validator, validator
 
 from datahub.configuration.common import AllowDenyPattern, ConfigModel, HiddenFromDocs
+from datahub.configuration.env_vars import get_bigquery_schema_parallelism
 from datahub.configuration.source_common import (
     EnvConfigMixin,
     LowerCaseDatasetUrnConfigMixin,
@@ -31,9 +31,7 @@ from datahub.ingestion.source.usage.usage_common import BaseUsageConfig
 
 logger = logging.getLogger(__name__)
 
-DEFAULT_BQ_SCHEMA_PARALLELISM = int(
-    os.getenv("DATAHUB_BIGQUERY_SCHEMA_PARALLELISM", 20)
-)
+DEFAULT_BQ_SCHEMA_PARALLELISM = get_bigquery_schema_parallelism()
 
 # Regexp for sharded tables.
 # A sharded table is a table that has a suffix of the form _yyyymmdd or yyyymmdd, where yyyymmdd is a date.

datahub/ingestion/source/grafana/models.py

@@ -92,6 +92,7 @@ class Dashboard(_GrafanaBaseModel):
         """Custom parsing to handle nested panel extraction."""
         dashboard_data = data.get("dashboard", {})
         _panel_data = dashboard_data.get("panels", [])
+        panels = []
         try:
             panels = cls.extract_panels(_panel_data)
         except Exception as e:
@@ -108,6 +109,10 @@ class Dashboard(_GrafanaBaseModel):
         if "meta" in dashboard_dict:
             del dashboard_dict["meta"]
 
+        # Handle refresh field type mismatch - convert boolean to string
+        if "refresh" in dashboard_dict and isinstance(dashboard_dict["refresh"], bool):
+            dashboard_dict["refresh"] = str(dashboard_dict["refresh"])
+
         return super().parse_obj(dashboard_dict)