acryl-datahub 0.15.0.4rc2__py3-none-any.whl → 0.15.0.5__py3-none-any.whl

datahub/cli/ingest_cli.py

@@ -12,7 +12,7 @@ import click_spinner
 from click_default_group import DefaultGroup
 from tabulate import tabulate
 
-import datahub as datahub_package
+from datahub._version import nice_version_name
 from datahub.cli import cli_utils
 from datahub.cli.config_utils import CONDENSED_DATAHUB_CONFIG_PATH
 from datahub.configuration.common import ConfigModel, GraphError
@@ -147,7 +147,7 @@ def run(
                 return ret
 
     # main function begins
-    logger.info("DataHub CLI version: %s", datahub_package.nice_version_name())
+    logger.info("DataHub CLI version: %s", nice_version_name())
 
     pipeline_config = load_config_file(
         config,

datahub/emitter/composite_emitter.py

@@ -0,0 +1,36 @@
+from typing import Callable, List, Optional, Union
+
+from datahub.emitter.generic_emitter import Emitter
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
+    MetadataChangeEvent,
+    MetadataChangeProposal,
+)
+
+
+# Experimental composite emitter that allows multiple emitters to be used in a single ingestion job
+class CompositeEmitter(Emitter):
+    def __init__(self, emitters: List[Emitter]) -> None:
+        self.emitters = emitters
+
+    def emit(
+        self,
+        item: Union[
+            MetadataChangeEvent,
+            MetadataChangeProposal,
+            MetadataChangeProposalWrapper,
+        ],
+        callback: Optional[Callable[[Exception, str], None]] = None,
+    ) -> None:
+        callback_called = False
+        for emitter in self.emitters:
+            if not callback_called:
+                # We want to ensure that the callback is only called once and we tie it to the first emitter
+                emitter.emit(item, callback)
+                callback_called = True
+            else:
+                emitter.emit(item)
+
+    def flush(self) -> None:
+        for emitter in self.emitters:
+            emitter.flush()

datahub/emitter/rest_emitter.py

@@ -22,7 +22,7 @@ from deprecated import deprecated
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import HTTPError, RequestException
 
-from datahub import nice_version_name
+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

datahub/entrypoints.py

@@ -1,4 +1,5 @@
 import logging
+import multiprocessing
 import os
 import platform
 import sys
@@ -6,7 +7,7 @@ from typing import ContextManager, Optional
 
 import click
 
-import datahub as datahub_package
+import datahub._version as datahub_version
 from datahub.cli.check_cli import check
 from datahub.cli.cli_utils import (
     fixup_gms_url,
@@ -74,8 +75,8 @@ if sys.version_info >= (3, 12):
     help="Write debug-level logs to a file.",
 )
 @click.version_option(
-    version=datahub_package.nice_version_name(),
-    prog_name=datahub_package.__package_name__,
+    version=datahub_version.nice_version_name(),
+    prog_name=datahub_version.__package_name__,
 )
 def datahub(
     debug: bool,
@@ -112,7 +113,7 @@ def datahub(
 def version(include_server: bool = False) -> None:
     """Print version number and exit."""
 
-    click.echo(f"DataHub CLI version: {datahub_package.nice_version_name()}")
+    click.echo(f"DataHub CLI version: {datahub_version.nice_version_name()}")
     click.echo(f"Models: {model_version_name()}")
     click.echo(f"Python version: {sys.version}")
     if include_server:
@@ -183,6 +184,18 @@ datahub.add_command(datacontract)
 datahub.add_command(assertions)
 datahub.add_command(container)
 
+try:
+    from datahub.cli.iceberg_cli import iceberg
+
+    datahub.add_command(iceberg)
+except ImportError as e:
+    logger.debug(f"Failed to load datahub iceberg command: {e}")
+    datahub.add_command(
+        make_shim_command(
+            "iceberg", "run `pip install 'acryl-datahub[iceberg-catalog]'`"
+        )
+    )
+
 try:
     from datahub.cli.lite_cli import lite
 
@@ -205,6 +218,14 @@ except ImportError as e:
 
 
 def main(**kwargs):
+    # We use threads in a variety of places within our CLI. The multiprocessing
+    # "fork" start method is not safe to use with threads.
+    # MacOS and Windows already default to "spawn", and Linux will as well starting in Python 3.14.
+    # https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
+    # Eventually it may make sense to use "forkserver" as the default where available,
+    # but we can revisit that in the future.
+    multiprocessing.set_start_method("spawn", force=True)
+
     # This wrapper prevents click from suppressing errors.
     try:
         sys.exit(datahub(standalone_mode=False, **kwargs))
@@ -223,7 +244,7 @@ def main(**kwargs):
             logger.exception(f"Command failed: {exc}")
 
         logger.debug(
-            f"DataHub CLI version: {datahub_package.__version__} at {datahub_package.__file__}"
+            f"DataHub CLI version: {datahub_version.__version__} at {__file__}"
         )
         logger.debug(
             f"Python version: {sys.version} at {sys.executable} on {platform.platform()}"

datahub/ingestion/api/incremental_lineage_helper.py

@@ -102,6 +102,10 @@ def convert_dashboard_info_to_patch(
     if aspect.datasets:
         patch_builder.add_datasets(aspect.datasets)
 
+    if aspect.dashboards:
+        for dashboard in aspect.dashboards:
+            patch_builder.add_dashboard(dashboard)
+
     if aspect.access:
         patch_builder.set_access(aspect.access)
 

datahub/ingestion/api/registry.py

@@ -17,7 +17,7 @@ from typing import (
 
 import typing_inspect
 
-from datahub import __package_name__
+from datahub._version import __package_name__
 from datahub.configuration.common import ConfigurationError
 
 if sys.version_info < (3, 10):
@@ -173,8 +173,10 @@ class PluginRegistry(Generic[T]):
 
         tp = self._ensure_not_lazy(key)
         if isinstance(tp, ModuleNotFoundError):
+            # TODO: Once we're on Python 3.11 (with PEP 678), we can use .add_note()
+            # to enrich the error instead of wrapping it.
             raise ConfigurationError(
-                f"{key} is disabled; try running: pip install '{__package_name__}[{key}]'"
+                f"{key} is disabled due to a missing dependency: {tp.name}; try running `pip install '{__package_name__}[{key}]'`"
             ) from tp
         elif isinstance(tp, Exception):
             raise ConfigurationError(

datahub/ingestion/glossary/classification_mixin.py

@@ -1,5 +1,6 @@
 import concurrent.futures
 import logging
+import multiprocessing
 from dataclasses import dataclass, field
 from functools import partial
 from math import ceil
@@ -182,6 +183,11 @@ class ClassificationHandler:
 
         with concurrent.futures.ProcessPoolExecutor(
             max_workers=self.config.classification.max_workers,
+            # The fork start method, which is the default on Linux for Python < 3.14, is not
+            # safe when the main process uses threads. The default start method on windows/macOS is
+            # already spawn, and will be changed to spawn for Linux in Python 3.14.
+            # https://docs.python.org/3/library/multiprocessing.html#contexts-and-start-methods
+            mp_context=multiprocessing.get_context("spawn"),
         ) as executor:
             column_info_proposal_futures = [
                 executor.submit(

datahub/ingestion/glossary/classifier.py

@@ -1,3 +1,4 @@
+import os
 from abc import ABCMeta, abstractmethod
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional
@@ -37,8 +38,8 @@ class ClassificationConfig(ConfigModel):
     )
 
     max_workers: int = Field(
-        default=1,
-        description="Number of worker processes to use for classification. Note that any number above 1 might lead to a deadlock. Set to 1 to disable.",
+        default=(os.cpu_count() or 4),
+        description="Number of worker processes to use for classification. Set to 1 to disable.",
     )
 
     table_pattern: AllowDenyPattern = Field(

datahub/ingestion/graph/client.py

@@ -40,6 +40,7 @@ from datahub.ingestion.graph.connections import (
     connections_gql,
     get_id_from_connection_urn,
 )
+from datahub.ingestion.graph.entity_versioning import EntityVersioningAPI
 from datahub.ingestion.graph.filters import (
     RemovedStatusFilter,
     SearchFilterRule,
@@ -125,7 +126,7 @@ def _graphql_entity_type(entity_type: str) -> str:
     return entity_type
 
 
-class DataHubGraph(DatahubRestEmitter):
+class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def __init__(self, config: DatahubClientConfig) -> None:
         self.config = config
         super().__init__(

datahub/ingestion/graph/entity_versioning.py

@@ -0,0 +1,201 @@
+import uuid
+from typing import Dict, Optional, Protocol, Type
+
+from datahub.emitter.mce_builder import Aspect
+from datahub.metadata.schema_classes import (
+    VersionPropertiesClass,
+    VersionSetPropertiesClass,
+)
+from datahub.metadata.urns import VersionSetUrn
+from datahub.utilities.urns.urn import guess_entity_type
+
+
+class DataHubGraphProtocol(Protocol):
+    def execute_graphql(
+        self,
+        query: str,
+        variables: Optional[Dict],
+        operation_name: Optional[str] = None,
+        format_exception: bool = True,
+    ) -> Dict: ...
+
+    def get_aspect(
+        self,
+        entity_urn: str,
+        aspect_type: Type[Aspect],
+        version: int = 0,
+    ) -> Optional[Aspect]: ...
+
+
+class EntityVersioningAPI(DataHubGraphProtocol):
+    LINK_VERSION_MUTATION = """
+        mutation($input: LinkVersionInput!) {
+            linkAssetVersion(input: $input) {
+                urn
+            }
+        }
+    """
+
+    UNLINK_VERSION_MUTATION = """
+        mutation($input: UnlinkVersionInput!) {
+            unlinkAssetVersion(input: $input) {
+                urn
+            }
+        }
+    """
+
+    def link_asset_to_version_set(
+        self,
+        asset_urn: str,
+        version_set_urn: Optional[str],
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> Optional[str]:
+        """Sets an entity as the latest version of a version set.
+        Can also be used to create a new version set, with `asset_urn` as the first version.
+
+        Args:
+            asset_urn: URN of the entity.
+            version_set_urn: URN of the version set, or None to generate a new version set urn
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `asset_urn` was linked,
+            or None if the `asset_urn` was already linked to `version_set_urn`.
+        """
+
+        entity_type = guess_entity_type(asset_urn)
+        if version_set_urn is None:
+            version_set_urn = VersionSetUrn(str(uuid.uuid4()), entity_type).urn()
+        elif guess_entity_type(version_set_urn) != "versionSet":
+            raise ValueError(f"Expected version set URN, got {version_set_urn}")
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if entity_version and entity_version.versionSet:
+            if entity_version.versionSet == version_set_urn:
+                return None
+            else:
+                raise ValueError(
+                    f"Asset {asset_urn} is already a version of {entity_version.versionSet}"
+                )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "linkedEntity": asset_urn,
+                "version": label,
+                "comment": comment,
+            }
+        }
+        response = self.execute_graphql(self.LINK_VERSION_MUTATION, variables)
+        try:
+            return response["linkAssetVersion"]["urn"]
+        except KeyError:
+            raise ValueError(f"Unexpected response: {response}")
+
+    def link_asset_to_versioned_asset(
+        self,
+        new_asset_urn: str,
+        old_asset_urn: str,
+        label: str,
+        *,
+        comment: Optional[str] = None,
+    ) -> Optional[str]:
+        """Sets an entity as the latest version of an existing versioned entity.
+
+        Args:
+            new_asset_urn: URN of the new latest entity.
+            old_asset_urn: URN of an existing versioned entity to link onto.
+            label: Label of the version.
+            comment: Comment about the version.
+
+        Returns:
+            URN of the version set to which `new_asset_urn` was linked,
+            or None if the `new_asset_urn` was already linked to `old_asset_urn`.
+        """
+
+        new_entity_type = guess_entity_type(new_asset_urn)
+        old_entity_type = guess_entity_type(old_asset_urn)
+        if new_entity_type != old_entity_type:
+            raise ValueError(
+                f"Expected URNs of the same type, got {new_entity_type} and {old_entity_type}"
+            )
+
+        old_entity_version = self.get_aspect(old_asset_urn, VersionPropertiesClass)
+        if not old_entity_version:
+            raise ValueError(f"Asset {old_asset_urn} is not versioned")
+
+        new_entity_version = self.get_aspect(new_asset_urn, VersionPropertiesClass)
+        if new_entity_version:
+            if new_entity_version.versionSet == old_entity_version.versionSet:
+                return None
+            else:
+                raise ValueError(
+                    f"Asset {new_asset_urn} is already a version of {new_entity_version.versionSet}"
+                )
+
+        return self.link_asset_to_version_set(
+            new_asset_urn, old_entity_version.versionSet, label, comment=comment
+        )
+
+    def unlink_asset_from_version_set(self, asset_urn: str) -> Optional[str]:
+        """Unlinks an entity from its version set.
+
+        Args:
+            asset_urn: URN of the entity to unlink from its version set.
+
+        Returns:
+            If successful, the URN of the version set from which `asset_urn` was unlinked,
+            or None if `asset_urn` was not linked to any version set.
+        """
+
+        entity_version = self.get_aspect(asset_urn, VersionPropertiesClass)
+        if not entity_version:
+            return None
+
+        variables = {
+            "input": {
+                "versionSet": entity_version.versionSet,
+                "unlinkedEntity": asset_urn,
+            }
+        }
+        response = self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables)
+        try:
+            return response["unlinkAssetVersion"]["urn"]
+        except KeyError:
+            raise ValueError(f"Unexpected response: {response}")
+
+    def unlink_latest_asset_from_version_set(
+        self, version_set_urn: str
+    ) -> Optional[str]:
+        """Unlinks the latest version of a version set.
+
+        Args:
+            version_set_urn: URN of the version set.
+
+        Returns:
+            If successful, the URN of the entity that was unlinked from `version_set_urn`,
+            or None if no entity was unlinked.
+        """
+
+        version_set_properties = self.get_aspect(
+            version_set_urn, VersionSetPropertiesClass
+        )
+        if not version_set_properties:
+            raise ValueError(
+                f"Version set {version_set_urn} does not exist or has no versions"
+            )
+
+        variables = {
+            "input": {
+                "versionSet": version_set_urn,
+                "unlinkedEntity": version_set_properties.latest,
+            }
+        }
+        response = self.execute_graphql(self.UNLINK_VERSION_MUTATION, variables)
+        try:
+            return response["unlinkAssetVersion"]["urn"]
+        except KeyError:
+            raise ValueError(f"Unexpected response: {response}")

datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -3,7 +3,7 @@ import logging
 import time
 from typing import Any, Dict, Optional
 
-from datahub import nice_version_name
+from datahub._version import nice_version_name
 from datahub.configuration.common import (
     ConfigModel,
     DynamicTypedConfig,

datahub/ingestion/run/connection.py

@@ -1,6 +1,6 @@
 import logging
 
-from datahub import __version__
+from datahub._version import __version__
 from datahub.ingestion.api.source import TestableSource, TestConnectionReport
 from datahub.ingestion.source.source_registry import source_registry
 

datahub/ingestion/run/pipeline.py

@@ -15,7 +15,7 @@ import click
 import humanfriendly
 import psutil
 
-import datahub
+from datahub._version import nice_version_name
 from datahub.configuration.common import (
     ConfigModel,
     IgnorableError,
@@ -144,8 +144,8 @@ def _add_init_error_context(step: str) -> Iterator[None]:
 
 @dataclass
 class CliReport(Report):
-    cli_version: str = datahub.nice_version_name()
-    cli_entry_location: str = datahub.__file__
+    cli_version: str = nice_version_name()
+    cli_entry_location: str = __file__
     models_version: str = model_version_name()
     py_version: str = sys.version
     py_exec_path: str = sys.executable

datahub/ingestion/source/abs/report.py

@@ -1,16 +1,16 @@
 import dataclasses
 from dataclasses import field as dataclass_field
-from typing import List
 
 from datahub.ingestion.source.state.stale_entity_removal_handler import (
     StaleEntityRemovalSourceReport,
 )
+from datahub.utilities.lossy_collections import LossyList
 
 
 @dataclasses.dataclass
 class DataLakeSourceReport(StaleEntityRemovalSourceReport):
     files_scanned = 0
-    filtered: List[str] = dataclass_field(default_factory=list)
+    filtered: LossyList[str] = dataclass_field(default_factory=LossyList)
 
     def report_file_scanned(self) -> None:
         self.files_scanned += 1