acryl-datahub 0.15.0.5rc5__py3-none-any.whl → 0.15.0.5rc7__py3-none-any.whl

datahub/ingestion/source/apply/datahub_apply.py

@@ -0,0 +1,223 @@
+import logging
+from functools import partial
+from typing import Any, Iterable, List, Optional, Union
+
+import progressbar
+from pydantic import Field
+
+from datahub.configuration.common import ConfigModel
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.decorators import (
+    SupportStatus,
+    config_class,
+    platform_name,
+    support_status,
+)
+from datahub.ingestion.api.source import MetadataWorkUnitProcessor, Source, SourceReport
+from datahub.ingestion.api.source_helpers import auto_workunit_reporter
+from datahub.ingestion.api.workunit import MetadataWorkUnit
+from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.metadata.schema_classes import (
+    DomainsClass,
+    GlossaryTermAssociationClass,
+    MetadataChangeProposalClass,
+    OwnerClass,
+    OwnershipTypeClass,
+    TagAssociationClass,
+)
+from datahub.specific.dataset import DatasetPatchBuilder
+
+logger = logging.getLogger(__name__)
+
+
+def apply_association_to_container(
+    container_urn: str,
+    association_urn: str,
+    association_type: str,
+    emit: bool = True,
+    graph: Optional[DataHubGraph] = None,
+) -> Optional[List[Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]]]:
+    """
+    Common function to add either tags, terms, domains, or owners to child datasets (for now).
+
+    Args:
+        container_urn: The URN of the container
+        association_urn: The URN of the tag, term, or user to apply
+        association_type: One of 'tag', 'term', 'domain' or 'owner'
+    """
+    urns: List[str] = [container_urn]
+    if not graph:
+        graph = get_default_graph()
+    logger.info(f"Using {graph}")
+    urns.extend(
+        graph.get_urns_by_filter(
+            container=container_urn,
+            batch_size=1000,
+            entity_types=["dataset", "container"],
+        )
+    )
+
+    all_patches: List[Any] = []
+    for urn in urns:
+        builder = DatasetPatchBuilder(urn)
+        patches: List[Any] = []
+        if association_type == "tag":
+            patches = builder.add_tag(TagAssociationClass(association_urn)).build()
+        elif association_type == "term":
+            patches = builder.add_term(
+                GlossaryTermAssociationClass(association_urn)
+            ).build()
+        elif association_type == "owner":
+            patches = builder.add_owner(
+                OwnerClass(
+                    owner=association_urn,
+                    type=OwnershipTypeClass.TECHNICAL_OWNER,
+                )
+            ).build()
+        elif association_type == "domain":
+            patches = [
+                MetadataChangeProposalWrapper(
+                    entityUrn=urn,
+                    aspect=DomainsClass(domains=[association_urn]),
+                )
+            ]
+        all_patches.extend(patches)
+    if emit:
+        mcps_iter = progressbar.progressbar(all_patches, redirect_stdout=True)
+        for mcp in mcps_iter:
+            graph.emit(mcp)
+        return None
+    else:
+        return all_patches
+
+
+class DomainApplyConfig(ConfigModel):
+    assets: List[str] = Field(
+        default_factory=list,
+        description="List of assets to apply domain hierarchichaly. Currently only containers and datasets are supported",
+    )
+    domain_urn: str = Field(default="")
+
+
+class TagApplyConfig(ConfigModel):
+    assets: List[str] = Field(
+        default_factory=list,
+        description="List of assets to apply tag hierarchichaly. Currently only containers and datasets are supported",
+    )
+    tag_urn: str = Field(default="")
+
+
+class TermApplyConfig(ConfigModel):
+    assets: List[str] = Field(
+        default_factory=list,
+        description="List of assets to apply term hierarchichaly. Currently only containers and datasets are supported",
+    )
+    term_urn: str = Field(default="")
+
+
+class OwnerApplyConfig(ConfigModel):
+    assets: List[str] = Field(
+        default_factory=list,
+        description="List of assets to apply owner hierarchichaly. Currently only containers and datasets are supported",
+    )
+    owner_urn: str = Field(default="")
+
+
+class DataHubApplyConfig(ConfigModel):
+    domain_apply: Optional[List[DomainApplyConfig]] = Field(
+        default=None,
+        description="List to apply domains to assets",
+    )
+    tag_apply: Optional[List[TagApplyConfig]] = Field(
+        default=None,
+        description="List to apply tags to assets",
+    )
+    term_apply: Optional[List[TermApplyConfig]] = Field(
+        default=None,
+        description="List to apply terms to assets",
+    )
+    owner_apply: Optional[List[OwnerApplyConfig]] = Field(
+        default=None,
+        description="List to apply owners to assets",
+    )
+
+
+@platform_name("DataHubApply")
+@config_class(DataHubApplyConfig)
+@support_status(SupportStatus.TESTING)
+class DataHubApplySource(Source):
+    """
+    This source is a helper over CLI
+    so people can use the helper to apply various metadata changes to DataHub
+    via Managed Ingestion
+    """
+
+    def __init__(self, ctx: PipelineContext, config: DataHubApplyConfig):
+        self.ctx = ctx
+        self.config = config
+        self.report = SourceReport()
+        self.graph = ctx.require_graph()
+
+    def _yield_workunits(
+        self,
+        proposals: List[
+            Union[MetadataChangeProposalWrapper, MetadataChangeProposalClass]
+        ],
+    ) -> Iterable[MetadataWorkUnit]:
+        for proposal in proposals:
+            if isinstance(proposal, MetadataChangeProposalWrapper):
+                yield proposal.as_workunit()
+            else:
+                yield MetadataWorkUnit(
+                    id=MetadataWorkUnit.generate_workunit_id(proposal),
+                    mcp_raw=proposal,
+                )
+
+    def _handle_assets(
+        self, assets: List[str], apply_urn: str, apply_type: str
+    ) -> Iterable[MetadataWorkUnit]:
+        for asset in assets:
+            change_proposals = apply_association_to_container(
+                asset, apply_urn, apply_type, emit=False, graph=self.graph
+            )
+            assert change_proposals is not None
+            yield from self._yield_workunits(change_proposals)
+
+    def _yield_domain(self) -> Iterable[MetadataWorkUnit]:
+        if not self.config.domain_apply:
+            return
+        for apply in self.config.domain_apply:
+            yield from self._handle_assets(apply.assets, apply.domain_urn, "domain")
+
+    def _yield_tag(self) -> Iterable[MetadataWorkUnit]:
+        if not self.config.tag_apply:
+            return
+        for apply in self.config.tag_apply:
+            yield from self._handle_assets(apply.assets, apply.tag_urn, "tag")
+
+    def _yield_term(self) -> Iterable[MetadataWorkUnit]:
+        if not self.config.term_apply:
+            return
+        for apply in self.config.term_apply:
+            yield from self._handle_assets(apply.assets, apply.term_urn, "term")
+
+    def _yield_owner(self) -> Iterable[MetadataWorkUnit]:
+        if not self.config.owner_apply:
+            return
+        for apply in self.config.owner_apply:
+            yield from self._handle_assets(apply.assets, apply.owner_urn, "owner")
+
+    def get_workunits_internal(
+        self,
+    ) -> Iterable[MetadataWorkUnit]:
+        yield from self._yield_domain()
+        yield from self._yield_tag()
+        yield from self._yield_term()
+        yield from self._yield_owner()
+
+    def get_workunit_processors(self) -> List[Optional[MetadataWorkUnitProcessor]]:
+        return [partial(auto_workunit_reporter, self.get_report())]
+
+    def get_report(self) -> SourceReport:
+        return self.report

