diff --git a/metadata-ingestion/src/datahub/ingestion/api/report.py b/metadata-ingestion/src/datahub/ingestion/api/report.py index c8c4bb46655..d2403301d85 100644 --- a/metadata-ingestion/src/datahub/ingestion/api/report.py +++ b/metadata-ingestion/src/datahub/ingestion/api/report.py @@ -2,7 +2,8 @@ import json import pprint import sys from dataclasses import dataclass -from typing import Dict +from enum import Enum +from typing import Any, Dict # The sort_dicts option was added in Python 3.8. if sys.version_info >= (3, 8): @@ -13,17 +14,36 @@ else: @dataclass class Report: + @staticmethod + def to_str(some_val: Any) -> str: + if isinstance(some_val, Enum): + return some_val.name + else: + return str(some_val) + + @staticmethod + def to_dict(some_val: Any) -> Any: + """A cheap way to generate a dictionary.""" + if hasattr(some_val, "as_obj"): + return some_val.as_obj() + if hasattr(some_val, "dict"): + return some_val.dict() + elif isinstance(some_val, list): + return [Report.to_dict(v) for v in some_val if v is not None] + elif isinstance(some_val, dict): + return { + Report.to_str(k): Report.to_dict(v) + for k, v in some_val.items() + if v is not None + } + else: + return Report.to_str(some_val) + def as_obj(self) -> dict: return { - key: value.as_obj() - if hasattr(value, "as_obj") - else value.dict() - if hasattr(value, "dict") # BaseModel extensions - else value - if isinstance(value, list) or isinstance(value, dict) # simple collections - else str(value) # stringify everything else + str(key): Report.to_dict(value) for (key, value) in self.__dict__.items() - if value # ignore nulls + if value is not None # ignore nulls } def as_string(self) -> str: diff --git a/metadata-ingestion/tests/unit/test_capability_report.py b/metadata-ingestion/tests/unit/test_capability_report.py new file mode 100644 index 00000000000..08ada0386b0 --- /dev/null +++ b/metadata-ingestion/tests/unit/test_capability_report.py @@ -0,0 +1,50 @@ +import json +from typing import cast + +from datahub.ingestion.api.source import ( + CapabilityReport, + SourceCapability, + TestConnectionReport, +) + + +def test_basic_capability_report(): + report = TestConnectionReport( + basic_connectivity=CapabilityReport( + capable=True, failure_reason=None, mitigation_message=None + ), + capability_report={ + "CONTAINERS": CapabilityReport( + capable=True, failure_reason=None, mitigation_message=None + ), + "SCHEMA_METADATA": CapabilityReport( + capable=True, failure_reason=None, mitigation_message=None + ), + "DESCRIPTIONS": CapabilityReport( + capable=False, + failure_reason="failed to get descriptions", + mitigation_message="Enable admin privileges for this account.", + ), + "DATA_PROFILING": CapabilityReport( + capable=True, failure_reason=None, mitigation_message=None + ), + SourceCapability.DOMAINS: CapabilityReport(capable=True), + }, + ) + print(report.as_obj()) + foo = cast(dict, report.as_obj()) + assert isinstance(foo, dict) + assert foo["capability_report"]["CONTAINERS"]["capable"] is True + assert foo["capability_report"]["SCHEMA_METADATA"]["capable"] is True + assert foo["capability_report"]["DESCRIPTIONS"]["capable"] is False + assert ( + foo["capability_report"]["DESCRIPTIONS"]["failure_reason"] + == "failed to get descriptions" + ) + assert ( + foo["capability_report"]["DESCRIPTIONS"]["mitigation_message"] + == "Enable admin privileges for this account." + ) + assert foo["capability_report"]["DOMAINS"]["capable"] is True + + assert isinstance(json.loads(report.as_json()), dict)