acryl-datahub 1.0.0rc18__py3-none-any.whl → 1.3.0.1rc9__py3-none-any.whl

datahub/ingestion/graph/config.py

@@ -1,13 +1,22 @@
+from enum import Enum, auto
 from typing import Dict, List, Optional
 
 from datahub.configuration.common import ConfigModel
+from datahub.configuration.env_vars import get_datahub_component
+
+
+class ClientMode(Enum):
+    INGESTION = auto()
+    CLI = auto()
+    SDK = auto()
+
+
+DATAHUB_COMPONENT_ENV: str = get_datahub_component().lower()
 
 
 class DatahubClientConfig(ConfigModel):
     """Configuration class for holding connectivity to datahub gms"""
 
-    # TODO: Having a default for the server doesn't make a ton of sense. This should be handled
-    # by callers / the CLI, but the actual client should not have any magic.
     server: str
     token: Optional[str] = None
     timeout_sec: Optional[float] = None
@@ -17,3 +26,10 @@ class DatahubClientConfig(ConfigModel):
     ca_certificate_path: Optional[str] = None
     client_certificate_path: Optional[str] = None
     disable_ssl_verification: bool = False
+    openapi_ingestion: Optional[bool] = None
+    client_mode: Optional[ClientMode] = None
+    datahub_component: Optional[str] = None
+    server_config_refresh_interval: Optional[int] = None
+
+    class Config:
+        extra = "ignore"

datahub/ingestion/graph/filters.py

@@ -1,30 +1,58 @@
 import dataclasses
 import enum
-from typing import Any, Dict, List, Optional
+import warnings
+from typing import Dict, List, Literal, Optional, Union
+
+from typing_extensions import TypeAlias
 
 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://docs.datahub.com/docs/graphql/enums#filteroperator
+FilterOperator: TypeAlias = Literal[
+    "CONTAIN",
+    "EQUAL",
+    "IEQUAL",
+    "IN",
+    "EXISTS",
+    "GREATER_THAN",
+    "GREATER_THAN_OR_EQUAL_TO",
+    "LESS_THAN",
+    "LESS_THAN_OR_EQUAL_TO",
+    "START_WITH",
+    "END_WITH",
+    "DESCENDANTS_INCL",
+    "ANCESTORS_INCL",
+    "RELATED_INCL",
+]
 
 
 @dataclasses.dataclass
 class SearchFilterRule:
     field: str
