fix(ingest): fix serialization of report to handle nesting (datahub-p…

…roject#5455)

metadata-ingestion/src/datahub/ingestion/api/report.py

@@ -2,7 +2,8 @@
 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 @@
 
 @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:

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)