acryl-datahub 1.0.0.1rc6__py3-none-any.whl → 1.0.0.1rc7__py3-none-any.whl

{acryl_datahub-1.0.0.1rc6.dist-info → acryl_datahub-1.0.0.1rc7.dist-info}/RECORD

@@ -1,9 +1,9 @@
-acryl_datahub-1.0.0.1rc6.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
+acryl_datahub-1.0.0.1rc7.dist-info/licenses/LICENSE,sha256=9xNHpsD0uYF5ONzXsKDCuHHB-xbiCrSbueWXqrTNsxk,11365
 datahub/__init__.py,sha256=aq_i5lVREmoLfYIqcx_pEQicO855YlhD19tWc1eZZNI,59
 datahub/__main__.py,sha256=pegIvQ9hzK7IhqVeUi1MeADSZ2QlP-D3K0OQdEg55RU,106
-datahub/_version.py,sha256=7kPZA7tlu74CuyNJP7xcQUj0ju3f89-XcI4tNSrX7u0,323
+datahub/_version.py,sha256=AZj-rwp4edRcZvS9Mq4fxTeV64QHFW-6zysNAtjc2qg,323
 datahub/entrypoints.py,sha256=2TYgHhs3sCxJlojIHjqfxzt3_ImPwPzq4vBtsUuMqu4,8885
-datahub/errors.py,sha256=bwtiNzFdVFze0IVKDEXQutkwk5j7cZkfXCUYCZIDSYg,565
+datahub/errors.py,sha256=BzKdcmYseHOt36zfjJXc17WNutFhp9Y23cU_L6cIkxc,612
 datahub/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/_codegen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/_codegen/aspect.py,sha256=PJRa-Z4ouXHq3OkulfyWhwZn-fFUBDK_UPvmqaWdbWk,1063
@@ -145,8 +145,8 @@ datahub/ingestion/api/registry.py,sha256=LbdZr89465Lj7ptQRVB4vI1JR1igWABvQFj9-WX
 datahub/ingestion/api/report.py,sha256=eM_TWWz6iJNd-c_S2_4eg2qKLGYP8vSROb_TMiCwBhY,4644
 datahub/ingestion/api/report_helpers.py,sha256=WbUC1kQeaKqIagGV3XzfPmPs7slAT1mfNY4og2BH2A8,994
 datahub/ingestion/api/sink.py,sha256=nfal7nsYY1AT2WQRjqO48uAHitpjax7TsRVzYXnqbeM,4918
-datahub/ingestion/api/source.py,sha256=Kz8xo0IY_5O3p5WE1i5dTmSK9IU20nqo4x6fvWcMAYw,19303
-datahub/ingestion/api/source_helpers.py,sha256=poP6EvkLsaiPM5fhjS5bcf-JMHMdPNMLv-eXCIVMUzM,19971
+datahub/ingestion/api/source.py,sha256=HrQahSEBeapMDnW8S6wSEyNLLE9RCs2R6eUrVaibuuc,19349
+datahub/ingestion/api/source_helpers.py,sha256=OhgBLdpUIuqF_gl4uV8Y2csp-z97zzXeFj2I5aUypCI,20158
 datahub/ingestion/api/transform.py,sha256=X0GpjMJzYkLuZx8MTWxH50cWGm9rGsnn3k188mmC8J8,582
 datahub/ingestion/api/workunit.py,sha256=e8n8RfSjHZZm2R4ShNH0UuMtUkMjyqqM2j2t7oL74lo,6327
 datahub/ingestion/api/auto_work_units/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -171,11 +171,11 @@ datahub/ingestion/glossary/classifier.py,sha256=daLxnVv_JlfB_jBOxH5LrU_xQRndrsGo
 datahub/ingestion/glossary/classifier_registry.py,sha256=yFOYLQhDgCLqXYMG3L1BquXafeLcZDcmp8meyw6k9ts,307
 datahub/ingestion/glossary/datahub_classifier.py,sha256=O7wm6gQT1Jf2QSKdWjJQbS5oSzJwplXzfza26Gdq5Mg,7555
 datahub/ingestion/graph/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/graph/client.py,sha256=791U-QMJXG3_RuNiQ4ennQ6NsOPQToSeKELHbncwzIQ,65573
