feat(datahub-lite): simplify get response (#7131)

datahub-project · Jan 25, 2023 · e842164 · e842164
1 parent 732c837
commit e842164
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 24 deletions.
diff --git a/metadata-ingestion/src/datahub/lite/duckdb_lite.py b/metadata-ingestion/src/datahub/lite/duckdb_lite.py
@@ -238,9 +238,7 @@ def get(
         typed: bool = False,
         as_of: Optional[int] = None,
         details: Optional[bool] = False,
-    ) -> Optional[
-        Dict[str, Union[str, Dict[str, Union[dict, _Aspect, SystemMetadataClass]]]]
-    ]:
+    ) -> Optional[Dict[str, Union[str, dict, _Aspect]]]:
         base_query = "SELECT urn, aspect_name, metadata, system_metadata from metadata_aspect_v2 WHERE urn = ?"
         if aspects:
             base_query += (
@@ -253,23 +251,21 @@ def get(
 
         self.duckdb_client.execute(base_query, [id])
         results = self.duckdb_client.fetchall()
-        result_map: Dict[
-            str, Union[str, Dict[str, Union[dict, _Aspect, SystemMetadataClass]]]
-        ] = {}
+        result_map: Dict[str, Union[str, dict, _Aspect]] = {}
         for r in results:
             aspect_name: str = r[1]
             aspect: Union[dict, _Aspect] = json.loads(r[2])
             if typed:
                 assert isinstance(aspect, dict)
                 aspect = ASPECT_MAP[aspect_name].from_obj(post_json_transform(aspect))
 
-            result_map[aspect_name] = {"value": aspect}
+            result_map[aspect_name] = aspect
             if details:
                 system_metadata: Union[dict, SystemMetadataClass] = json.loads(r[3])
                 if typed:
                     assert isinstance(system_metadata, dict)
                     system_metadata = SystemMetadataClass.from_obj(system_metadata)
-                result_map[aspect_name].update({"systemMetadata": system_metadata})  # type: ignore
+                result_map[aspect_name].update({"__systemMetadata": system_metadata})  # type: ignore
         if result_map:
             result_map = {**{"urn": id}, **result_map}
             return result_map

diff --git a/metadata-ingestion/src/datahub/lite/lite_local.py b/metadata-ingestion/src/datahub/lite/lite_local.py
@@ -5,11 +5,7 @@
 from datahub.configuration.common import ConfigModel
 from datahub.emitter.mcp import MetadataChangeProposalWrapper
 from datahub.ingestion.api.closeable import Closeable
-from datahub.metadata.schema_classes import (
-    MetadataChangeEventClass,
-    SystemMetadataClass,
-    _Aspect,
-)
+from datahub.metadata.schema_classes import MetadataChangeEventClass, _Aspect
 from datahub.utilities.type_annotations import get_class_from_annotation
 
 LiteConfig = TypeVar("LiteConfig", bound=ConfigModel)
@@ -86,9 +82,7 @@ def get(
         typed: bool = False,
         as_of: Optional[int] = None,
         details: Optional[bool] = False,
-    ) -> Optional[
-        Dict[str, Union[str, Dict[str, Union[dict, _Aspect, SystemMetadataClass]]]]
-    ]:
+    ) -> Optional[Dict[str, Union[str, dict, _Aspect]]]:
         pass
 
     @abstractmethod

diff --git a/metadata-ingestion/src/datahub/lite/lite_util.py b/metadata-ingestion/src/datahub/lite/lite_util.py
@@ -14,11 +14,7 @@
     SearchFlavor,
 )
 from datahub.lite.lite_registry import lite_registry
-from datahub.metadata.schema_classes import (
-    MetadataChangeEventClass,
-    SystemMetadataClass,
-    _Aspect,
-)
+from datahub.metadata.schema_classes import MetadataChangeEventClass, _Aspect
 
 logger = logging.getLogger(__name__)
 
@@ -67,9 +63,7 @@ def get(
         typed: bool = False,
         as_of: Optional[int] = None,
         details: Optional[bool] = False,
-    ) -> Optional[
-        Dict[str, Union[str, Dict[str, Union[dict, _Aspect, SystemMetadataClass]]]]
-    ]:
+    ) -> Optional[Dict[str, Union[str, dict, _Aspect]]]:
         return self.get(id, aspects, typed, as_of, details)
 
     def search(