acryl-datahub 1.2.0.11rc1__py3-none-any.whl → 1.2.0.11rc3__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -30,7 +30,14 @@ from datahub.ingestion.graph.filters import (
     _get_status_filter,
 )
 from datahub.metadata.schema_classes import EntityTypeName
-from datahub.metadata.urns import ContainerUrn, 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]}
@@ -235,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."""
 
@@ -407,6 +502,9 @@ if TYPE_CHECKING or not PYDANTIC_SUPPORTS_CALLABLE_DISCRIMINATOR:
         _DomainFilter,
         _ContainerFilter,
         _EnvFilter,
+        _OwnerFilter,
+        _GlossaryTermFilter,
+        _TagFilter,
         _CustomCondition,
     ]
 
@@ -448,6 +546,11 @@ else:
                     _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())
                 ],
@@ -551,6 +654,24 @@ class FilterDsl:
     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/structured_properties.py

@@ -70,3 +70,30 @@ class HasStructuredPropertiesPatch(MetadataPatchProposal):
             ),
         )
         return self
+
+    def set_structured_property_manual(
+        self, property: StructuredPropertyValueAssignmentClass
+    ) -> Self:
+        """Add or update a structured property, using a StructuredPropertyValueAssignmentClass object."""
+
+        self.remove_structured_property(property.propertyUrn)
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "add",
+            path=("properties", property.propertyUrn),
+            value=property,
+        )
+        return self
+
+    def add_structured_property_manual(
+        self, property: StructuredPropertyValueAssignmentClass
+    ) -> Self:
+        """Add a structured property, using a StructuredPropertyValueAssignmentClass object."""
+
+        self._add_patch(
+            StructuredPropertiesClass.ASPECT_NAME,
+            "add",
+            path=("properties", property.propertyUrn),
+            value=property,
+        )
+        return self

datahub/sql_parsing/sqlglot_lineage.py

@@ -1176,7 +1176,12 @@ def _try_extract_select(
             statement = sqlglot.exp.Select().select("*").from_(statement)
     elif isinstance(statement, sqlglot.exp.Insert):
         # TODO Need to map column renames in the expressions part of the statement.
-        statement = statement.expression
+        # Preserve CTEs when extracting the SELECT expression from INSERT
+        original_ctes = statement.ctes
+        statement = statement.expression  # Get the SELECT expression from the INSERT
+        if isinstance(statement, sqlglot.exp.Query) and original_ctes:
+            for cte in original_ctes:
+                statement = statement.with_(alias=cte.alias, as_=cte.this)
     elif isinstance(statement, sqlglot.exp.Update):
         # Assumption: the output table is already captured in the modified tables list.
         statement = _extract_select_from_update(statement)