+datahub/ingestion/graph/client.py,sha256=Qtjf5YrQeQzcTb0qxr6-y4MSEKSJm8f0hO6BoeRA_yI,65916
 datahub/ingestion/graph/config.py,sha256=_oha8Je7P80ZmrkZUAaRHyYbdMmTkMI5JkYjEP2Ri1Q,751
 datahub/ingestion/graph/connections.py,sha256=9462L0ZWGKURyypAln25eMPhK3pcufBar9tNDoqspXs,741
 datahub/ingestion/graph/entity_versioning.py,sha256=nrcNz0Qm6kpE6oTu_mrYUQDx14KPspBTc6R9SyFUY6c,6901
-datahub/ingestion/graph/filters.py,sha256=VFZKmef7ay1sQ5zRDDC1M_i6T96VzIgs-FzMs5eibiQ,7347
+datahub/ingestion/graph/filters.py,sha256=hZ8YOQRxC0_mbAx_SLkgqyYXr0Fw3O4U2wo2UMuDHJY,8653
 datahub/ingestion/reporting/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py,sha256=iEulcZMLBQuUfe9MAYyobMekvMcNm4dqVcS_C_2KfrI,9736
 datahub/ingestion/reporting/file_reporter.py,sha256=tiWukmMxHrTQI3rOAumsq6lRlw8T6spqpS6XBDYnrZU,1640
@@ -263,7 +263,7 @@ datahub/ingestion/source/bigquery_v2/queries.py,sha256=c1BpeQP8p8y-FOhmiQkkY2IqG
 datahub/ingestion/source/bigquery_v2/queries_extractor.py,sha256=_5cAXVU8b8T_nAPDsvN2JRd2dmM1t1J1mRylfKiPen4,19530
 datahub/ingestion/source/bigquery_v2/usage.py,sha256=A9c-ofclaRk0NSnc4IRaqJYqMPv6ecCld_TPy3V2qFs,40748
 datahub/ingestion/source/cassandra/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-datahub/ingestion/source/cassandra/cassandra.py,sha256=eIB2qiE6UaXCDmogWHhEafoPrBu3nlHsgc_1-2udOcc,14834
+datahub/ingestion/source/cassandra/cassandra.py,sha256=lKvPP0Uahi9xw_yh9cArPPtwvAauXolaEk-6f-jhpz4,14558
 datahub/ingestion/source/cassandra/cassandra_api.py,sha256=UVGQTsk6O57Q6wrWo54bQPLtStTWhw_Fq6fgW3Bjgk8,12515
 datahub/ingestion/source/cassandra/cassandra_config.py,sha256=vIMUOzazWTGi03B51vI0-YMxaMJHUGmCxJJgd8pKhC8,3791
 datahub/ingestion/source/cassandra/cassandra_profiling.py,sha256=DkSIryZNwLei5PaKuu9fNEKxEbhIrPI-T9gaVoM87NQ,11063
@@ -906,8 +906,8 @@ datahub/sdk/entity.py,sha256=Q29AbpS58L4gD8ETwoNIwG-ouytz4c0MSSFi6-jLl_4,6742
 datahub/sdk/entity_client.py,sha256=Sxe6H6Vr_tqLJu5KW7MJfLWJ6mgh4mbsx7u7MOBpM64,5052
 datahub/sdk/main_client.py,sha256=h2MKRhR-BO0zGCMhF7z2bTncX4hagKrAYwR3wTNTtzA,3666
 datahub/sdk/resolver_client.py,sha256=nKMAZJt2tRSGfKSzoREIh43PXqjM3umLiYkYHJjo1io,3243