-    condition: str  # TODO: convert to an enum
+    condition: FilterOperator
     values: List[str]
     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(
@@ -48,15 +76,25 @@ class RemovedStatusFilter(enum.Enum):
     """Search only soft-deleted entities."""
 
 
+def _validate_or_filter_structure(
+    or_filters: List[Dict[str, List[SearchFilterRule]]],
+) -> None:
+    for filter_list in or_filters:
+        if "and" not in filter_list:
+            raise ValueError(f"Invalid or filter: {filter_list}")
+        if not isinstance(filter_list["and"], list):
+            raise ValueError(f"Invalid or filter: {filter_list}")
+
+
 def generate_filter(
-    platform: Optional[str],
+    platform: Union[None, str, List[str]],
     platform_instance: Optional[str],
     env: Optional[str],
-    container: Optional[str],
-    status: RemovedStatusFilter,
+    container: Union[None, str, List[str]],
+    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.
@@ -65,8 +103,7 @@ def generate_filter(
     :param container: The container to filter by.
     :param status: The status to filter by.
     :param extra_filters: Extra AND filters to apply.
-    :param extra_or_filters: Extra OR filters to apply. These are combined with
-    the AND filters using an OR at the top level.
+    :param extra_or_filters: Extra OR filters to apply. These are combined with the AND filters using an OR at the top level.
     """
     and_filters: List[RawSearchFilterRule] = []
 
@@ -85,15 +122,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:
@@ -107,11 +145,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
 
@@ -123,7 +177,7 @@ def _get_env_filters(env: str) -> List[RawSearchFilterRule]:
         # For most entity types, we look at the origin field.
         {
             "field": "origin",
-            "value": env,
+            "values": [env],
             "condition": "EQUAL",
         },
         # For containers, we look at the customProperties field.
@@ -131,15 +185,15 @@ def _get_env_filters(env: str) -> List[RawSearchFilterRule]:
         # we look for the "env" property. Otherwise, we use the "instance" property.
         {
             "field": "customProperties",
-            "value": f"env={env}",
+            "values": [f"env={env}"],
         },
         {
             "field": "customProperties",
-            "value": f"instance={env}",
+            "values": [f"instance={env}"],
         },
         {
             "field": "env",
-            "value": env,
+            "values": [env],
         },
         # Note that not all entity types have an env (e.g. dashboards / charts).
         # If the env filter is specified, these will be excluded.
@@ -173,23 +227,31 @@ def _get_status_filter(status: RemovedStatusFilter) -> Optional[SearchFilterRule
         raise ValueError(f"Invalid status filter: {status}")
 
 
-def _get_container_filter(container: str) -> SearchFilterRule:
+def _get_container_filter(container: Union[str, List[str]]) -> SearchFilterRule:
+    if not isinstance(container, list):
+        container = [container]
+
     # Warn if container is not a fully qualified urn.
     # TODO: Change this once we have a first-class container urn type.
-    if guess_entity_type(container) != "container":
-        raise ValueError(f"Invalid container urn: {container}")
+    for cont in container:
+        if guess_entity_type(cont) != "container":
+            raise ValueError(f"Invalid container urn: {cont}")
 
     return SearchFilterRule(
         field="browsePathV2",
-        values=[container],
+        values=container,
         condition="CONTAIN",
     )
 
 
 def _get_platform_instance_filter(
-    platform: Optional[str], platform_instance: str
+    platform: Union[None, str, List[str]], platform_instance: str
 ) -> SearchFilterRule:
     if platform:
+        if isinstance(platform, list):
+            raise ValueError(
+                "Platform instance filter cannot be combined with a multi-value platform filter."
+            )
         # Massage the platform instance into a fully qualified urn, if necessary.
         platform_instance = make_dataplatform_instance_urn(platform, platform_instance)
 
@@ -205,9 +267,11 @@ def _get_platform_instance_filter(
     )
 
 
-def _get_platform_filter(platform: str) -> SearchFilterRule:
+def _get_platform_filter(platform: Union[str, List[str]]) -> SearchFilterRule:
+    if not isinstance(platform, list):
+        platform = [platform]
     return SearchFilterRule(
         field="platform.keyword",
         condition="EQUAL",
-        values=[make_data_platform_urn(platform)],
+        values=[make_data_platform_urn(plt) for plt in platform],
     )

datahub/ingestion/graph/links.py

@@ -0,0 +1,55 @@
+import urllib.parse
+from typing import Optional
+
+import datahub.metadata.urns as urns
+from datahub.utilities.urns.urn import guess_entity_type
+
+_url_prefixes = {
+    # Atypical mappings.
+    urns.DataJobUrn.ENTITY_TYPE: "tasks",
+    urns.DataFlowUrn.ENTITY_TYPE: "pipelines",
+    urns.CorpUserUrn.ENTITY_TYPE: "user",
+    urns.CorpGroupUrn.ENTITY_TYPE: "group",
+    # Normal mappings - matches the entity type.
+    urns.ChartUrn.ENTITY_TYPE: "chart",
+    urns.ContainerUrn.ENTITY_TYPE: "container",
+    urns.DataProductUrn.ENTITY_TYPE: "dataProduct",
+    urns.DatasetUrn.ENTITY_TYPE: "dataset",
+    urns.DashboardUrn.ENTITY_TYPE: "dashboard",
+    urns.DomainUrn.ENTITY_TYPE: "domain",
+    urns.GlossaryNodeUrn.ENTITY_TYPE: "glossaryNode",
+    urns.GlossaryTermUrn.ENTITY_TYPE: "glossaryTerm",
+    urns.TagUrn.ENTITY_TYPE: "tag",
+}
+
+
+def make_url_for_urn(
+    frontend_base_url: str,
+    entity_urn: str,
+    *,
+    tab: Optional[str] = None,
+) -> str:
+    """Build the public-facing URL for an entity urn.
+
+    Args:
+        frontend_url: The public-facing base url of the frontend.
+        entity_urn: The urn of the entity to get the url for.
+        tab: The tab to deep link into. If not provided, the default tab for the entity will be shown.
+
+    Returns:
+        The public-facing url for the entity.
+
+    Examples:
+        >>> make_url_for_urn("https://demo.datahub.com", "urn:li:container:b41c14bc5cb3ccfbb0433c8cbdef2992", tab="Contents")
+        'https://demo.datahub.com/container/urn%3Ali%3Acontainer%3Ab41c14bc5cb3ccfbb0433c8cbdef2992/Contents'
+        >>> make_url_for_urn("https://demo.datahub.com", "urn:li:dataset:(urn:li:dataPlatform:snowflake,long_tail_companions.adoption.actuating,PROD)")
+        'https://demo.datahub.com/dataset/urn%3Ali%3Adataset%3A%28urn%3Ali%3AdataPlatform%3Asnowflake%2Clong_tail_companions.adoption.actuating%2CPROD%29/'
+    """
+    entity_type = guess_entity_type(entity_urn)
+    encoded_entity_urn = urllib.parse.quote(entity_urn, safe="")
+
+    url_prefix = _url_prefixes.get(entity_type, entity_type)
+    url = f"{frontend_base_url}/{url_prefix}/{encoded_entity_urn}/"
+    if tab:
+        url += f"{tab}"
+    return url

datahub/ingestion/reporting/datahub_ingestion_run_summary_provider.py

@@ -13,6 +13,7 @@ from datahub.configuration.common import (
 from datahub.emitter.aspect import JSON_CONTENT_TYPE
 from datahub.emitter.mce_builder import datahub_guid, make_data_platform_urn
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
+from datahub.emitter.rest_emitter import EmitMode
 from datahub.ingestion.api.common import PipelineContext, RecordEnvelope
 from datahub.ingestion.api.pipeline_run_listener import PipelineRunListener
 from datahub.ingestion.api.sink import NoopWriteCallback, Sink
@@ -111,6 +112,7 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
     def __init__(self, sink: Sink, report_recipe: bool, ctx: PipelineContext) -> None:
         assert ctx.pipeline_config is not None
 
+        self.ctx = ctx
         self.sink: Sink = sink
         self.report_recipe = report_recipe
         ingestion_source_key = self.generate_unique_key(ctx.pipeline_config)
@@ -191,18 +193,25 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
             )
             return json.dumps(converted_recipe)
 
-    def _emit_aspect(self, entity_urn: Urn, aspect_value: _Aspect) -> None:
-        self.sink.write_record_async(
-            RecordEnvelope(
-                record=MetadataChangeProposalWrapper(
-                    entityUrn=str(entity_urn),
-                    aspect=aspect_value,
-                ),
-                metadata={},
-            ),
-            NoopWriteCallback(),
+    def _emit_aspect(
+        self, entity_urn: Urn, aspect_value: _Aspect, try_sync: bool = False
+    ) -> None:
+        mcp = MetadataChangeProposalWrapper(
+            entityUrn=str(entity_urn),
+            aspect=aspect_value,
         )
 
+        if try_sync and self.ctx.graph:
+            self.ctx.graph.emit_mcp(mcp, emit_mode=EmitMode.SYNC_PRIMARY)
+        else:
+            self.sink.write_record_async(
+                RecordEnvelope(
+                    record=mcp,
+                    metadata={},
+                ),
+                NoopWriteCallback(),
+            )
+
     def on_start(self, ctx: PipelineContext) -> None:
         assert ctx.pipeline_config is not None
         # Construct the dataHubExecutionRequestInput aspect
@@ -223,6 +232,7 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
         self._emit_aspect(
             entity_urn=self.execution_request_input_urn,
             aspect_value=execution_input_aspect,
+            try_sync=True,
         )
 
     def on_completion(
@@ -258,4 +268,4 @@ class DatahubIngestionRunSummaryProvider(PipelineRunListener):
             entity_urn=self.execution_request_input_urn,
             aspect_value=execution_result_aspect,
         )
-        self.sink.close()
+        # Note: sink.close() is handled by the pipeline's context manager

datahub/ingestion/run/pipeline.py

@@ -31,6 +31,7 @@ from datahub.ingestion.api.source import Extractor, Source
 from datahub.ingestion.api.transform import Transformer
 from datahub.ingestion.extractor.extractor_registry import extractor_registry
 from datahub.ingestion.graph.client import DataHubGraph, get_default_graph
+from datahub.ingestion.graph.config import ClientMode
 from datahub.ingestion.reporting.reporting_provider_registry import (
     reporting_provider_registry,
 )
@@ -39,13 +40,14 @@ from datahub.ingestion.run.sink_callback import DeadLetterQueueCallback, Logging
 from datahub.ingestion.sink.datahub_rest import DatahubRestSink
 from datahub.ingestion.sink.sink_registry import sink_registry
 from datahub.ingestion.source.source_registry import source_registry
-from datahub.ingestion.transformer.system_metadata_transformer import (
-    SystemMetadataTransformer,
-)
 from datahub.ingestion.transformer.transform_registry import transform_registry
 from datahub.sdk._attribution import KnownAttribution, change_default_attribution
 from datahub.telemetry import stats
 from datahub.telemetry.telemetry import telemetry_instance
+from datahub.upgrade.upgrade import (
+    is_server_default_cli_ahead,
+    retrieve_version_stats,
+)
 from datahub.utilities._custom_package_loader import model_version_name
 from datahub.utilities.global_warning_util import (
     clear_global_warnings,
@@ -139,9 +141,8 @@ class CliReport(Report):
 
 
 def _make_default_rest_sink(ctx: PipelineContext) -> DatahubRestSink:
-    graph = get_default_graph()
+    graph = get_default_graph(ClientMode.INGESTION)
     sink_config = graph._make_rest_sink_config()
-
     return DatahubRestSink(ctx, sink_config)
 
 
@@ -174,10 +175,14 @@ class Pipeline:
         self.last_time_printed = int(time.time())
         self.cli_report = CliReport()
 
-        with contextlib.ExitStack() as exit_stack, contextlib.ExitStack() as inner_exit_stack:
+        with (
+            contextlib.ExitStack() as exit_stack,
+            contextlib.ExitStack() as inner_exit_stack,
+        ):
             self.graph: Optional[DataHubGraph] = None
             with _add_init_error_context("connect to DataHub"):
                 if self.config.datahub_api:
+                    self.config.datahub_api.client_mode = ClientMode.INGESTION
                     self.graph = exit_stack.enter_context(
                         DataHubGraph(self.config.datahub_api)
                     )
@@ -260,6 +265,11 @@ class Pipeline:
                 with _add_init_error_context("configure transformers"):
                     self._configure_transforms()
 
+                # Register completion callback with sink to handle final reporting
+                self.sink.register_pre_shutdown_callback(
+                    self._notify_reporters_on_ingestion_completion
+                )
+
             # If all of the initialization succeeds, we can preserve the exit stack until the pipeline run.
             # We need to use an exit stack so that if we have an exception during initialization,
             # things that were already initialized are still cleaned up.
@@ -286,9 +296,6 @@ class Pipeline:
                     f"Transformer type:{transformer_type},{transformer_class} configured"
                 )
 
-        # Add the system metadata transformer at the end of the list.
-        self.transformers.append(SystemMetadataTransformer(self.ctx))
-
     def _configure_reporting(self, report_to: Optional[str]) -> None:
         if self.dry_run:
             # In dry run mode, we don't want to report anything.
@@ -342,8 +349,48 @@ class Pipeline:
         for reporter in self.reporters:
             try:
                 reporter.on_start(ctx=self.ctx)
-            except Exception as e:
-                logger.warning("Reporting failed on start", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on start", exc_info=True)
+
+    def _warn_old_cli_version(self) -> None:
+        """
+        Check if the server default CLI version is ahead of the CLI version being used.
+        If so, add a warning to the report.
+        """
+
+        try:
+            version_stats = retrieve_version_stats(timeout=2.0, graph=self.graph)
+        except RuntimeError as e:
+            # Handle case where there's no event loop available (e.g., in ThreadPoolExecutor)
+            if "no current event loop" in str(e):
+                logger.debug("Skipping version check - no event loop available")
+                return
+            raise
+
+        if not version_stats or not self.graph:
+            return
+
+        if is_server_default_cli_ahead(version_stats):
+            server_default_version = (
+                version_stats.server.current_server_default_cli_version.version
+                if version_stats.server.current_server_default_cli_version
+                else None
+            )
+            current_version = version_stats.client.current.version
+
+            logger.debug(
+                f"""
+                client_version: {current_version}
+                server_default_version: {server_default_version}
+                server_default_cli_ahead: True
+            """
+            )
+
+            self.source.get_report().warning(
+                title="Server default CLI version is ahead of CLI version",
+                message="Please upgrade the CLI version being used",
+                context=f"Server Default CLI version: {server_default_version}, Used CLI version: {current_version}",
+            )
 
     def _notify_reporters_on_ingestion_completion(self) -> None:
         for reporter in self.reporters:
@@ -365,8 +412,8 @@ class Pipeline:
                     report=self._get_structured_report(),
                     ctx=self.ctx,
                 )
-            except Exception as e:
-                logger.warning("Reporting failed on completion", exc_info=e)
+            except Exception:
+                logger.warning("Reporting failed on completion", exc_info=True)
 
     @classmethod
     def create(
@@ -400,7 +447,20 @@ class Pipeline:
             return True
         return False
 
+    def _set_platform(self) -> None:
+        platform = self.source.infer_platform()
+        if platform:
+            self.source.get_report().set_platform(platform)
+        else:
+            self.source.get_report().warning(
+                message="Platform not found",
+                title="Platform not found",
+                context="Platform not found",
+            )
+
     def run(self) -> None:
+        self._set_platform()
+        self._warn_old_cli_version()
         with self.exit_stack, self.inner_exit_stack:
             if self.config.flags.generate_memory_profiles:
                 import memray
@@ -466,10 +526,10 @@ class Pipeline:
 
                     except (RuntimeError, SystemExit):
                         raise
-                    except Exception as e:
+                    except Exception:
                         logger.error(
                             "Failed to process some records. Continuing.",
-                            exc_info=e,
+                            exc_info=True,
                         )
                         # TODO: Transformer errors should be reported more loudly / as part of the pipeline report.
 
@@ -498,9 +558,9 @@ class Pipeline:
 
                 self.process_commits()
                 self.final_status = PipelineStatus.COMPLETED
-            except (SystemExit, KeyboardInterrupt) as e:
+            except (SystemExit, KeyboardInterrupt):
                 self.final_status = PipelineStatus.CANCELLED
-                logger.error("Caught error", exc_info=e)
+                logger.error("Caught error", exc_info=True)
                 raise
             except Exception as exc:
                 self.final_status = PipelineStatus.ERROR
@@ -508,8 +568,6 @@ class Pipeline:
             finally:
                 clear_global_warnings()
 
-                self._notify_reporters_on_ingestion_completion()
-
     def transform(self, records: Iterable[RecordEnvelope]) -> Iterable[RecordEnvelope]:
         """
         Transforms the given sequence of records by passing the records through the transformers
@@ -561,18 +619,20 @@ class Pipeline:
     def raise_from_status(self, raise_warnings: bool = False) -> None:
         if self.source.get_report().failures:
             raise PipelineExecutionError(
-                "Source reported errors", self.source.get_report()
+                "Source reported errors", self.source.get_report().failures
             )
         if self.sink.get_report().failures:
-            raise PipelineExecutionError("Sink reported errors", self.sink.get_report())
+            raise PipelineExecutionError(
+                "Sink reported errors", self.sink.get_report().failures
+            )
         if raise_warnings:
             if self.source.get_report().warnings:
                 raise PipelineExecutionError(
-                    "Source reported warnings", self.source.get_report()
+                    "Source reported warnings", self.source.get_report().warnings
                 )
             if self.sink.get_report().warnings:
                 raise PipelineExecutionError(
-                    "Sink reported warnings", self.sink.get_report()
+                    "Sink reported warnings", self.sink.get_report().warnings
                 )
 
     def log_ingestion_stats(self) -> None:
@@ -581,15 +641,22 @@ class Pipeline:
         sink_failures = len(self.sink.get_report().failures)
         sink_warnings = len(self.sink.get_report().warnings)
         global_warnings = len(get_global_warnings())
+        source_aspects = self.source.get_report().get_aspects_dict()
+        source_aspects_by_subtype = (
+            self.source.get_report().get_aspects_by_subtypes_dict()
+        )
 
         telemetry_instance.ping(
             "ingest_stats",
             {
                 "source_type": self.source_type,
+                "source_aspects": source_aspects,
+                "source_aspects_by_subtype": source_aspects_by_subtype,
                 "sink_type": self.sink_type,
                 "transformer_types": [
                     transformer.type for transformer in self.config.transformers or []
                 ],
+                "extractor_type": self.config.source.extractor,
                 "records_written": stats.discretize(
                     self.sink.get_report().total_records_written
                 ),

datahub/ingestion/run/pipeline_config.py

@@ -6,8 +6,8 @@ from typing import Any, Dict, List, Optional
 
 from pydantic import Field, validator
 
-from datahub.configuration.common import ConfigModel, DynamicTypedConfig
-from datahub.ingestion.graph.client import DatahubClientConfig
+from datahub.configuration.common import ConfigModel, DynamicTypedConfig, HiddenFromDocs
+from datahub.ingestion.graph.config import DatahubClientConfig
 from datahub.ingestion.sink.file import FileSinkConfig
 
 logger = logging.getLogger(__name__)
@@ -85,7 +85,7 @@ class PipelineConfig(ConfigModel):
     source: SourceConfig
     sink: Optional[DynamicTypedConfig] = None
     transformers: Optional[List[DynamicTypedConfig]] = None
-    flags: FlagsConfig = Field(default=FlagsConfig(), hidden_from_docs=True)
+    flags: HiddenFromDocs[FlagsConfig] = FlagsConfig()
     reporting: List[ReporterConfig] = []
     run_id: str = DEFAULT_RUN_ID
     datahub_api: Optional[DatahubClientConfig] = None

datahub/ingestion/sink/datahub_kafka.py

@@ -74,4 +74,5 @@ class DatahubKafkaSink(Sink[KafkaSinkConfig, SinkReport]):
             callback(err, f"Failed to write record: {err}")
 
     def close(self) -> None:
+        super().close()
         self.emitter.flush()