acryl-datahub 1.1.1rc4__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -1,7 +1,10 @@
 from __future__ import annotations
 
 import abc
+import json
 from typing import (
+    TYPE_CHECKING,
+    Annotated,
     Any,
     ClassVar,
     Iterator,
@@ -15,7 +18,10 @@ from typing import (
 import pydantic
 
 from datahub.configuration.common import ConfigModel
-from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
+from datahub.configuration.pydantic_migration_helpers import (
+    PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR,
+    PYDANTIC_VERSION_2,
+)
 from datahub.ingestion.graph.client import flexible_entity_type_to_graphql
 from datahub.ingestion.graph.filters import (
     FilterOperator,
@@ -24,7 +30,14 @@ from datahub.ingestion.graph.filters import (
     _get_status_filter,
 )
 from datahub.metadata.schema_classes import EntityTypeName
-from datahub.metadata.urns import DataPlatformUrn, DomainUrn
+from datahub.metadata.urns import (
+    ContainerUrn,
+    CorpGroupUrn,
+    CorpUserUrn,
+    DataPlatformUrn,
+    DomainUrn,
+)
+from datahub.utilities.urns.urn import guess_entity_type
 
 _AndSearchFilterRule = TypedDict(
     "_AndSearchFilterRule", {"and": List[SearchFilterRule]}
@@ -33,21 +46,32 @@ _OrFilters = List[_AndSearchFilterRule]
 
 
 class _BaseFilter(ConfigModel):
-    class Config:
-        # We can't wrap this in a TYPE_CHECKING block because the pydantic plugin
-        # doesn't recognize it properly. So unfortunately we'll need to live
-        # with the deprecation warning w/ pydantic v2.
-        allow_population_by_field_name = True
-        if PYDANTIC_VERSION_2:
-            populate_by_name = True
+    model_config = pydantic.ConfigDict(populate_by_name=True)
 
     @abc.abstractmethod
-    def compile(self) -> _OrFilters:
-        pass
+    def compile(self) -> _OrFilters: ...
 
     def dfs(self) -> Iterator[_BaseFilter]:
         yield self
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        if cls is _BaseFilter:
+            raise ValueError("Cannot get discriminator for _BaseFilter")
+        if PYDANTIC_VERSION_2:
+            fields: dict = cls.model_fields  # type: ignore
+        else:
+            fields = cls.__fields__  # type: ignore
+
+        # Assumes that there's only one field name per filter.
+        # If that's not the case, this method should be overridden.
+        if len(fields.keys()) != 1:
+            raise ValueError(
+                f"Found multiple fields that could be the discriminator for this filter: {list(fields.keys())}"
+            )
+        name, field = next(iter(fields.items()))
+        return field.alias or name  # type: ignore
+
 
 class _EntityTypeFilter(_BaseFilter):
     """Filter for specific entity types.
@@ -59,7 +83,7 @@ class _EntityTypeFilter(_BaseFilter):
     ENTITY_TYPE_FIELD: ClassVar[str] = "_entityType"
 
     entity_type: List[str] = pydantic.Field(
-        description="The entity type to filter on. Can be 'dataset', 'chart', 'dashboard', 'corpuser', etc.",
+        description="The entity type to filter on. Can be 'dataset', 'chart', 'dashboard', 'corpuser', 'dataProduct', etc.",
     )
 
     def _build_rule(self) -> SearchFilterRule:
@@ -74,15 +98,19 @@ class _EntityTypeFilter(_BaseFilter):
 
 
 class _EntitySubtypeFilter(_BaseFilter):
-    entity_subtype: str = pydantic.Field(
+    entity_subtype: List[str] = pydantic.Field(
         description="The entity subtype to filter on. Can be 'Table', 'View', 'Source', etc. depending on the native platform's concepts.",
     )
 
+    @pydantic.validator("entity_subtype", pre=True)
+    def validate_entity_subtype(cls, v: str) -> List[str]:
+        return [v] if not isinstance(v, list) else v
+
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
             field="typeNames",
             condition="EQUAL",
-            values=[self.entity_subtype],
+            values=self.entity_subtype,
         )
 
     def compile(self) -> _OrFilters:
@@ -148,6 +176,39 @@ class _DomainFilter(_BaseFilter):
         return [{"and": [self._build_rule()]}]
 
 
+class _ContainerFilter(_BaseFilter):
+    container: List[str]
+    direct_descendants_only: bool = pydantic.Field(
+        default=False,
+        description="If true, only entities that are direct descendants of the container will be returned.",
+    )
+
+    @pydantic.validator("container", each_item=True)
+    def validate_container(cls, v: str) -> str:
+        return str(ContainerUrn.from_string(v))
+
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        return "container"
+
+    def _build_rule(self) -> SearchFilterRule:
+        if self.direct_descendants_only:
+            return SearchFilterRule(
+                field="container",
+                condition="EQUAL",
+                values=self.container,
+            )
+        else:
+            return SearchFilterRule(
+                field="browsePathV2",
+                condition="CONTAIN",
+                values=self.container,
+            )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
 class _EnvFilter(_BaseFilter):
     # Note that not all entity types have an env (e.g. dashboards / charts).
     # If the env filter is specified, these will be excluded.
@@ -181,6 +242,94 @@ class _EnvFilter(_BaseFilter):
         ]
 
 
+class _OwnerFilter(_BaseFilter):
+    """Filter for entities owned by specific users or groups."""
+
+    owner: List[str] = pydantic.Field(
+        description="The owner to filter on. Should be user or group URNs.",
+    )
+
+    @pydantic.validator("owner", each_item=True)
+    def validate_owner(cls, v: str) -> str:
+        if not v.startswith("urn:li:"):
+            raise ValueError(f"Owner must be a valid User or Group URN, got: {v}")
+        _type = guess_entity_type(v)
+        if _type == CorpUserUrn.ENTITY_TYPE:
+            return str(CorpUserUrn.from_string(v))
+        elif _type == CorpGroupUrn.ENTITY_TYPE:
+            return str(CorpGroupUrn.from_string(v))
+        else:
+            raise ValueError(f"Owner must be a valid User or Group URN, got: {v}")
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="owners",
+            condition="EQUAL",
+            values=self.owner,
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
+class _GlossaryTermFilter(_BaseFilter):
+    """Filter for entities associated with specific glossary terms."""
+
+    glossary_term: List[str] = pydantic.Field(
+        description="The glossary term to filter on. Should be glossary term URNs.",
+    )
+
+    @pydantic.validator("glossary_term", each_item=True)
+    def validate_glossary_term(cls, v: str) -> str:
+        if not v.startswith("urn:li:"):
+            raise ValueError(f"Glossary term must be a valid URN, got: {v}")
+        # Validate that it's a glossary term URN
+        _type = guess_entity_type(v)
+        if _type != "glossaryTerm":
+            raise ValueError(
+                f"Glossary term must be a valid glossary term URN, got: {v}"
+            )
+        return v
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="glossaryTerms",
+            condition="EQUAL",
+            values=self.glossary_term,
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
+class _TagFilter(_BaseFilter):
+    """Filter for entities associated with specific tags."""
+
+    tag: List[str] = pydantic.Field(
+        description="The tag to filter on. Should be tag URNs.",
+    )
+
+    @pydantic.validator("tag", each_item=True)
+    def validate_tag(cls, v: str) -> str:
+        if not v.startswith("urn:li:"):
+            raise ValueError(f"Tag must be a valid URN, got: {v}")
+        # Validate that it's a tag URN
+        _type = guess_entity_type(v)
+        if _type != "tag":
+            raise ValueError(f"Tag must be a valid tag URN, got: {v}")
+        return v
+
+    def _build_rule(self) -> SearchFilterRule:
+        return SearchFilterRule(
+            field="tags",
+            condition="EQUAL",
+            values=self.tag,
+        )
+
+    def compile(self) -> _OrFilters:
+        return [{"and": [self._build_rule()]}]
+
+
 class _CustomCondition(_BaseFilter):
     """Represents a single field condition."""
 
@@ -196,6 +345,10 @@ class _CustomCondition(_BaseFilter):
         )
         return [{"and": [rule]}]
 
+    @classmethod
+    def _field_discriminator(cls) -> str:
+        return "_custom"
+
 
 class _And(_BaseFilter):
     """Represents an AND conjunction of filters."""
@@ -302,31 +455,116 @@ class _Not(_BaseFilter):
         yield from self.not_.dfs()
 
 
-# TODO: With pydantic 2, we can use a RootModel with a
-# discriminated union to make the error messages more informative.
-Filter = Union[
-    _And,
-    _Or,
-    _Not,
-    _EntityTypeFilter,
-    _EntitySubtypeFilter,
-    _StatusFilter,
-    _PlatformFilter,
-    _DomainFilter,
-    _EnvFilter,
-    _CustomCondition,
-]
-
-
-# Required to resolve forward references to "Filter"
-if PYDANTIC_VERSION_2:
-    _And.model_rebuild()  # type: ignore
-    _Or.model_rebuild()  # type: ignore
-    _Not.model_rebuild()  # type: ignore
-else:
+def _filter_discriminator(v: Any) -> Optional[str]:
+    if isinstance(v, _BaseFilter):
+        return v._field_discriminator()
+
+    if not isinstance(v, dict):
+        return None
+
+    keys = list(v.keys())
+    if len(keys) == 1:
+        return keys[0]
+    elif set(keys).issuperset({"container"}):
+        return _ContainerFilter._field_discriminator()
+    elif set(keys).issuperset({"field", "condition"}):
+        return _CustomCondition._field_discriminator()
+
+    return None
+
+
+def _parse_and_like_filter(value: Any) -> Any:
+    # Do not parse if filter is already of type and/or/not or a custom condition
+    # also do not parse container filter if direct_descendants_only is specified
+    if (
+        isinstance(value, dict)
+        and not set(value.keys()).intersection(
+            {"and", "or", "not", "field", "condition", "direct_descendants_only"}
+        )
+        and len(value) > 1
+    ):
+        return {"and": [{k: v} for k, v in value.items()]}
+
+    return value
+
+
+if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
+    # The `not TYPE_CHECKING` bit is required to make the linter happy,
+    # since we currently only run mypy with pydantic v1.
+    Filter = Union[
+        _And,
+        _Or,
+        _Not,
+        _EntityTypeFilter,
+        _EntitySubtypeFilter,
+        _StatusFilter,
+        _PlatformFilter,
+        _DomainFilter,
+        _ContainerFilter,
+        _EnvFilter,
+        _OwnerFilter,
+        _GlossaryTermFilter,
+        _TagFilter,
+        _CustomCondition,
+    ]
+
     _And.update_forward_refs()
     _Or.update_forward_refs()
     _Not.update_forward_refs()
+else:
+    from pydantic import Discriminator, Tag
+
+    def _parse_json_from_string(value: Any) -> Any:
+        if isinstance(value, str):
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError:
+                return value
+        else:
+            return value
+
+    # TODO: Once we're fully on pydantic 2, we can use a RootModel here.
+    # That way we'd be able to attach methods to the Filter type.
+    # e.g. replace load_filters(...) with Filter.load(...)
+    Filter = Annotated[
+        Annotated[
+            Union[
+                Annotated[_And, Tag(_And._field_discriminator())],
+                Annotated[_Or, Tag(_Or._field_discriminator())],
+                Annotated[_Not, Tag(_Not._field_discriminator())],
+                Annotated[
+                    _EntityTypeFilter, Tag(_EntityTypeFilter._field_discriminator())
+                ],
+                Annotated[
+                    _EntitySubtypeFilter,
+                    Tag(_EntitySubtypeFilter._field_discriminator()),
+                ],
+                Annotated[_StatusFilter, Tag(_StatusFilter._field_discriminator())],
+                Annotated[_PlatformFilter, Tag(_PlatformFilter._field_discriminator())],
+                Annotated[_DomainFilter, Tag(_DomainFilter._field_discriminator())],
+                Annotated[
+                    _ContainerFilter, Tag(_ContainerFilter._field_discriminator())
+                ],
+                Annotated[_EnvFilter, Tag(_EnvFilter._field_discriminator())],
+                Annotated[_OwnerFilter, Tag(_OwnerFilter._field_discriminator())],
+                Annotated[
+                    _GlossaryTermFilter, Tag(_GlossaryTermFilter._field_discriminator())
+                ],
+                Annotated[_TagFilter, Tag(_TagFilter._field_discriminator())],
+                Annotated[
+                    _CustomCondition, Tag(_CustomCondition._field_discriminator())
+                ],
+            ],
+            Discriminator(_filter_discriminator),
+        ],
+        pydantic.BeforeValidator(_parse_and_like_filter),
+        pydantic.BeforeValidator(_parse_json_from_string),
+    ]
+
+    # Required to resolve forward references to "Filter"
+    _And.model_rebuild()  # type: ignore
+    _Or.model_rebuild()  # type: ignore
+    _Not.model_rebuild()  # type: ignore
 
 
 def load_filters(obj: Any) -> Filter:
@@ -400,10 +638,40 @@ class FilterDsl:
     def domain(domain: Union[str, Sequence[str]], /) -> _DomainFilter:
         return _DomainFilter(domain=[domain] if isinstance(domain, str) else domain)
 
+    @staticmethod
+    def container(
+        container: Union[str, Sequence[str]],
+        /,
+        *,
+        direct_descendants_only: bool = False,
+    ) -> _ContainerFilter:
+        return _ContainerFilter(
+            container=[container] if isinstance(container, str) else container,
+            direct_descendants_only=direct_descendants_only,
+        )
+
     @staticmethod
     def env(env: Union[str, Sequence[str]], /) -> _EnvFilter:
         return _EnvFilter(env=[env] if isinstance(env, str) else env)
 
+    @staticmethod
+    def owner(owner: Union[str, Sequence[str]], /) -> _OwnerFilter:
+        return _OwnerFilter(owner=[owner] if isinstance(owner, str) else owner)
+
+    @staticmethod
+    def glossary_term(
+        glossary_term: Union[str, Sequence[str]], /
+    ) -> _GlossaryTermFilter:
+        return _GlossaryTermFilter(
+            glossary_term=[glossary_term]
+            if isinstance(glossary_term, str)
+            else glossary_term
+        )
+
+    @staticmethod
+    def tag(tag: Union[str, Sequence[str]], /) -> _TagFilter:
+        return _TagFilter(tag=[tag] if isinstance(tag, str) else tag)
+
     @staticmethod
     def has_custom_property(key: str, value: str) -> _CustomCondition:
         return _CustomCondition(

datahub/secret/datahub_secret_store.py

@@ -65,3 +65,6 @@ class DataHubSecretStore(SecretStore):
     def create(cls, config: Any) -> "DataHubSecretStore":
         config = DataHubSecretStoreConfig.parse_obj(config)
         return cls(config)
+
+    def close(self) -> None:
+        self.client.graph.close()

datahub/secret/environment_secret_store.py

@@ -0,0 +1,29 @@
+import os
+from typing import Dict, List, Union
+
+from datahub.secret.secret_store import SecretStore
+
+
+# Simple SecretStore implementation that fetches Secret values from the local environment.
+class EnvironmentSecretStore(SecretStore):
+    def __init__(self, config):
+        pass
+
+    def close(self) -> None:
+        return
+
+    def get_secret_values(self, secret_names: List[str]) -> Dict[str, Union[str, None]]:
+        values = {}
+        for secret_name in secret_names:
+            values[secret_name] = os.getenv(secret_name)
+        return values
+
+    def get_secret_value(self, secret_name: str) -> Union[str, None]:
+        return os.getenv(secret_name)
+
+    def get_id(self) -> str:
+        return "env"
+
+    @classmethod
+    def create(cls, config: Dict) -> "EnvironmentSecretStore":
+        return cls(config)

datahub/secret/file_secret_store.py

@@ -0,0 +1,49 @@
+import logging
+import os
+from typing import Any, Dict, List, Union
+
+from pydantic import BaseModel
+
+from datahub.secret.secret_store import SecretStore
+
+logger = logging.getLogger(__name__)
+
+
+class FileSecretStoreConfig(BaseModel):
+    basedir: str = "/mnt/secrets"
+    max_length: int = 1024768
+
+
+# Simple SecretStore implementation that fetches Secret values from the local files.
+class FileSecretStore(SecretStore):
+    def __init__(self, config: FileSecretStoreConfig):
+        self.config = config
+
+    def get_secret_values(self, secret_names: List[str]) -> Dict[str, Union[str, None]]:
+        values = {}
+        for secret_name in secret_names:
+            values[secret_name] = self.get_secret_value(secret_name)
+        return values
+
+    def get_secret_value(self, secret_name: str) -> Union[str, None]:
+        secret_path = os.path.join(self.config.basedir, secret_name)
+        if os.path.exists(secret_path):
+            with open(secret_path, "r") as f:
+                secret_value = f.read(self.config.max_length + 1)
+                if len(secret_value) > self.config.max_length:
+                    logger.warning(
+                        f"Secret {secret_name} is longer than {self.config.max_length} and will be truncated."
+                    )
+                return secret_value[: self.config.max_length].rstrip()
+        return None
+
+    def get_id(self) -> str:
+        return "file"
+
+    def close(self) -> None:
+        return
+
+    @classmethod
+    def create(cls, config: Any) -> "FileSecretStore":
+        config = FileSecretStoreConfig.parse_obj(config)
+        return cls(config)

datahub/specific/aspect_helpers/fine_grained_lineage.py

@@ -0,0 +1,76 @@
+from abc import abstractmethod
+from typing import List, Tuple
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal, PatchPath
+from datahub.metadata.schema_classes import (
+    FineGrainedLineageClass as FineGrainedLineage,
+)
+
+
+class HasFineGrainedLineagePatch(MetadataPatchProposal):
+    @abstractmethod
+    def _fine_grained_lineage_location(self) -> Tuple[str, PatchPath]:
+        """Return the aspect name where fine-grained lineage is stored."""
+        raise NotImplementedError("Subclasses must implement this method.")
+
+    @staticmethod
+    def _get_fine_grained_key(
+        fine_grained_lineage: FineGrainedLineage,
+    ) -> Tuple[str, str, str]:
+        downstreams = fine_grained_lineage.downstreams or []
+        if len(downstreams) != 1:
+            raise TypeError("Cannot patch with more or less than one downstream.")
+        transform_op = fine_grained_lineage.transformOperation or "NONE"
+        downstream_urn = downstreams[0]
+        query_id = fine_grained_lineage.query or "NONE"
+        return transform_op, downstream_urn, query_id
+
+    def add_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "add",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={"confidenceScore": fine_grained_lineage.confidenceScore},
+            )
+        return self
+
+    def remove_fine_grained_lineage(
+        self, fine_grained_lineage: FineGrainedLineage
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        (
+            transform_op,
+            downstream_urn,
+            query_id,
+        ) = self._get_fine_grained_key(fine_grained_lineage)
+        for upstream_urn in fine_grained_lineage.upstreams or []:
+            self._add_patch(
+                aspect_name,
+                "remove",
+                path=(*path, transform_op, downstream_urn, query_id, upstream_urn),
+                value={},
+            )
+        return self
+
+    def set_fine_grained_lineages(
+        self, fine_grained_lineages: List[FineGrainedLineage]
+    ) -> Self:
+        aspect_name, path = self._fine_grained_lineage_location()
+        self._add_patch(
+            aspect_name,
+            "add",
+            path=path,
+            value=fine_grained_lineages,
+        )
+        return self

datahub/specific/aspect_helpers/siblings.py

@@ -0,0 +1,73 @@
+from typing import List
+
+from typing_extensions import Self
+
+from datahub.emitter.mcp_patch_builder import MetadataPatchProposal
+from datahub.metadata.com.linkedin.pegasus2avro.common import Siblings
+
+
+class HasSiblingsPatch(MetadataPatchProposal):
+    def add_sibling(self, sibling_urn: str, primary: bool = False) -> Self:
+        """Add a sibling relationship to the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to add.
+            primary: Whether this entity should be marked as primary in the relationship.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "add",
+            path=("siblings", sibling_urn),
+            value=sibling_urn,
+        )
+
+        # Set primary flag if specified
+        if primary:
+            self._add_patch(
+                Siblings.ASPECT_NAME,
+                "add",
+                path=("primary",),
+                value=primary,
+            )
+
+        return self
+
+    def remove_sibling(self, sibling_urn: str) -> Self:
+        """Remove a sibling relationship from the entity.
+
+        Args:
+            sibling_urn: The URN of the sibling entity to remove.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME,
+            "remove",
+            path=("siblings", sibling_urn),
+            value={},
+        )
+        return self
+
+    def set_siblings(self, sibling_urns: List[str], primary: bool = False) -> Self:
+        """Set the complete list of siblings for the entity.
+
+        This will replace all existing siblings with the new list.
+
+        Args:
+            sibling_urns: The list of sibling URNs to set.
+            primary: Whether this entity should be marked as primary.
+
+        Returns:
+            The patch builder instance.
+        """
+        self._add_patch(
+            Siblings.ASPECT_NAME, "add", path=("siblings",), value=sibling_urns
+        )
+
+        self._add_patch(Siblings.ASPECT_NAME, "add", path=("primary",), value=primary)
+
+        return self