-datahub/sdk/search_client.py,sha256=h9O_rsphkTdpd5hMPay3xSXfJM761cf4PjNCBwCnFzU,1309
-datahub/sdk/search_filters.py,sha256=WaJKFUKT9P70NBkh36f44rZrJ7zRJpRwu8mN-rEx5y0,11364
+datahub/sdk/search_client.py,sha256=BJR5t7Ff2oDNOGLcSCp9YHzrGKbgOQr7T8XQKGEpucw,3437
+datahub/sdk/search_filters.py,sha256=BcMhvG5hGYAATtLPLz4WLRjKApX2oLYrrcGn-CG__ek,12901
 datahub/secret/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 datahub/secret/datahub_secret_store.py,sha256=9u9S87-15jwhj4h0EsAVIMdQLgvstKc8voQux2slxgU,2477
 datahub/secret/datahub_secrets_client.py,sha256=nDmhziKdvseJHlaDVUcAwK8Fv8maeAaG-ktZtWG2b70,1316
@@ -996,7 +996,7 @@ datahub/utilities/sqllineage_patch.py,sha256=0Buh50bmEqJFg1HFRCknCnePo1cecI4JmGx
 datahub/utilities/stats_collections.py,sha256=CxaTcrF7J6am7iX5jPhFKne535UcyDk_oreVwR013fU,1625
 datahub/utilities/str_enum.py,sha256=EsqCLPbrqyQ2YU_wt7QP-a6P5fnpIshXJ3AI8gLBlVA,474
 datahub/utilities/tee_io.py,sha256=jBrsUfTPTk9IICntfGOG0HR-Fjp8BQMde-FPQ4r3kuI,601
-datahub/utilities/threaded_iterator_executor.py,sha256=WC4tvJ4TQRkH0VO_FD91GbedcKUqx0lc4tHDNOiF6ps,1770
+datahub/utilities/threaded_iterator_executor.py,sha256=6BpCE0os3d-uMYxHBilPQC-JvEBkU6JQY4bGs06JKYI,2004
 datahub/utilities/threading_timeout.py,sha256=hOzDI55E3onXblHNwGsePJUWMXo5zqaWCnoYdL2-KPM,1316
 datahub/utilities/time.py,sha256=Q7S_Zyom8C2zcl2xFbjNw6K8nZsCub5XGAB4OEmIS34,1847
 datahub/utilities/topological_sort.py,sha256=kcK5zPSR393fgItr-KSLV3bDqfJfBRS8E5kkCpPBgUY,1358
@@ -1043,8 +1043,8 @@ datahub_provider/operators/datahub_assertion_operator.py,sha256=uvTQ-jk2F0sbqqxp
 datahub_provider/operators/datahub_assertion_sensor.py,sha256=lCBj_3x1cf5GMNpHdfkpHuyHfVxsm6ff5x2Z5iizcAo,140
 datahub_provider/operators/datahub_operation_operator.py,sha256=aevDp2FzX7FxGlXrR0khoHNbxbhKR2qPEX5e8O2Jyzw,174
 datahub_provider/operators/datahub_operation_sensor.py,sha256=8fcdVBCEPgqy1etTXgLoiHoJrRt_nzFZQMdSzHqSG7M,168
-acryl_datahub-1.0.0.1rc6.dist-info/METADATA,sha256=8vR0C58PC28apLdQmFrOElrjYK-ArD0warcqjSc-fD8,176849
-acryl_datahub-1.0.0.1rc6.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-acryl_datahub-1.0.0.1rc6.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
-acryl_datahub-1.0.0.1rc6.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
-acryl_datahub-1.0.0.1rc6.dist-info/RECORD,,
+acryl_datahub-1.0.0.1rc7.dist-info/METADATA,sha256=QkeMAnAXXez9FFTnJVpxhOJMNDadF0gfVQ3uz3Fh4i8,176849
+acryl_datahub-1.0.0.1rc7.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+acryl_datahub-1.0.0.1rc7.dist-info/entry_points.txt,sha256=o3mDeJXSKhsy7XLkuogihraiabBdLn9HaizYXPrxmk0,9710
+acryl_datahub-1.0.0.1rc7.dist-info/top_level.txt,sha256=iLjSrLK5ox1YVYcglRUkcvfZPvKlobBWx7CTUXx8_GI,25
+acryl_datahub-1.0.0.1rc7.dist-info/RECORD,,