datahub/ingestion/source/dbt/dbt_core.py

@@ -488,7 +488,7 @@ class DBTCoreSource(DBTSourceBase, TestableSource):
     ) -> Dict:
         if re.match("^https?://", uri):
             return json.loads(requests.get(uri).text)
-        elif re.match("^s3://", uri):
+        elif is_s3_uri(uri):
             u = urlparse(uri)
             assert aws_connection
             response = aws_connection.get_s3_client().get_object(

datahub/ingestion/source/sql/sql_config.py

@@ -2,8 +2,6 @@ import logging
 from abc import abstractmethod
 from typing import Any, Dict, Optional
 
-import cachetools
-import cachetools.keys
 import pydantic
 from pydantic import Field
 from sqlalchemy.engine import URL
@@ -29,7 +27,6 @@ from datahub.ingestion.source.state.stateful_ingestion_base import (
     StatefulIngestionConfigBase,
 )
 from datahub.ingestion.source_config.operation_config import is_profiling_enabled
-from datahub.utilities.cachetools_keys import self_methodkey
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -118,13 +115,6 @@ class SQLCommonConfig(
     # Custom Stateful Ingestion settings
     stateful_ingestion: Optional[StatefulStaleMetadataRemovalConfig] = None
 
-    # TRICKY: The operation_config is time-dependent. Because we don't want to change
-    # whether or not we're running profiling mid-ingestion, we cache the result of this method.
-    # TODO: This decorator should be moved to the is_profiling_enabled(operation_config) method.
-    @cachetools.cached(
-        cache=cachetools.LRUCache(maxsize=1),
-        key=self_methodkey,
-    )
     def is_profiling_enabled(self) -> bool:
         return self.profiling.enabled and is_profiling_enabled(
             self.profiling.operation_config

datahub/ingestion/source_config/operation_config.py

@@ -2,10 +2,12 @@ import datetime
 import logging
 from typing import Any, Dict, Optional
 
+import cachetools
 import pydantic
 from pydantic.fields import Field
 
 from datahub.configuration.common import ConfigModel
+from datahub.utilities.cachetools_keys import self_methodkey
 
 logger = logging.getLogger(__name__)
 
@@ -62,6 +64,13 @@ class OperationConfig(ConfigModel):
         return profile_date_of_month
 
 
+# TRICKY: The operation_config is time-dependent. Because we don't want to change
+# whether or not we're running profiling mid-ingestion, we cache the result of this method.
+# An additional benefit is that we only print the log lines on the first call.
+@cachetools.cached(
+    cache=cachetools.LRUCache(maxsize=1),
+    key=self_methodkey,
+)
 def is_profiling_enabled(operation_config: OperationConfig) -> bool:
     if operation_config.lower_freq_profile_enabled is False:
         return True

datahub/metadata/_schema_classes.py

@@ -971,13 +971,15 @@ class AssertionSourceTypeClass(object):
     # No docs available.
     
     NATIVE = "NATIVE"
-    """The assertion was defined natively on DataHub by a user."""
+    """The assertion was defined natively on DataHub by a user.
+    DataHub Cloud only"""
     
     EXTERNAL = "EXTERNAL"
     """The assertion was defined and managed externally of DataHub."""
     
     INFERRED = "INFERRED"
-    """The assertion was inferred, e.g. from offline AI / ML models."""
+    """The assertion was inferred, e.g. from offline AI / ML models.
+    DataHub Cloud only"""
     
     
     
@@ -4120,6 +4122,7 @@ class DeprecationClass(_Aspect):
         note: str,
         actor: str,
         decommissionTime: Union[None, int]=None,
+        replacement: Union[None, str]=None,
     ):
         super().__init__()
         
@@ -4127,12 +4130,14 @@ class DeprecationClass(_Aspect):
         self.decommissionTime = decommissionTime
         self.note = note
         self.actor = actor
+        self.replacement = replacement
     
     def _restore_defaults(self) -> None:
         self.deprecated = bool()
         self.decommissionTime = self.RECORD_SCHEMA.fields_dict["decommissionTime"].default
         self.note = str()
         self.actor = str()
+        self.replacement = self.RECORD_SCHEMA.fields_dict["replacement"].default
     
     
     @property
@@ -4175,6 +4180,58 @@ class DeprecationClass(_Aspect):
         self._inner_dict['actor'] = value
     
     
+    @property
+    def replacement(self) -> Union[None, str]:
+        # No docs available.
+        return self._inner_dict.get('replacement')  # type: ignore
+    
+    @replacement.setter
+    def replacement(self, value: Union[None, str]) -> None:
+        self._inner_dict['replacement'] = value
+    
+    
+class DisplayPropertiesClass(_Aspect):
+    """Properties related to how the entity is displayed in the Datahub UI"""
+
+
+    ASPECT_NAME = 'displayProperties'
+    ASPECT_INFO = {}
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.common.DisplayProperties")
+
+    def __init__(self,
+        colorHex: Union[None, str]=None,
+        icon: Union[None, "IconPropertiesClass"]=None,
+    ):
+        super().__init__()
+        
+        self.colorHex = colorHex
+        self.icon = icon
+    
+    def _restore_defaults(self) -> None:
+        self.colorHex = self.RECORD_SCHEMA.fields_dict["colorHex"].default
+        self.icon = self.RECORD_SCHEMA.fields_dict["icon"].default
+    
+    
+    @property
+    def colorHex(self) -> Union[None, str]:
+        """The color associated with the entity in Hex. For example #FFFFFF."""
+        return self._inner_dict.get('colorHex')  # type: ignore
+    
+    @colorHex.setter
+    def colorHex(self, value: Union[None, str]) -> None:
+        self._inner_dict['colorHex'] = value
+    
+    
+    @property
+    def icon(self) -> Union[None, "IconPropertiesClass"]:
+        """The icon associated with the entity"""
+        return self._inner_dict.get('icon')  # type: ignore
+    
+    @icon.setter
+    def icon(self, value: Union[None, "IconPropertiesClass"]) -> None:
+        self._inner_dict['icon'] = value
+    
+    
 class DocumentationClass(_Aspect):
     """Aspect used for storing all applicable documentations on assets.
     This aspect supports multiple documentations from different sources.
@@ -4830,6 +4887,65 @@ class GlossaryTermsClass(_Aspect):
         self._inner_dict['auditStamp'] = value
     
     
+class IconLibraryClass(object):
+    """Enum of possible icon sources"""
+    
+    MATERIAL = "MATERIAL"
+    """Material UI"""
+    
+    
+    
+class IconPropertiesClass(DictWrapper):
+    """Properties describing an icon associated with an entity"""
+    
+    RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.common.IconProperties")
+    def __init__(self,
+        iconLibrary: Union[str, "IconLibraryClass"],
+        name: str,
+        style: str,
+    ):
+        super().__init__()
+        
+        self.iconLibrary = iconLibrary
+        self.name = name
+        self.style = style
+    
+    def _restore_defaults(self) -> None:
+        self.iconLibrary = IconLibraryClass.MATERIAL
+        self.name = str()
+        self.style = str()
+    
+    
+    @property
+    def iconLibrary(self) -> Union[str, "IconLibraryClass"]:
+        """The source of the icon: e.g. Antd, Material, etc"""
+        return self._inner_dict.get('iconLibrary')  # type: ignore
+    
+    @iconLibrary.setter
+    def iconLibrary(self, value: Union[str, "IconLibraryClass"]) -> None:
+        self._inner_dict['iconLibrary'] = value
+    
+    
+    @property
+    def name(self) -> str:
+        """The name of the icon"""
+        return self._inner_dict.get('name')  # type: ignore
+    
+    @name.setter
+    def name(self, value: str) -> None:
+        self._inner_dict['name'] = value
+    
+    
+    @property
+    def style(self) -> str:
+        """Any modifier for the icon, this will be library-specific, e.g. filled/outlined, etc"""
+        return self._inner_dict.get('style')  # type: ignore
+    
+    @style.setter
+    def style(self, value: str) -> None:
+        self._inner_dict['style'] = value
+    
+    
 class IncidentSummaryDetailsClass(DictWrapper):
     """Summary statistics about incidents on an entity."""
     
@@ -12965,13 +13081,16 @@ class CorpUserAppearanceSettingsClass(DictWrapper):
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.identity.CorpUserAppearanceSettings")
     def __init__(self,
         showSimplifiedHomepage: Union[None, bool]=None,
+        showThemeV2: Union[None, bool]=None,
     ):
         super().__init__()
         
         self.showSimplifiedHomepage = showSimplifiedHomepage
+        self.showThemeV2 = showThemeV2
     
     def _restore_defaults(self) -> None:
         self.showSimplifiedHomepage = self.RECORD_SCHEMA.fields_dict["showSimplifiedHomepage"].default
+        self.showThemeV2 = self.RECORD_SCHEMA.fields_dict["showThemeV2"].default
     
     
     @property
@@ -12985,6 +13104,16 @@ class CorpUserAppearanceSettingsClass(DictWrapper):
         self._inner_dict['showSimplifiedHomepage'] = value
     
     
+    @property
+    def showThemeV2(self) -> Union[None, bool]:
+        """Flag controlling whether the V2 UI for DataHub is shown."""
+        return self._inner_dict.get('showThemeV2')  # type: ignore
+    
+    @showThemeV2.setter
+    def showThemeV2(self, value: Union[None, bool]) -> None:
+        self._inner_dict['showThemeV2'] = value
+    
+    
 class CorpUserCredentialsClass(_Aspect):
     """Corp user credentials"""
 
@@ -15780,7 +15909,7 @@ class SchemaFieldKeyClass(_Aspect):
 
 
     ASPECT_NAME = 'schemaFieldKey'
-    ASPECT_INFO = {'keyForEntity': 'schemaField', 'entityCategory': 'core', 'entityAspects': ['schemafieldInfo', 'structuredProperties', 'forms', 'businessAttributes', 'status', 'schemaFieldAliases', 'documentation', 'testResults']}
+    ASPECT_INFO = {'keyForEntity': 'schemaField', 'entityCategory': 'core', 'entityAspects': ['schemafieldInfo', 'structuredProperties', 'forms', 'businessAttributes', 'status', 'schemaFieldAliases', 'documentation', 'testResults', 'deprecation']}
     RECORD_SCHEMA = get_schema_type("com.linkedin.pegasus2avro.metadata.key.SchemaFieldKey")
 
     def __init__(self,
@@ -20824,7 +20953,7 @@ class PostInfoClass(_Aspect):
     
     @property
     def target(self) -> Union[None, str]:
-        """Optional URN that the post is associated with."""
+        """Optional Entity URN that the post is associated with."""
         return self._inner_dict.get('target')  # type: ignore
     
     @target.setter
@@ -20838,6 +20967,9 @@ class PostTypeClass(object):
     HOME_PAGE_ANNOUNCEMENT = "HOME_PAGE_ANNOUNCEMENT"
     """The Post is an Home Page announcement."""
     
+    ENTITY_ANNOUNCEMENT = "ENTITY_ANNOUNCEMENT"
+    """The Post is an Entity level announcement."""
+    
     
     
 class QueryLanguageClass(object):
@@ -25030,6 +25162,7 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.common.DataTransform': DataTransformClass,
     'com.linkedin.pegasus2avro.common.DataTransformLogic': DataTransformLogicClass,
     'com.linkedin.pegasus2avro.common.Deprecation': DeprecationClass,
+    'com.linkedin.pegasus2avro.common.DisplayProperties': DisplayPropertiesClass,
     'com.linkedin.pegasus2avro.common.Documentation': DocumentationClass,
     'com.linkedin.pegasus2avro.common.DocumentationAssociation': DocumentationAssociationClass,
     'com.linkedin.pegasus2avro.common.Edge': EdgeClass,
@@ -25044,6 +25177,8 @@ __SCHEMA_TYPES = {
     'com.linkedin.pegasus2avro.common.GlobalTags': GlobalTagsClass,
     'com.linkedin.pegasus2avro.common.GlossaryTermAssociation': GlossaryTermAssociationClass,
     'com.linkedin.pegasus2avro.common.GlossaryTerms': GlossaryTermsClass,
+    'com.linkedin.pegasus2avro.common.IconLibrary': IconLibraryClass,
+    'com.linkedin.pegasus2avro.common.IconProperties': IconPropertiesClass,
     'com.linkedin.pegasus2avro.common.IncidentSummaryDetails': IncidentSummaryDetailsClass,
     'com.linkedin.pegasus2avro.common.IncidentsSummary': IncidentsSummaryClass,
     'com.linkedin.pegasus2avro.common.InputField': InputFieldClass,
@@ -25503,6 +25638,7 @@ __SCHEMA_TYPES = {
     'DataTransform': DataTransformClass,
     'DataTransformLogic': DataTransformLogicClass,
     'Deprecation': DeprecationClass,
+    'DisplayProperties': DisplayPropertiesClass,
     'Documentation': DocumentationClass,
     'DocumentationAssociation': DocumentationAssociationClass,
     'Edge': EdgeClass,
@@ -25517,6 +25653,8 @@ __SCHEMA_TYPES = {
     'GlobalTags': GlobalTagsClass,
     'GlossaryTermAssociation': GlossaryTermAssociationClass,
     'GlossaryTerms': GlossaryTermsClass,
+    'IconLibrary': IconLibraryClass,
+    'IconProperties': IconPropertiesClass,
     'IncidentSummaryDetails': IncidentSummaryDetailsClass,
     'IncidentsSummary': IncidentsSummaryClass,
     'InputField': InputFieldClass,
@@ -25926,6 +26064,7 @@ ASPECT_CLASSES: List[Type[_Aspect]] = [
     DocumentationClass,
     DataPlatformInstanceClass,
     InputFieldsClass,
+    DisplayPropertiesClass,
     OwnershipClass,
     OperationClass,
     FormsClass,
@@ -26144,6 +26283,7 @@ class AspectBag(TypedDict, total=False):
     documentation: DocumentationClass
     dataPlatformInstance: DataPlatformInstanceClass
     inputFields: InputFieldsClass
+    displayProperties: DisplayPropertiesClass
     ownership: OwnershipClass
     operation: OperationClass
     forms: FormsClass

datahub/metadata/com/linkedin/pegasus2avro/common/__init__.py

@@ -22,6 +22,7 @@ from .....schema_classes import DataPlatformInstanceClass
 from .....schema_classes import DataTransformClass
 from .....schema_classes import DataTransformLogicClass
 from .....schema_classes import DeprecationClass
+from .....schema_classes import DisplayPropertiesClass
 from .....schema_classes import DocumentationClass
 from .....schema_classes import DocumentationAssociationClass
 from .....schema_classes import EdgeClass
@@ -36,6 +37,8 @@ from .....schema_classes import FormsClass
 from .....schema_classes import GlobalTagsClass
 from .....schema_classes import GlossaryTermAssociationClass
 from .....schema_classes import GlossaryTermsClass
+from .....schema_classes import IconLibraryClass
+from .....schema_classes import IconPropertiesClass
 from .....schema_classes import IncidentSummaryDetailsClass
 from .....schema_classes import IncidentsSummaryClass
 from .....schema_classes import InputFieldClass
@@ -85,6 +88,7 @@ DataPlatformInstance = DataPlatformInstanceClass
 DataTransform = DataTransformClass
 DataTransformLogic = DataTransformLogicClass
 Deprecation = DeprecationClass
+DisplayProperties = DisplayPropertiesClass
 Documentation = DocumentationClass
 DocumentationAssociation = DocumentationAssociationClass
 Edge = EdgeClass
@@ -99,6 +103,8 @@ Forms = FormsClass
 GlobalTags = GlobalTagsClass
 GlossaryTermAssociation = GlossaryTermAssociationClass
 GlossaryTerms = GlossaryTermsClass
+IconLibrary = IconLibraryClass
+IconProperties = IconPropertiesClass
 IncidentSummaryDetails = IncidentSummaryDetailsClass
 IncidentsSummary = IncidentsSummaryClass
 InputField = InputFieldClass