acryl-datahub 1.0.0.1rc6__py3-none-any.whl → 1.0.0.2__py3-none-any.whl

datahub/sdk/search_filters.py

@@ -3,7 +3,10 @@ from __future__ import annotations
 import abc
 from typing import (
     Any,
+    ClassVar,
+    Iterator,
     List,
+    Optional,
     Sequence,
     TypedDict,
     Union,
@@ -13,8 +16,13 @@ import pydantic
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.pydantic_migration_helpers import PYDANTIC_VERSION_2
-from datahub.ingestion.graph.client import entity_type_to_graphql
-from datahub.ingestion.graph.filters import FilterOperator, SearchFilterRule
+from datahub.ingestion.graph.client import flexible_entity_type_to_graphql
+from datahub.ingestion.graph.filters import (
+    FilterOperator,
+    RemovedStatusFilter,
+    SearchFilterRule,
+    _get_status_filter,
+)
 from datahub.metadata.schema_classes import EntityTypeName
 from datahub.metadata.urns import DataPlatformUrn, DomainUrn
 
@@ -37,25 +45,28 @@ class _BaseFilter(ConfigModel):
     def compile(self) -> _OrFilters:
         pass
 
-
-def _flexible_entity_type_to_graphql(entity_type: str) -> str:
-    if entity_type.upper() == entity_type:
-        # Assume that we were passed a graphql EntityType enum value,
-        # so no conversion is needed.
-        return entity_type
-    return entity_type_to_graphql(entity_type)
+    def dfs(self) -> Iterator[_BaseFilter]:
+        yield self
 
 
 class _EntityTypeFilter(_BaseFilter):
+    """Filter for specific entity types.
+
+    If no entity type filter is specified, we will search all entity types in the
+    default search set, mirroring the behavior of the DataHub UI.
+    """
+
+    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.",
     )
 
     def _build_rule(self) -> SearchFilterRule:
         return SearchFilterRule(
-            field="_entityType",
+            field=self.ENTITY_TYPE_FIELD,
             condition="EQUAL",
-            values=[_flexible_entity_type_to_graphql(t) for t in self.entity_type],
+            values=[flexible_entity_type_to_graphql(t) for t in self.entity_type],
         )
 
     def compile(self) -> _OrFilters:
@@ -78,6 +89,26 @@ class _EntitySubtypeFilter(_BaseFilter):
         return [{"and": [self._build_rule()]}]
 
 
+class _StatusFilter(_BaseFilter):
+    """Filter for the status of entities during search.
+
+    If not explicitly specified, the NOT_SOFT_DELETED status filter will be applied.
+    """
+
+    status: RemovedStatusFilter
+
+    def _build_rule(self) -> Optional[SearchFilterRule]:
+        return _get_status_filter(self.status)
+
+    def compile(self) -> _OrFilters:
+        rule = self._build_rule()
+        if rule:
+            return [{"and": [rule]}]
+        else:
+            # Our boolean algebra logic requires something here - returning [] would cause errors.
+            return FilterDsl.true().compile()
+
+
 class _PlatformFilter(_BaseFilter):
     platform: List[str]
     # TODO: Add validator to convert string -> list of strings
@@ -213,6 +244,11 @@ class _And(_BaseFilter):
             ]
         }
 
+    def dfs(self) -> Iterator[_BaseFilter]:
+        yield self
+        for filter in self.and_:
+            yield from filter.dfs()
+
 
 class _Or(_BaseFilter):
     """Represents an OR conjunction of filters."""
@@ -226,6 +262,11 @@ class _Or(_BaseFilter):
             merged_filter.extend(filter.compile())
         return merged_filter
 
+    def dfs(self) -> Iterator[_BaseFilter]:
+        yield self
+        for filter in self.or_:
+            yield from filter.dfs()
+
 
 class _Not(_BaseFilter):
     """Represents a NOT filter."""
@@ -256,6 +297,10 @@ class _Not(_BaseFilter):
 
         return final_filters
 