datahub/_version.py

@@ -1,6 +1,6 @@
 # Published at https://pypi.org/project/acryl-datahub/.
 __package_name__ = "acryl-datahub"
-__version__ = "1.0.0.1rc6"
+__version__ = "1.0.0.1rc7"
 
 
 def is_dev_mode() -> bool:

datahub/errors.py

@@ -31,6 +31,10 @@ class MultipleSubtypesWarning(Warning):
     pass
 
 
+class SearchFilterWarning(Warning):
+    pass
+
+
 class ExperimentalWarning(Warning):
     pass
 

datahub/ingestion/api/source.py

@@ -51,6 +51,7 @@ from datahub.ingestion.api.source_helpers import (
 from datahub.ingestion.api.workunit import MetadataWorkUnit
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import MetadataChangeEvent
 from datahub.metadata.schema_classes import UpstreamLineageClass
+from datahub.sdk.entity import Entity
 from datahub.utilities.lossy_collections import LossyDict, LossyList
 from datahub.utilities.type_annotations import get_class_from_annotation
 
@@ -480,7 +481,7 @@ class Source(Closeable, metaclass=ABCMeta):
 
     def get_workunits_internal(
         self,
-    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper]]:
+    ) -> Iterable[Union[MetadataWorkUnit, MetadataChangeProposalWrapper, Entity]]:
         raise NotImplementedError(
             "get_workunits_internal must be implemented if get_workunits is not overriden."
         )

datahub/ingestion/api/source_helpers.py

@@ -35,6 +35,7 @@ from datahub.metadata.schema_classes import (
     TimeWindowSizeClass,
 )
 from datahub.metadata.urns import DatasetUrn, GlossaryTermUrn, TagUrn, Urn
+from datahub.sdk.entity import Entity
 from datahub.specific.dataset import DatasetPatchBuilder
 from datahub.telemetry import telemetry
 from datahub.utilities.urns.error import InvalidUrnError
@@ -49,7 +50,12 @@ logger = logging.getLogger(__name__)
 
 def auto_workunit(
     stream: Iterable[
-        Union[MetadataChangeEventClass, MetadataChangeProposalWrapper, MetadataWorkUnit]
+        Union[
+            MetadataChangeEventClass,
+            MetadataChangeProposalWrapper,
+            MetadataWorkUnit,
+            Entity,
+        ]
     ],
 ) -> Iterable[MetadataWorkUnit]:
     """Convert a stream of MCEs and MCPs to a stream of :class:`MetadataWorkUnit`s."""
@@ -62,6 +68,8 @@ def auto_workunit(
             )
         elif isinstance(item, MetadataChangeProposalWrapper):
             yield item.as_workunit()
+        elif isinstance(item, Entity):
+            yield from item.as_workunits()
         else:
             yield item
 

datahub/ingestion/graph/client.py

