acryl-datahub 1.0.0.2rc4__py3-none-any.whl → 1.0.0.3__py3-none-any.whl

datahub/cli/put_cli.py

@@ -8,6 +8,7 @@ from datahub.cli.cli_utils import post_entity
 from datahub.configuration.config_loader import load_config_file
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import (
     DataPlatformInfoClass as DataPlatformInfo,
     PlatformTypeClass,
@@ -53,7 +54,7 @@ def aspect(urn: str, aspect: str, aspect_data: str, run_id: Optional[str]) -> No
         aspect_data, allow_stdin=True, resolve_env_vars=False, process_directives=False
     )
 
-    client = get_default_graph()
+    client = get_default_graph(ClientMode.CLI)
 
     system_metadata: Union[None, SystemMetadataClass] = None
     if run_id:
@@ -118,7 +119,7 @@ def platform(
         displayName=display_name or platform_name,
         logoUrl=logo,
     )
-    datahub_graph = get_default_graph()
+    datahub_graph = get_default_graph(ClientMode.CLI)
     mcp = MetadataChangeProposalWrapper(
         entityUrn=str(platform_urn),
         aspect=data_platform_info,

datahub/cli/specific/assertions_cli.py

@@ -15,6 +15,7 @@ from datahub.api.entities.assertion.compiler_interface import (
 from datahub.emitter.mce_builder import make_assertion_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.integrations.assertion.registry import ASSERTION_PLATFORMS
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -39,7 +40,7 @@ def upsert(file: str) -> None:
 
     assertions_spec: AssertionsConfigSpec = AssertionsConfigSpec.from_yaml(file)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         for assertion_spec in assertions_spec.assertions:
             try:
                 mcp = MetadataChangeProposalWrapper(

datahub/cli/specific/datacontract_cli.py

@@ -6,6 +6,7 @@ from click_default_group import DefaultGroup
 
 from datahub.api.entities.datacontract.datacontract import DataContract
 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
 
@@ -28,7 +29,7 @@ def upsert(file: str) -> None:
     data_contract: DataContract = DataContract.from_yaml(file)
     urn = data_contract.urn
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if not graph.exists(data_contract.entity):
             raise ValueError(
                 f"Cannot define a data contract for non-existent entity {data_contract.entity}"
@@ -72,7 +73,7 @@ def delete(urn: Optional[str], file: Optional[str], hard: bool) -> None:
         data_contract = DataContract.from_yaml(file)
         urn = data_contract.urn
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if not graph.exists(urn):
             raise ValueError(f"Data Contract {urn} does not exist")
 

datahub/cli/specific/dataproduct_cli.py

@@ -20,6 +20,7 @@ from datahub.emitter.mce_builder import (
     validate_ownership_type,
 )
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.schema_classes import OwnerClass, OwnershipTypeClass
 from datahub.specific.dataproduct import DataProductPatchBuilder
 from datahub.telemetry import telemetry
@@ -81,7 +82,7 @@ def mutate(file: Path, validate_assets: bool, external_url: str, upsert: bool) -
 
     config_dict = load_file(pathlib.Path(file))
     id = config_dict.get("id") if isinstance(config_dict, dict) else None
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product: DataProduct = DataProduct.from_yaml(file, graph)
         external_url_override = (
             external_url
@@ -162,7 +163,7 @@ def upsert(file: Path, validate_assets: bool, external_url: str) -> None:
 def diff(file: Path, update: bool) -> None:
     """Diff a Data Product file with its twin in DataHub"""
 
-    with get_default_graph() as emitter:
+    with get_default_graph(ClientMode.CLI) as emitter:
         id: Optional[str] = None
         try:
             data_product_local: DataProduct = DataProduct.from_yaml(file, emitter)
@@ -216,7 +217,7 @@ def delete(urn: str, file: Path, hard: bool) -> None:
         raise click.Abort()
 
     graph: DataHubGraph
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         data_product_urn = (
             urn if urn.startswith("urn:li:dataProduct") else f"urn:li:dataProduct:{urn}"
         )
@@ -248,7 +249,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataProduct:"):
         urn = f"urn:li:dataProduct:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataproduct: DataProduct = DataProduct.from_datahub(graph=graph, id=urn)
             click.secho(
@@ -306,7 +307,7 @@ def set_description(urn: str, description: str, md_file: Path) -> None:
 
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.set_description(description)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "set description")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -342,7 +343,7 @@ def add_owner(urn: str, owner: str, owner_type: str) -> None:
             owner=_get_owner_urn(owner), type=owner_type, typeUrn=owner_type_urn
         )
     )
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add owners")
         for mcp in dataproduct_patcher.build():
             graph.emit(mcp)
@@ -360,7 +361,7 @@ def remove_owner(urn: str, owner_urn: str) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_owner(owner=_get_owner_urn(owner_urn))
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove owners")
         for mcp in dataproduct_patcher.build():
             click.echo(json.dumps(mcp.to_obj()))
@@ -382,7 +383,7 @@ def add_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.add_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "add assets")
         if validate_assets:
             _abort_if_non_existent_urn(
@@ -409,7 +410,7 @@ def remove_asset(urn: str, asset: str, validate_assets: bool) -> None:
         urn = f"urn:li:dataProduct:{urn}"
     dataproduct_patcher: DataProductPatchBuilder = DataProduct.get_patch_builder(urn)
     dataproduct_patcher.remove_asset(asset)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         _abort_if_non_existent_urn(graph, urn, "remove assets")
         if validate_assets:
             _abort_if_non_existent_urn(

datahub/cli/specific/dataset_cli.py

@@ -12,6 +12,7 @@ from click_default_group import DefaultGroup
 from datahub.api.entities.dataset.dataset import Dataset, DatasetRetrievalConfig
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
@@ -54,7 +55,7 @@ def get(urn: str, to_file: str) -> None:
     if not urn.startswith("urn:li:dataset:"):
         urn = f"urn:li:dataset:{urn}"
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             dataset: Dataset = Dataset.from_datahub(graph=graph, urn=urn)
             click.secho(
@@ -82,7 +83,7 @@ def add_sibling(urn: str, sibling_urns: Tuple[str]) -> None:
     all_urns.add(urn)
     for sibling_urn in sibling_urns:
         all_urns.add(sibling_urn)
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         for _urn in all_urns:
             _emit_sibling(graph, urn, _urn, all_urns)
 
@@ -181,7 +182,7 @@ def sync(file: str, to_datahub: bool, dry_run: bool) -> None:
     dry_run_prefix = "[dry-run]: " if dry_run else ""  # prefix to use in messages
 
     failures: List[str] = []
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         datasets = Dataset.from_yaml(file)
         for dataset in datasets:
             assert (

datahub/cli/specific/forms_cli.py

@@ -7,6 +7,7 @@ from click_default_group import DefaultGroup
 
 from datahub.api.entities.forms.forms import Forms
 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
 
@@ -40,7 +41,7 @@ def upsert(file: Path) -> None:
 @telemetry.with_telemetry()
 def get(urn: str, to_file: str) -> None:
     """Get form from DataHub"""
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             form: Forms = Forms.from_datahub(graph=graph, urn=urn)
             click.secho(

datahub/cli/specific/group_cli.py

@@ -10,6 +10,7 @@ from datahub.api.entities.corpgroup.corpgroup import (
 )
 from datahub.cli.specific.file_loader import load_file
 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
 
@@ -40,7 +41,7 @@ def upsert(file: Path, override_editable: bool) -> None:
 
     config_dict = load_file(file)
     group_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 group_config in group_configs:
             try:
                 datahub_group = CorpGroup.parse_obj(group_config)

datahub/cli/specific/structuredproperties_cli.py

@@ -11,6 +11,7 @@ from datahub.api.entities.structuredproperties.structuredproperties import (
     StructuredProperties,
 )
 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
 from datahub.utilities.urns.urn import Urn
@@ -33,7 +34,7 @@ def properties() -> None:
 def upsert(file: Path) -> None:
     """Upsert structured properties in DataHub."""
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         StructuredProperties.create(str(file), graph)
 
 
@@ -48,7 +49,7 @@ def get(urn: str, to_file: str) -> None:
     """Get structured properties from DataHub"""
     urn = Urn.make_structured_property_urn(urn)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if graph.exists(urn):
             structuredproperties: StructuredProperties = (
                 StructuredProperties.from_datahub(graph=graph, urn=urn)
@@ -117,7 +118,7 @@ def list(details: bool, to_file: str) -> None:
         with open(file, "w") as fp:
             yaml.dump(serialized_objects, fp)
 
-    with get_default_graph() as graph:
+    with get_default_graph(ClientMode.CLI) as graph:
         if details:
             logger.info(
                 "Listing structured properties with details. Use --no-details for urns only"

datahub/cli/specific/user_cli.py

@@ -8,6 +8,7 @@ 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.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 
@@ -38,7 +39,7 @@ def upsert(file: Path, override_editable: bool) -> None:
 
     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)

datahub/cli/state_cli.py

@@ -5,6 +5,7 @@ import click
 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
 
@@ -28,7 +29,7 @@ def inspect(pipeline_name: str, platform: str) -> None:
     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,6 +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.ingestion.graph.config import ClientMode
 from datahub.telemetry import telemetry
 from datahub.upgrade import upgrade
 from datahub.utilities.urns.urn import Urn
@@ -63,7 +64,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"):

datahub/configuration/common.py

@@ -33,10 +33,15 @@ REDACT_KEYS = {
 }
 REDACT_SUFFIXES = {
     "_password",
+    "-password",
     "_secret",
+    "-secret",
     "_token",
+    "-token",
     "_key",
+    "-key",
     "_key_id",
+    "-key-id",
 }
 
 

datahub/configuration/source_common.py

@@ -16,7 +16,7 @@ class PlatformInstanceConfigMixin(ConfigModel):
         default=None,
         description="The instance of the platform that all assets produced by this recipe belong to. "
         "This should be unique within the platform. "
-        "See https://datahubproject.io/docs/platform-instances/ for more details.",
+        "See https://docs.datahub.com/docs/platform-instances/ for more details.",
     )
 
 

datahub/emitter/mcp.py

@@ -1,9 +1,11 @@
 import dataclasses
 import json
-from typing import TYPE_CHECKING, List, Optional, Sequence, Tuple, Union
+import warnings
+from typing import TYPE_CHECKING, Dict, List, Optional, Sequence, Tuple, Union
 
 from datahub.emitter.aspect import ASPECT_MAP, JSON_CONTENT_TYPE
 from datahub.emitter.serialization_helper import post_json_transform, pre_json_transform
+from datahub.errors import DataHubDeprecationWarning
 from datahub.metadata.schema_classes import (
     ChangeTypeClass,
     DictWrapper,
@@ -69,18 +71,28 @@ class MetadataChangeProposalWrapper:
     aspectName: Union[None, str] = None
     aspect: Union[None, _Aspect] = None
     systemMetadata: Union[None, SystemMetadataClass] = None
+    headers: Union[None, Dict[str, str]] = None
 
     def __post_init__(self) -> None:
         if self.entityUrn and self.entityType == _ENTITY_TYPE_UNSET:
             self.entityType = guess_entity_type(self.entityUrn)
         elif self.entityUrn and self.entityType:
-            guessed_entity_type = guess_entity_type(self.entityUrn).lower()
-            # Entity type checking is actually case insensitive.
-            # Note that urns are case sensitive, but entity types are not.
-            if self.entityType.lower() != guessed_entity_type:
+            guessed_entity_type = guess_entity_type(self.entityUrn)
+            if self.entityType.lower() != guessed_entity_type.lower():
+                # If they aren't a case-ignored match, raise an error.
                 raise ValueError(
                     f"entityType {self.entityType} does not match the entity type {guessed_entity_type} from entityUrn {self.entityUrn}",
                 )
+            elif self.entityType != guessed_entity_type:
+                # If they only differ in case, normalize and print a warning.
+                self.entityType = guessed_entity_type
+                warnings.warn(
+                    f"The passed entityType {self.entityType} differs in case from the expected entity type {guessed_entity_type}. "
+                    "This will be automatically corrected for now, but will become an error in a future release. "
+                    "Note that the entityType field is optional and will be automatically inferred from the entityUrn.",
+                    DataHubDeprecationWarning,
+                    stacklevel=3,
+                )
         elif self.entityType == _ENTITY_TYPE_UNSET:
             raise ValueError("entityType must be set if entityUrn is not set")
 
@@ -112,6 +124,7 @@ class MetadataChangeProposalWrapper:
             auditHeader=self.auditHeader,
             aspectName=self.aspectName,
             systemMetadata=self.systemMetadata,
+            headers=self.headers,
         )
 
     def make_mcp(self) -> MetadataChangeProposalClass:
@@ -211,6 +224,7 @@ class MetadataChangeProposalWrapper:
                 aspectName=mcpc.aspectName,
                 aspect=aspect,
                 systemMetadata=mcpc.systemMetadata,
+                headers=mcpc.headers,
             )
         else:
             return None
@@ -228,6 +242,7 @@ class MetadataChangeProposalWrapper:
             changeType=mcl.changeType,
             auditHeader=mcl.auditHeader,
             systemMetadata=mcl.systemMetadata,
+            headers=mcl.headers,
         )
         return cls.try_from_mcpc(mcpc) or mcpc
 

datahub/emitter/request_helper.py

@@ -1,14 +1,31 @@
+import json
 import shlex
-from typing import List, Optional, Union
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Union
 
 import requests
 from requests.auth import HTTPBasicAuth
 
+from datahub.emitter.aspect import JSON_CONTENT_TYPE, JSON_PATCH_CONTENT_TYPE
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.serialization_helper import pre_json_transform
+from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
+    MetadataChangeProposal,
+)
+from datahub.metadata.schema_classes import ChangeTypeClass
+
+
+def _decode_bytes(value: Union[str, bytes]) -> str:
+    """Decode bytes to string, if necessary."""
+    if isinstance(value, bytes):
+        return value.decode()
+    return value
+
 
 def _format_header(name: str, value: Union[str, bytes]) -> str:
     if name == "Authorization":
         return f"{name!s}: <redacted>"
-    return f"{name!s}: {value!s}"
+    return f"{name!s}: {_decode_bytes(value)}"
 
 
 def make_curl_command(
@@ -21,7 +38,9 @@ def make_curl_command(
 
     if session.auth:
         if isinstance(session.auth, HTTPBasicAuth):
-            fragments.extend(["-u", f"{session.auth.username}:<redacted>"])
+            fragments.extend(
+                ["-u", f"{_decode_bytes(session.auth.username)}:<redacted>"]
+            )
         else:
             # For other auth types, they should be handled via headers
             fragments.extend(["-H", "<unknown auth type>"])
@@ -31,3 +50,97 @@ def make_curl_command(
 
     fragments.append(url)
     return shlex.join(fragments)
+
+
+@dataclass
+class OpenApiRequest:
+    """Represents an OpenAPI request for entity operations."""
+
+    method: str
+    url: str
+    payload: List[Dict[str, Any]]
+
+    @classmethod
+    def from_mcp(
+        cls,
+        mcp: Union[MetadataChangeProposal, MetadataChangeProposalWrapper],
+        gms_server: str,
+        async_flag: Optional[bool] = None,
+        async_default: bool = False,
+    ) -> Optional["OpenApiRequest"]:
+        """Factory method to create an OpenApiRequest from a MetadataChangeProposal."""
+        if not mcp.aspectName or (
+            mcp.changeType != ChangeTypeClass.DELETE and not mcp.aspect
+        ):
+            return None
+
+        resolved_async_flag = async_flag if async_flag is not None else async_default
+
+        method = "post"
+        url = f"{gms_server}/openapi/v3/entity/{mcp.entityType}?async={'true' if resolved_async_flag else 'false'}"
+        payload = []
+
+        if mcp.changeType == ChangeTypeClass.DELETE:
+            method = "delete"
+            url = f"{gms_server}/openapi/v3/entity/{mcp.entityType}/{mcp.entityUrn}"
+        else:
+            if mcp.aspect:
+                if mcp.changeType == ChangeTypeClass.PATCH:
+                    method = "patch"
+                    obj = mcp.aspect.to_obj()
+                    content_type = obj.get("contentType")
+                    if obj.get("value") and content_type == JSON_PATCH_CONTENT_TYPE:
+                        # Undo double serialization.
+                        obj = json.loads(obj["value"])
+                        patch_value = obj
+                    else:
+                        raise NotImplementedError(
+                            f"ChangeType {mcp.changeType} only supports context type {JSON_PATCH_CONTENT_TYPE}, found {content_type}."
+                        )
+
+                    if isinstance(patch_value, list):
+                        patch_value = {"patch": patch_value}
+
+                    payload = [
+                        {
+                            "urn": mcp.entityUrn,
+                            mcp.aspectName: {
+                                "value": patch_value,
+                                "systemMetadata": mcp.systemMetadata.to_obj()
+                                if mcp.systemMetadata
+                                else None,
+                            },
+                        }
+                    ]
+                else:
+                    if isinstance(mcp, MetadataChangeProposalWrapper):
+                        aspect_value = pre_json_transform(
+                            mcp.to_obj(simplified_structure=True)
+                        )["aspect"]["json"]
+                    else:
+                        obj = mcp.aspect.to_obj()
+                        content_type = obj.get("contentType")
+                        if obj.get("value") and content_type == JSON_CONTENT_TYPE:
+                            # Undo double serialization.
+                            obj = json.loads(obj["value"])
+                        elif content_type == JSON_PATCH_CONTENT_TYPE:
+                            raise NotImplementedError(
+                                f"ChangeType {mcp.changeType} does not support patch."
+                            )
+                        aspect_value = pre_json_transform(obj)
+
+                    payload = [
+                        {
+                            "urn": mcp.entityUrn,
+                            mcp.aspectName: {
+                                "value": aspect_value,
+                                "systemMetadata": mcp.systemMetadata.to_obj()
+                                if mcp.systemMetadata
+                                else None,
+                            },
+                        }
+                    ]
+            else:
+                raise ValueError(f"ChangeType {mcp.changeType} requires a value.")
+
+        return cls(method=method, url=url, payload=payload)