+    def dfs(self) -> Iterator[_BaseFilter]:
+        yield self
+        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.
@@ -265,6 +310,7 @@ Filter = Union[
     _Not,
     _EntityTypeFilter,
     _EntitySubtypeFilter,
+    _StatusFilter,
     _PlatformFilter,
     _DomainFilter,
     _EnvFilter,
@@ -312,6 +358,18 @@ class FilterDsl:
     def not_(arg: "Filter") -> _Not:
         return _Not(not_=arg)
 
+    @staticmethod
+    def true() -> "Filter":
+        return _CustomCondition(
+            field="urn",
+            condition="EXISTS",
+            values=[],
+        )
+
+    @staticmethod
+    def false() -> "Filter":
+        return FilterDsl.not_(FilterDsl.true())
+
     @staticmethod
     def entity_type(
         entity_type: Union[EntityTypeName, Sequence[EntityTypeName]],
@@ -354,6 +412,10 @@ class FilterDsl:
             values=[f"{key}={value}"],
         )
 
+    @staticmethod
+    def soft_deleted(status: RemovedStatusFilter) -> _StatusFilter:
+        return _StatusFilter(status=status)
+
     # TODO: Add a soft-deletion status filter
     # TODO: add a container / browse path filter
     # TODO add shortcut for custom filters

datahub/testing/mcp_diff.py

@@ -189,7 +189,7 @@ class MCPDiff:
         """
         aspect_diffs = [v for d in self.aspect_changes.values() for v in d.values()]
         for aspect_diff in aspect_diffs:
-            for _, old, new in aspect_diff.aspects_changed.keys():
+            for _, old, new in aspect_diff.aspects_changed:
                 golden[old.delta_info.idx] = new.delta_info.original
 
         indices_to_remove = set()

datahub/utilities/file_backed_collections.py

@@ -250,7 +250,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                 rowid INTEGER PRIMARY KEY AUTOINCREMENT,
                 key TEXT UNIQUE,
                 value BLOB
-                {"".join(f", {column_name} BLOB" for column_name in self.extra_columns.keys())}
+                {"".join(f", {column_name} BLOB" for column_name in self.extra_columns)}
             )"""
         )
 
@@ -267,7 +267,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
         if self.indexes_created:
             return
         # The key column will automatically be indexed, but we need indexes for the extra columns.
-        for column_name in self.extra_columns.keys():
+        for column_name in self.extra_columns:
             self._conn.execute(
                 f"CREATE INDEX {self.tablename}_{column_name} ON {self.tablename} ({column_name})"
             )
@@ -305,12 +305,12 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                 f"""INSERT INTO {self.tablename} (
                     key,
                     value
-                    {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name}" for column_name in self.extra_columns)}
                 )
                 VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})
                 ON CONFLICT (key) DO UPDATE SET
                     value = excluded.value
-                    {"".join(f", {column_name} = excluded.{column_name}" for column_name in self.extra_columns.keys())}
+                    {"".join(f", {column_name} = excluded.{column_name}" for column_name in self.extra_columns)}
                 """,
                 items_to_write,
             )
@@ -321,7 +321,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                         f"""INSERT INTO {self.tablename} (
                             key,
                             value
-                            {"".join(f", {column_name}" for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name}" for column_name in self.extra_columns)}
                         )
                         VALUES ({", ".join(["?"] * (2 + len(self.extra_columns)))})""",
                         item,
@@ -330,7 +330,7 @@ class FileBackedDict(MutableMapping[str, _VT], Closeable, Generic[_VT]):
                     self._conn.execute(
                         f"""UPDATE {self.tablename} SET
                             value = ?
-                            {"".join(f", {column_name} = ?" for column_name in self.extra_columns.keys())}
+                            {"".join(f", {column_name} = ?" for column_name in self.extra_columns)}
                         WHERE key = ?""",
                         (*item[1:], item[0]),
                     )

datahub/utilities/hive_schema_to_avro.py

@@ -155,7 +155,7 @@ class HiveColumnToAvroConverter:
 
     @staticmethod
     def _parse_basic_datatype_string(s: str) -> Dict[str, object]:
-        if s in HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE.keys():
+        if s in HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE:
             return {
                 "type": HiveColumnToAvroConverter._PRIVIMITE_HIVE_TYPE_TO_AVRO_TYPE[s],
                 "native_data_type": s,
@@ -218,7 +218,7 @@ class HiveColumnToAvroConverter:
         buf = ""
         level = 0
         for c in s:
-            if c in HiveColumnToAvroConverter._BRACKETS.keys():
+            if c in HiveColumnToAvroConverter._BRACKETS:
                 level += 1
                 buf += c
             elif c in HiveColumnToAvroConverter._BRACKETS.values():

datahub/utilities/ingest_utils.py

@@ -32,10 +32,10 @@ def deploy_source_vars(
     name: Optional[str],
     config: str,
     urn: Optional[str],
-    executor_id: str,
+    executor_id: Optional[str],
     cli_version: Optional[str],
     schedule: Optional[str],
-    time_zone: str,
+    time_zone: Optional[str],
     extra_pip: Optional[str],
     debug: bool = False,
 ) -> dict:

datahub/utilities/threaded_iterator_executor.py

@@ -1,7 +1,15 @@
 import concurrent.futures
 import contextlib
 import queue
-from typing import Any, Callable, Generator, Iterable, Tuple, TypeVar
+from typing import (
+    Any,
+    Callable,
+    Iterable,
+    Iterator,
+    Optional,
+    Tuple,
+    TypeVar,
+)
 
 T = TypeVar("T")
 
@@ -18,8 +26,13 @@ class ThreadedIteratorExecutor:
         worker_func: Callable[..., Iterable[T]],
         args_list: Iterable[Tuple[Any, ...]],
         max_workers: int,
-    ) -> Generator[T, None, None]:
-        out_q: queue.Queue[T] = queue.Queue()
+        max_backpressure: Optional[int] = None,
+    ) -> Iterator[T]:
+        if max_backpressure is None:
+            max_backpressure = 10 * max_workers
+        assert max_backpressure >= max_workers
+
+        out_q: queue.Queue[T] = queue.Queue(maxsize=max_backpressure)
 
         def _worker_wrapper(
             worker_func: Callable[..., Iterable[T]], *args: Any

datahub/ingestion/transformer/system_metadata_transformer.py

@@ -1,45 +0,0 @@
-import functools
-from typing import Iterable
-
-from datahub.emitter.mce_builder import get_sys_time
-from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
-from datahub.ingestion.api.transform import Transformer
-from datahub.ingestion.api.workunit import MetadataWorkUnit
-from datahub.ingestion.transformer.auto_helper_transformer import AutoHelperTransformer
-from datahub.metadata.schema_classes import SystemMetadataClass
-
-
-def auto_system_metadata(
-    ctx: PipelineContext,
-    stream: Iterable[MetadataWorkUnit],
-) -> Iterable[MetadataWorkUnit]:
-    if not ctx.pipeline_config:
-        raise ValueError("Pipeline config is required for system metadata")
-    set_system_metadata = ctx.pipeline_config.flags.set_system_metadata
-    set_pipeline_name = ctx.pipeline_config.flags.set_system_metadata_pipeline_name
-
-    for workunit in stream:
-        if set_system_metadata:
-            workunit.metadata.systemMetadata = SystemMetadataClass(
-                lastObserved=get_sys_time(), runId=ctx.run_id
-            )
-            if set_pipeline_name:
-                workunit.metadata.systemMetadata.pipelineName = ctx.pipeline_name
-
-        yield workunit
-
-
-class SystemMetadataTransformer(Transformer):
-    def __init__(self, ctx: PipelineContext):
-        self._inner_transformer = AutoHelperTransformer(
-            functools.partial(auto_system_metadata, ctx)
-        )
-
-    def transform(
-        self, record_envelopes: Iterable[RecordEnvelope]
-    ) -> Iterable[RecordEnvelope]:
-        yield from self._inner_transformer.transform(record_envelopes)
-
-    @classmethod
-    def create(cls, config_dict: dict, ctx: PipelineContext) -> Transformer:
-        raise NotImplementedError(f"{cls.__name__} cannot be created from config")