@@ -49,6 +49,7 @@ from datahub.ingestion.graph.connections import (
 )
 from datahub.ingestion.graph.entity_versioning import EntityVersioningAPI
 from datahub.ingestion.graph.filters import (
+    RawSearchFilter,
     RawSearchFilterRule,
     RemovedStatusFilter,
     generate_filter,
@@ -75,10 +76,11 @@ from datahub.metadata.schema_classes import (
     SystemMetadataClass,
     TelemetryClientIdClass,
 )
+from datahub.metadata.urns import CorpUserUrn, Urn
 from datahub.telemetry.telemetry import telemetry_instance
 from datahub.utilities.perf_timer import PerfTimer
 from datahub.utilities.str_enum import StrEnum
-from datahub.utilities.urns.urn import Urn, guess_entity_type
+from datahub.utilities.urns.urn import guess_entity_type
 
 if TYPE_CHECKING:
     from datahub.ingestion.sink.datahub_rest import (
@@ -116,7 +118,7 @@ def entity_type_to_graphql(entity_type: str) -> str:
     """Convert the entity types into GraphQL "EntityType" enum values."""
 
     # Hard-coded special cases.
-    if entity_type == "corpuser":
+    if entity_type == CorpUserUrn.ENTITY_TYPE:
         return "CORP_USER"
 
     # Convert camelCase to UPPER_UNDERSCORE.
@@ -133,6 +135,14 @@ def entity_type_to_graphql(entity_type: str) -> str:
     return entity_type
 
 
+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)
+
+
 class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def __init__(self, config: DatahubClientConfig) -> None:
         self.config = config
@@ -805,7 +815,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
 
         :return: An iterable of (urn, schema info) tuple that match the filters.
         """
-        types = [entity_type_to_graphql("dataset")]
+        types = self._get_types(["dataset"])
 
         # Add the query default of * if no query is specified.
         query = query or "*"
@@ -873,10 +883,10 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         env: Optional[str] = None,
         query: Optional[str] = None,
         container: Optional[str] = None,
-        status: RemovedStatusFilter = RemovedStatusFilter.NOT_SOFT_DELETED,
+        status: Optional[RemovedStatusFilter] = RemovedStatusFilter.NOT_SOFT_DELETED,
         batch_size: int = 10000,
         extraFilters: Optional[List[RawSearchFilterRule]] = None,
-        extra_or_filters: Optional[List[Dict[str, List[RawSearchFilterRule]]]] = None,
+        extra_or_filters: Optional[RawSearchFilter] = None,
     ) -> Iterable[str]:
         """Fetch all urns that match all of the given filters.
 
@@ -968,7 +978,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         status: RemovedStatusFilter = RemovedStatusFilter.NOT_SOFT_DELETED,
         batch_size: int = 10000,
         extra_and_filters: Optional[List[RawSearchFilterRule]] = None,
-        extra_or_filters: Optional[List[Dict[str, List[RawSearchFilterRule]]]] = None,
+        extra_or_filters: Optional[RawSearchFilter] = None,
         extra_source_fields: Optional[List[str]] = None,
         skip_cache: bool = False,
     ) -> Iterable[dict]:
@@ -1121,7 +1131,8 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
                 )
 
             types = [
-                entity_type_to_graphql(entity_type) for entity_type in entity_types
+                flexible_entity_type_to_graphql(entity_type)
+                for entity_type in entity_types
             ]
         return types
 

datahub/ingestion/graph/filters.py

@@ -1,6 +1,7 @@
 import dataclasses
 import enum
-from typing import Any, Dict, List, Literal, Optional
+import warnings
+from typing import Dict, List, Literal, Optional, Union
 
 from typing_extensions import TypeAlias
 
@@ -8,9 +9,14 @@ from datahub.emitter.mce_builder import (
     make_data_platform_urn,
     make_dataplatform_instance_urn,
 )
+from datahub.errors import SearchFilterWarning
 from datahub.utilities.urns.urn import guess_entity_type
 
