acryl-datahub 1.0.0rc17__py3-none-any.whl → 1.0.0.1__py3-none-any.whl

datahub/api/entities/dataset/dataset.py

@@ -506,7 +506,7 @@ class Dataset(StrictModel):
         # We don't check references for tags
         return list(set(references))
 
-    def generate_mcp(  # noqa: C901
+    def generate_mcp(
         self,
     ) -> Iterable[Union[MetadataChangeProposalClass, MetadataChangeProposalWrapper]]:
         mcp = MetadataChangeProposalWrapper(
@@ -643,33 +643,6 @@ class Dataset(StrictModel):
                     )
                     assert field_urn.startswith("urn:li:schemaField:")
 
-                    if field.globalTags:
-                        mcp = MetadataChangeProposalWrapper(
-                            entityUrn=field_urn,
-                            aspect=GlobalTagsClass(
-                                tags=[
-                                    TagAssociationClass(tag=make_tag_urn(tag))
-                                    for tag in field.globalTags
-                                ]
-                            ),
-                        )
-                        yield mcp
-
-                    if field.glossaryTerms:
-                        mcp = MetadataChangeProposalWrapper(
-                            entityUrn=field_urn,
-                            aspect=GlossaryTermsClass(
-                                terms=[
-                                    GlossaryTermAssociationClass(
-                                        urn=make_term_urn(term)
-                                    )
-                                    for term in field.glossaryTerms
-                                ],
-                                auditStamp=self._mint_auditstamp("yaml"),
-                            ),
-                        )
-                        yield mcp
-
                     if field.structured_properties:
                         urn_prefix = f"{StructuredPropertyUrn.URN_PREFIX}:{StructuredPropertyUrn.LI_DOMAIN}:{StructuredPropertyUrn.ENTITY_TYPE}"
                         mcp = MetadataChangeProposalWrapper(

datahub/cli/specific/dataset_cli.py

@@ -29,13 +29,16 @@ def dataset() -> None:
     name="upsert",
 )
 @click.option("-f", "--file", required=True, type=click.Path(exists=True))
+@click.option(
+    "-n", "--dry-run", type=bool, is_flag=True, default=False, help="Perform a dry run"
+)
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
-def upsert(file: Path) -> None:
+def upsert(file: Path, dry_run: bool) -> None:
     """Upsert attributes to a Dataset in DataHub."""
     # Call the sync command with to_datahub=True to perform the upsert operation
     ctx = click.get_current_context()
-    ctx.invoke(sync, file=str(file), to_datahub=True)
+    ctx.invoke(sync, file=str(file), dry_run=dry_run, to_datahub=True)
 
 
 @dataset.command(
@@ -167,11 +170,16 @@ def file(lintcheck: bool, lintfix: bool, file: str) -> None:
 )
 @click.option("-f", "--file", required=True, type=click.Path(exists=True))
 @click.option("--to-datahub/--from-datahub", required=True, is_flag=True)
+@click.option(
+    "-n", "--dry-run", type=bool, is_flag=True, default=False, help="Perform a dry run"
+)
 @upgrade.check_upgrade
 @telemetry.with_telemetry()
-def sync(file: str, to_datahub: bool) -> None:
+def sync(file: str, to_datahub: bool, dry_run: bool) -> None:
     """Sync a Dataset file to/from DataHub"""
 
+    dry_run_prefix = "[dry-run]: " if dry_run else ""  # prefix to use in messages
+
     failures: List[str] = []
     with get_default_graph() as graph:
         datasets = Dataset.from_yaml(file)
@@ -189,7 +197,7 @@ def sync(file: str, to_datahub: bool) -> None:
                     click.secho(
                         "\n\t- ".join(
                             [
-                                f"Skipping Dataset {dataset.urn} due to missing entity references: "
+                                f"{dry_run_prefix}Skipping Dataset {dataset.urn} due to missing entity references: "
                             ]
                             + missing_entity_references
                         ),
@@ -199,13 +207,18 @@ def sync(file: str, to_datahub: bool) -> None:
                     continue
                 try:
                     for mcp in dataset.generate_mcp():
-                        graph.emit(mcp)
-                    click.secho(f"Update succeeded for urn {dataset.urn}.", fg="green")
+                        if not dry_run:
+                            graph.emit(mcp)
+                    click.secho(
+                        f"{dry_run_prefix}Update succeeded for urn {dataset.urn}.",
+                        fg="green",
+                    )
                 except Exception as e:
                     click.secho(
-                        f"Update failed for id {id}. due to {e}",
+                        f"{dry_run_prefix}Update failed for id {id}. due to {e}",
                         fg="red",
                     )
+                    failures.append(dataset.urn)
             else:
                 # Sync from DataHub
                 if graph.exists(dataset.urn):
@@ -215,13 +228,16 @@ def sync(file: str, to_datahub: bool) -> None:
                     existing_dataset: Dataset = Dataset.from_datahub(
                         graph=graph, urn=dataset.urn, config=dataset_get_config
                     )
-                    existing_dataset.to_yaml(Path(file))
+                    if not dry_run:
+                        existing_dataset.to_yaml(Path(file))
+                    else:
+                        click.secho(f"{dry_run_prefix}Will update file {file}")
                 else:
-                    click.secho(f"Dataset {dataset.urn} does not exist")
+                    click.secho(f"{dry_run_prefix}Dataset {dataset.urn} does not exist")
                     failures.append(dataset.urn)
     if failures:
         click.secho(
-            f"\nFailed to sync the following Datasets: {', '.join(failures)}",
+            f"\n{dry_run_prefix}Failed to sync the following Datasets: {', '.join(failures)}",
             fg="red",
         )
         raise click.Abort()

datahub/emitter/mce_builder.py

@@ -125,9 +125,7 @@ def parse_ts_millis(ts: Optional[float]) -> Optional[datetime]:
 
 
 def make_data_platform_urn(platform: str) -> str:
-    if platform.startswith("urn:li:dataPlatform:"):
-        return platform
-    return DataPlatformUrn.create_from_id(platform).urn()
+    return DataPlatformUrn(platform).urn()
 
 
 def make_dataset_urn(platform: str, name: str, env: str = DEFAULT_ENV) -> str:

datahub/emitter/mcp_builder.py

@@ -117,6 +117,14 @@ class ContainerKey(DatahubKey):
 PlatformKey = ContainerKey
 
 
+class NamespaceKey(ContainerKey):
+    """
+    For Iceberg namespaces (databases/schemas)
+    """
+
+    namespace: str
+
+
 class DatabaseKey(ContainerKey):
     database: str
 

datahub/emitter/request_helper.py

@@ -1,8 +1,8 @@
-import itertools
 import shlex
-from typing import List, Union
+from typing import List, Optional, Union
 
 import requests
+from requests.auth import HTTPBasicAuth
 
 
 def _format_header(name: str, value: Union[str, bytes]) -> str:
@@ -12,17 +12,22 @@ def _format_header(name: str, value: Union[str, bytes]) -> str:
 
 
 def make_curl_command(
-    session: requests.Session, method: str, url: str, payload: str
+    session: requests.Session, method: str, url: str, payload: Optional[str] = None
 ) -> str:
-    fragments: List[str] = [
-        "curl",
-        *itertools.chain(
-            *[
-                ("-X", method),
-                *[("-H", _format_header(k, v)) for (k, v) in session.headers.items()],
-                ("--data", payload),
-            ]
-        ),
-        url,
-    ]
+    fragments: List[str] = ["curl", "-X", method]
+
+    for header_name, header_value in session.headers.items():
+        fragments.extend(["-H", _format_header(header_name, header_value)])
+
+    if session.auth:
+        if isinstance(session.auth, HTTPBasicAuth):
+            fragments.extend(["-u", f"{session.auth.username}:<redacted>"])
+        else:
+            # For other auth types, they should be handled via headers
+            fragments.extend(["-H", "<unknown auth type>"])
+
+    if payload:
+        fragments.extend(["--data", payload])
+
+    fragments.append(url)
     return shlex.join(fragments)

datahub/emitter/response_helper.py

@@ -1,17 +1,21 @@
 import json
 import logging
+import warnings
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Sequence, Union
 
 from requests import Response
 
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.errors import APITracingWarning
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeProposal,
 )
 
 logger = logging.getLogger(__name__)
 
+_TRACE_HEADER_NAME = "traceparent"
+
 
 @dataclass
 class TraceData:
@@ -25,14 +29,11 @@ class TraceData:
             raise TypeError("data must be a dictionary")
 
 
-def _extract_trace_id(
-    response: Response, trace_header: str = "traceparent"
-) -> Optional[str]:
+def _extract_trace_id(response: Response) -> Optional[str]:
     """
     Extract trace ID from response headers.
     Args:
         response: HTTP response object
-        trace_header: Name of the trace header to use
     Returns:
         Trace ID if found and response is valid, None otherwise
     """
@@ -40,9 +41,17 @@ def _extract_trace_id(
         logger.debug(f"Invalid status code: {response.status_code}")
         return None
 
-    trace_id = response.headers.get(trace_header)
+    trace_id = response.headers.get(_TRACE_HEADER_NAME)
     if not trace_id:
-        logger.debug(f"Missing trace header: {trace_header}")
+        # This will only be printed if
+        # 1. we're in async mode (checked by the caller)
+        # 2. the server did not return a trace ID
+        logger.debug(f"Missing trace header: {_TRACE_HEADER_NAME}")
+        warnings.warn(
+            "No trace ID found in response headers. API tracing is not active - likely due to an outdated server version.",
+            APITracingWarning,
+            stacklevel=3,
+        )
         return None
 
     return trace_id
@@ -51,20 +60,19 @@ def _extract_trace_id(
 def extract_trace_data(
     response: Response,
     aspects_to_trace: Optional[List[str]] = None,
-    trace_header: str = "traceparent",
 ) -> Optional[TraceData]:
-    """
-    Extract trace data from a response object.
+    """Extract trace data from a response object.
+
+    If we run into a JSONDecodeError, we'll log an error and return None.
+
     Args:
         response: HTTP response object
         aspects_to_trace: Optional list of aspect names to extract. If None, extracts all aspects.
-        trace_header: Name of the trace header to use (default: "traceparent")
+
     Returns:
         TraceData object if successful, None otherwise
-    Raises:
-        JSONDecodeError: If response body cannot be decoded as JSON
     """
-    trace_id = _extract_trace_id(response, trace_header)
+    trace_id = _extract_trace_id(response)
     if not trace_id:
         return None
 
@@ -104,19 +112,18 @@ def extract_trace_data_from_mcps(
     response: Response,
     mcps: Sequence[Union[MetadataChangeProposal, MetadataChangeProposalWrapper]],
     aspects_to_trace: Optional[List[str]] = None,
-    trace_header: str = "traceparent",
 ) -> Optional[TraceData]:
-    """
-    Extract trace data from a response object and populate data from provided MCPs.
+    """Extract trace data from a response object and populate data from provided MCPs.
+
     Args:
         response: HTTP response object used only for trace_id extraction
         mcps: List of MCP URN and aspect data
         aspects_to_trace: Optional list of aspect names to extract. If None, extracts all aspects.
-        trace_header: Name of the trace header to use (default: "traceparent")
+
     Returns:
         TraceData object if successful, None otherwise
     """
-    trace_id = _extract_trace_id(response, trace_header)
+    trace_id = _extract_trace_id(response)
     if not trace_id:
         return None
 

datahub/emitter/rest_emitter.py

@@ -5,6 +5,7 @@ import json
 import logging
 import os
 import time
+import warnings
 from collections import defaultdict
 from dataclasses import dataclass
 from datetime import datetime, timedelta
@@ -24,9 +25,9 @@ from typing import (
 
 import pydantic
 import requests
-from deprecated import deprecated
 from requests.adapters import HTTPAdapter, Retry
 from requests.exceptions import HTTPError, RequestException
+from typing_extensions import deprecated
 
 from datahub._version import nice_version_name
 from datahub.cli import config_utils
@@ -40,7 +41,7 @@ from datahub.configuration.common import (
     TraceTimeoutError,
     TraceValidationError,
 )
-from datahub.emitter.aspect import JSON_CONTENT_TYPE
+from datahub.emitter.aspect import JSON_CONTENT_TYPE, JSON_PATCH_CONTENT_TYPE
 from datahub.emitter.generic_emitter import Emitter
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.request_helper import make_curl_command
@@ -50,6 +51,7 @@ from datahub.emitter.response_helper import (
     extract_trace_data_from_mcps,
 )
 from datahub.emitter.serialization_helper import pre_json_transform
+from datahub.errors import APITracingWarning
 from datahub.ingestion.api.closeable import Closeable
 from datahub.metadata.com.linkedin.pegasus2avro.mxe import (
     MetadataChangeEvent,
@@ -107,9 +109,9 @@ class RestSinkEndpoint(ConfigEnum):
     OPENAPI = auto()
 
 
-DEFAULT_REST_SINK_ENDPOINT = pydantic.parse_obj_as(
+DEFAULT_REST_EMITTER_ENDPOINT = pydantic.parse_obj_as(
     RestSinkEndpoint,
-    os.getenv("DATAHUB_REST_SINK_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
+    os.getenv("DATAHUB_REST_EMITTER_DEFAULT_ENDPOINT", RestSinkEndpoint.RESTLI),
 )
 
 
@@ -227,7 +229,9 @@ class DataHubRestEmitter(Closeable, Emitter):
         ca_certificate_path: Optional[str] = None,
         client_certificate_path: Optional[str] = None,
         disable_ssl_verification: bool = False,
-        openapi_ingestion: bool = False,
+        openapi_ingestion: bool = (
+            DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI
+        ),
         default_trace_mode: bool = False,
     ):
         if not gms_server:
@@ -357,8 +361,14 @@ class DataHubRestEmitter(Closeable, Emitter):
                 )["aspect"]["json"]
             else:
                 obj = mcp.aspect.to_obj()
-                if obj.get("value") and obj.get("contentType") == JSON_CONTENT_TYPE:
+                content_type = obj.get("contentType")
+                if obj.get("value") and content_type == JSON_CONTENT_TYPE:
+                    # Undo double serialization.
                     obj = json.loads(obj["value"])
+                elif content_type == JSON_PATCH_CONTENT_TYPE:
+                    raise NotImplementedError(
+                        "Patches are not supported for OpenAPI ingestion. Set the endpoint to RESTLI."
+                    )
                 aspect_value = pre_json_transform(obj)
             return (
                 url,
@@ -597,7 +607,7 @@ class DataHubRestEmitter(Closeable, Emitter):
 
         return len(mcp_obj_chunks)
 
-    @deprecated
+    @deprecated("Use emit with a datasetUsageStatistics aspect instead")
     def emit_usage(self, usageStats: UsageAggregation) -> None:
         url = f"{self._gms_server}/usageStats?action=batchIngest"
 
@@ -749,6 +759,12 @@ class DataHubRestEmitter(Closeable, Emitter):
             trace_flag if trace_flag is not None else self._default_trace_mode
         )
         resolved_async_flag = async_flag if async_flag is not None else async_default
+        if resolved_trace_flag and not resolved_async_flag:
+            warnings.warn(
+                "API tracing is only available with async ingestion. For sync mode, API errors will be surfaced as exceptions.",
+                APITracingWarning,
+                stacklevel=3,
+            )
         return resolved_trace_flag and resolved_async_flag
 
     def __repr__(self) -> str:

datahub/errors.py

@@ -31,5 +31,13 @@ class MultipleSubtypesWarning(Warning):
     pass
 
 
+class SearchFilterWarning(Warning):
+    pass
+
+
 class ExperimentalWarning(Warning):
     pass
+
+
+class APITracingWarning(Warning):
+    pass

datahub/ingestion/api/source.py

@@ -27,6 +27,7 @@ from typing_extensions import LiteralString, Self
 
 from datahub.configuration.common import ConfigModel
 from datahub.configuration.source_common import PlatformInstanceConfigMixin
+from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.mcp_builder import mcps_from_mce
 from datahub.ingestion.api.auto_work_units.auto_dataset_properties_aspect import (
     auto_patch_last_modified,
@@ -44,11 +45,13 @@ from datahub.ingestion.api.source_helpers import (
     auto_lowercase_urns,
     auto_materialize_referenced_tags_terms,
     auto_status_aspect,
+    auto_workunit,
     auto_workunit_reporter,
 )
 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
 
@@ -473,10 +476,12 @@ class Source(Closeable, metaclass=ABCMeta):
 
     def get_workunits(self) -> Iterable[MetadataWorkUnit]:
         return self._apply_workunit_processors(
-            self.get_workunit_processors(), self.get_workunits_internal()
+            self.get_workunit_processors(), auto_workunit(self.get_workunits_internal())
         )
 
-    def get_workunits_internal(self) -> Iterable[MetadataWorkUnit]:
+    def get_workunits_internal(
+        self,
+    ) -> 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
@@ -48,7 +49,14 @@ logger = logging.getLogger(__name__)
 
 
 def auto_workunit(
-    stream: Iterable[Union[MetadataChangeEventClass, MetadataChangeProposalWrapper]],
+    stream: Iterable[
+        Union[
+            MetadataChangeEventClass,
+            MetadataChangeProposalWrapper,
+            MetadataWorkUnit,
+            Entity,
+        ]
+    ],
 ) -> Iterable[MetadataWorkUnit]:
     """Convert a stream of MCEs and MCPs to a stream of :class:`MetadataWorkUnit`s."""
 
@@ -58,8 +66,12 @@ def auto_workunit(
                 id=MetadataWorkUnit.generate_workunit_id(item),
                 mce=item,
             )
-        else:
+        elif isinstance(item, MetadataChangeProposalWrapper):
             yield item.as_workunit()
+        elif isinstance(item, Entity):
+            yield from item.as_workunits()
+        else:
+            yield item
 
 
 def create_dataset_props_patch_builder(

datahub/ingestion/extractor/schema_util.py

@@ -362,6 +362,7 @@ class AvroToMceSchemaConverter:
                 merged_props: Dict[str, Any] = {}
                 merged_props.update(self._schema.other_props)
                 merged_props.update(schema.other_props)
+                merged_props.update(actual_schema.other_props)
 
                 # Parse meta_mapping
                 meta_aspects: Dict[str, Any] = {}

datahub/ingestion/graph/client.py

@@ -23,9 +23,9 @@ from typing import (
 )
 
 from avro.schema import RecordSchema
-from deprecated import deprecated
 from pydantic import BaseModel
 from requests.models import HTTPError
+from typing_extensions import deprecated
 
 from datahub.cli import config_utils
 from datahub.configuration.common import ConfigModel, GraphError, OperationalError
@@ -33,7 +33,7 @@ from datahub.emitter.aspect import TIMESERIES_ASPECT_MAP
 from datahub.emitter.mce_builder import DEFAULT_ENV, Aspect
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.emitter.rest_emitter import (
-    DEFAULT_REST_SINK_ENDPOINT,
+    DEFAULT_REST_EMITTER_ENDPOINT,
     DEFAULT_REST_TRACE_MODE,
     DatahubRestEmitter,
     RestSinkEndpoint,
@@ -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
@@ -147,7 +157,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
             ca_certificate_path=self.config.ca_certificate_path,
             client_certificate_path=self.config.client_certificate_path,
             disable_ssl_verification=self.config.disable_ssl_verification,
-            openapi_ingestion=DEFAULT_REST_SINK_ENDPOINT == RestSinkEndpoint.OPENAPI,
+            openapi_ingestion=DEFAULT_REST_EMITTER_ENDPOINT == RestSinkEndpoint.OPENAPI,
             default_trace_mode=DEFAULT_REST_TRACE_MODE == RestTraceMode.ENABLED,
         )
 
@@ -330,7 +340,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
                 f"Failed to find {aspect_type_name} in response {response_json}"
             )
 
-    @deprecated(reason="Use get_aspect instead which makes aspect string name optional")
+    @deprecated("Use get_aspect instead which makes aspect string name optional")
     def get_aspect_v2(
         self,
         entity_urn: str,
@@ -355,7 +365,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def get_schema_metadata(self, entity_urn: str) -> Optional[SchemaMetadataClass]:
         return self.get_aspect(entity_urn=entity_urn, aspect_type=SchemaMetadataClass)
 
-    @deprecated(reason="Use get_aspect directly.")
+    @deprecated("Use get_aspect directly.")
     def get_domain_properties(self, entity_urn: str) -> Optional[DomainPropertiesClass]:
         return self.get_aspect(entity_urn=entity_urn, aspect_type=DomainPropertiesClass)
 
@@ -376,7 +386,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def get_domain(self, entity_urn: str) -> Optional[DomainsClass]:
         return self.get_aspect(entity_urn=entity_urn, aspect_type=DomainsClass)
 
-    @deprecated(reason="Use get_aspect directly.")
+    @deprecated("Use get_aspect directly.")
     def get_browse_path(self, entity_urn: str) -> Optional[BrowsePathsClass]:
         return self.get_aspect(entity_urn=entity_urn, aspect_type=BrowsePathsClass)
 
@@ -505,7 +515,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
         return response.json()
 
     @deprecated(
-        reason="Use get_aspect for a single aspect or get_entity_semityped for a full entity."
+        "Use get_aspect for a single aspect or get_entity_semityped for a full entity."
     )
     def get_aspects_for_entity(
         self,
@@ -635,9 +645,6 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
     def _aspect_count_endpoint(self):
         return f"{self.config.server}/aspects?action=getCount"
 
-    # def _session(self) -> Session:
-    #    return super()._session
-
     def get_domain_urn_by_name(self, domain_name: str) -> Optional[str]:
         """Retrieve a domain urn based on its name. Returns None if there is no match found"""
 
@@ -749,9 +756,7 @@ class DataHubGraph(DatahubRestEmitter, EntityVersioningAPI):
 
         assert res["upsertConnection"]["urn"] == urn
 
-    @deprecated(
-        reason='Use get_urns_by_filter(entity_types=["container"], ...) instead'
-    )
+    @deprecated('Use get_urns_by_filter(entity_types=["container"], ...) instead')
     def get_container_urns_by_filter(
         self,
         env: Optional[str] = None,
@@ -810,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 "*"
@@ -878,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.
 
@@ -973,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]:
@@ -1126,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