-RawSearchFilterRule = Dict[str, Any]
+RawSearchFilterRule: TypeAlias = Dict[str, Union[str, bool, List[str]]]
+
+# This is a list of OR filters, each of which is a list of AND filters.
+# This can be put directly into the orFilters parameter in GraphQL.
+RawSearchFilter: TypeAlias = List[Dict[Literal["and"], List[RawSearchFilterRule]]]
 
 # Mirrors our GraphQL enum: https://datahubproject.io/docs/graphql/enums#filteroperator
 FilterOperator: TypeAlias = Literal[
@@ -39,12 +45,14 @@ class SearchFilterRule:
     negated: bool = False
 
     def to_raw(self) -> RawSearchFilterRule:
-        return {
+        rule: RawSearchFilterRule = {
             "field": self.field,
             "condition": self.condition,
             "values": self.values,
-            "negated": self.negated,
         }
+        if self.negated:
+            rule["negated"] = True
+        return rule
 
     def negate(self) -> "SearchFilterRule":
         return SearchFilterRule(
@@ -73,10 +81,10 @@ def generate_filter(
     platform_instance: Optional[str],
     env: Optional[str],
     container: Optional[str],
-    status: RemovedStatusFilter,
+    status: Optional[RemovedStatusFilter],
     extra_filters: Optional[List[RawSearchFilterRule]],
-    extra_or_filters: Optional[List[RawSearchFilterRule]] = None,
-) -> List[Dict[str, List[RawSearchFilterRule]]]:
+    extra_or_filters: Optional[RawSearchFilter] = None,
+) -> RawSearchFilter:
     """
     Generate a search filter based on the provided parameters.
     :param platform: The platform to filter by.
@@ -105,15 +113,16 @@ def generate_filter(
         and_filters.append(_get_container_filter(container).to_raw())
 
     # Status filter.
-    status_filter = _get_status_filter(status)
-    if status_filter:
-        and_filters.append(status_filter.to_raw())
+    if status:
+        status_filter = _get_status_filter(status)
+        if status_filter:
+            and_filters.append(status_filter.to_raw())
 
     # Extra filters.
     if extra_filters:
         and_filters += extra_filters
 
-    or_filters: List[Dict[str, List[RawSearchFilterRule]]] = [{"and": and_filters}]
+    or_filters: RawSearchFilter = [{"and": and_filters}]
 
     # Env filter
     if env:
@@ -127,11 +136,27 @@ def generate_filter(
 
     # Extra OR filters are distributed across the top level and lists.
     if extra_or_filters:
-        or_filters = [
-            {"and": and_filter["and"] + [extra_or_filter]}
-            for extra_or_filter in extra_or_filters
-            for and_filter in or_filters
-        ]
+        new_or_filters: RawSearchFilter = []
+        for and_filter in or_filters:
+            for extra_or_filter in extra_or_filters:
+                if isinstance(extra_or_filter, dict) and "and" in extra_or_filter:
+                    new_or_filters.append(
+                        {"and": and_filter["and"] + extra_or_filter["and"]}
+                    )
+                else:
+                    # Hack for backwards compatibility.
+                    # We have some code that erroneously passed a List[RawSearchFilterRule]
+                    # instead of a List[Dict["and", List[RawSearchFilterRule]]].
+                    warnings.warn(
+                        "Passing a List[RawSearchFilterRule] to extra_or_filters is deprecated. "
+                        "Please pass a List[Dict[str, List[RawSearchFilterRule]]] instead.",
+                        SearchFilterWarning,
+                        stacklevel=3,
+                    )
+                    new_or_filters.append(
+                        {"and": and_filter["and"] + [extra_or_filter]}  # type: ignore
+                    )
+        or_filters = new_or_filters
 
     return or_filters
 

datahub/ingestion/source/cassandra/cassandra.py

@@ -123,16 +123,7 @@ class CassandraSource(StatefulIngestionSourceBase):
             ).workunit_processor,
         ]
 
-    def get_workunits_internal(
-        self,
-    ) -> Iterable[MetadataWorkUnit]:
-        for metadata in self._get_metadata():
-            if isinstance(metadata, MetadataWorkUnit):
-                yield metadata
-            else:
-                yield from metadata.as_workunits()
-
-    def _get_metadata(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
+    def get_workunits_internal(self) -> Iterable[Union[MetadataWorkUnit, Entity]]:
         if not self.cassandra_api.authenticate():
             return
         keyspaces: List[CassandraKeyspace] = self.cassandra_api.get_keyspaces()

datahub/sdk/search_client.py

@@ -2,36 +2,106 @@ from __future__ import annotations
 
 from typing import (
     TYPE_CHECKING,
-    Dict,
     Iterable,
     List,
     Optional,
+    Tuple,
+    Type,
+    TypeVar,
 )
 
-from datahub.ingestion.graph.filters import RawSearchFilterRule
+from datahub.ingestion.graph.filters import RawSearchFilter, RemovedStatusFilter
 from datahub.metadata.urns import Urn
-from datahub.sdk.search_filters import Filter
+from datahub.sdk.search_filters import (
+    Filter,
+    FilterDsl,
+    _EntityTypeFilter,
+    _OrFilters,
+    _StatusFilter,
+)
 
 if TYPE_CHECKING:
     from datahub.sdk.main_client import DataHubClient
 
 
+_FilterType = TypeVar("_FilterType", bound=Filter)
+
+
+def _typed_dfs(
+    filter: Optional[_FilterType], type: Type[_FilterType]
+) -> Optional[List[_FilterType]]:
+    if filter is None:
+        return None
+
+    found: Optional[List[_FilterType]] = None
+    for f in filter.dfs():
+        if isinstance(f, type):
+            if found is None:
+                found = []
+            found.append(f)
+    return found
+
+
 def compile_filters(
     filter: Optional[Filter],
-) -> Optional[List[Dict[str, List[RawSearchFilterRule]]]]:
+) -> Tuple[Optional[List[str]], RawSearchFilter]:
     # TODO: Not every filter type is supported for every entity type.
     # If we can detect issues with the filters at compile time, we should
     # raise an error.
 
-    if filter is None:
-        return None
+    existing_soft_deleted_filter = _typed_dfs(filter, _StatusFilter)
+    if existing_soft_deleted_filter is None:
+        soft_deleted_filter = FilterDsl.soft_deleted(
+            RemovedStatusFilter.NOT_SOFT_DELETED
+        )
+        if filter is None:
+            filter = soft_deleted_filter
+        else:
+            filter = FilterDsl.and_(filter, soft_deleted_filter)
+
+    # This should be safe - if filter were None coming in, then we would replace it
+    # with the soft-deleted filter.
+    assert filter is not None
 
     initial_filters = filter.compile()
-    return [
+
+    compiled_filters: RawSearchFilter = [
         {"and": [rule.to_raw() for rule in andClause["and"]]}
         for andClause in initial_filters
     ]
 
+    entity_types = compute_entity_types(initial_filters)
+
+    return entity_types, compiled_filters
+
+
+def compute_entity_types(
+    filters: _OrFilters,
+) -> Optional[List[str]]:
+    found_filters = False
+    found_positive_filters = False
+    entity_types: List[str] = []
+    for ands in filters:
+        for clause in ands["and"]:
+            if clause.field == _EntityTypeFilter.ENTITY_TYPE_FIELD:
+                found_filters = True
+                if not clause.negated:
+                    found_positive_filters = True
+
+                entity_types.extend(clause.values)
+
+    if not found_filters:
+        # If we didn't find any filters, use None so we use the default set.
+        return None
+
+    if not found_positive_filters:
+        # If we only found negated filters, then it's probably a query like
+        # "find me all entities except for dashboards". In that case, we
+        # still want to use the default set.
+        return None
+
+    return entity_types
+
 
 class SearchClient:
     def __init__(self, client: DataHubClient):
@@ -43,8 +113,11 @@ class SearchClient:
         filter: Optional[Filter] = None,
     ) -> Iterable[Urn]:
         # TODO: Add better limit / pagination support.
+        types, compiled_filters = compile_filters(filter)
         for urn in self._client._graph.get_urns_by_filter(
             query=query,
-            extra_or_filters=compile_filters(filter),
+            status=None,
+            extra_or_filters=compiled_filters,
+            entity_types=types,
         ):
             yield Urn.from_string(urn)

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