Refactor create_hierarchy

metaphor/common/hierarchy.py

@@ -0,0 +1,33 @@
+from typing import List
+
+from metaphor.models.metadata_change_event import (
+    AssetPlatform,
+    Hierarchy,
+    HierarchyInfo,
+    HierarchyLogicalID,
+    HierarchyType,
+)
+
+
+def create_hierarchy(
+    platform: AssetPlatform,
+    path: List[str],
+    name: str = "",
+    hierarchy_type: HierarchyType = HierarchyType.VIRTUAL_HIERARCHY,
+) -> Hierarchy:
+    """
+    Create a hierarchy with name
+    """
+    return Hierarchy(
+        logical_id=HierarchyLogicalID(
+            path=[platform.value] + path,
+        ),
+        hierarchy_info=(
+            HierarchyInfo(
+                name=name,
+                type=hierarchy_type,
+            )
+            if name
+            else None
+        ),
+    )

metaphor/looker/folder.py

@@ -1,12 +1,11 @@
 from dataclasses import dataclass
 from typing import Dict, List, Optional
 
+from metaphor.common.hierarchy import create_hierarchy
 from metaphor.common.logger import get_logger
 from metaphor.models.metadata_change_event import (
-    DashboardPlatform,
+    AssetPlatform,
     Hierarchy,
-    HierarchyInfo,
-    HierarchyLogicalID,
     HierarchyType,
 )
 
@@ -58,13 +57,9 @@ def _build_hierarchies(
         if folder_id in folder_hierarchies or folder is None:
             continue
 
-        hierarchy = Hierarchy(
-            logical_id=HierarchyLogicalID(
-                path=[DashboardPlatform.LOOKER.value] + directories[: i + 1]
-            ),
-            hierarchy_info=HierarchyInfo(
-                type=HierarchyType.LOOKER_FOLDER, name=folder.name
-            ),
+        folder_hierarchies[folder_id] = create_hierarchy(
+            platform=AssetPlatform.LOOKER,
+            name=folder.name,
+            path=directories[: i + 1],
+            hierarchy_type=HierarchyType.LOOKER_FOLDER,
         )
-
-        folder_hierarchies[folder_id] = hierarchy

metaphor/quick_sight/folder.py

@@ -1,26 +1,16 @@
 from typing import List
 
-from metaphor.models.metadata_change_event import (
-    DashboardPlatform,
-    Hierarchy,
-    HierarchyInfo,
-    HierarchyLogicalID,
-    HierarchyType,
-)
+from metaphor.common.hierarchy import create_hierarchy
+from metaphor.models.metadata_change_event import AssetPlatform, Hierarchy
 
 DASHBOARD_DIRECTORIES = ["DASHBOARD"]
 DATA_SET_DIRECTORIES = ["DATA_SET"]
 
 
-def _create_virtual_hierarchy(name: str, path: List[str]) -> Hierarchy:
-    return Hierarchy(
-        logical_id=HierarchyLogicalID(path=[DashboardPlatform.QUICK_SIGHT.name] + path),
-        hierarchy_info=HierarchyInfo(name=name, type=HierarchyType.VIRTUAL_HIERARCHY),
-    )
-
-
 def create_top_level_folders() -> List[Hierarchy]:
+    platform = AssetPlatform.QUICK_SIGHT
+
     return [
-        _create_virtual_hierarchy("Dashboards", DASHBOARD_DIRECTORIES),
-        _create_virtual_hierarchy("DataSets", DATA_SET_DIRECTORIES),
+        create_hierarchy(platform, DASHBOARD_DIRECTORIES, "Dashboards"),
+        create_hierarchy(platform, DATA_SET_DIRECTORIES, "DataSets"),
     ]

metaphor/thought_spot/extractor.py

@@ -16,10 +16,12 @@
     to_virtual_view_entity_id,
 )
 from metaphor.common.event_util import ENTITY_TYPES
+from metaphor.common.hierarchy import create_hierarchy
 from metaphor.common.logger import get_logger
 from metaphor.common.utils import unique_list
 from metaphor.models.crawler_run_metadata import Platform
 from metaphor.models.metadata_change_event import (
+    AssetPlatform,
     AssetStructure,
     Chart,
     Dashboard,
@@ -30,6 +32,7 @@
     EntityType,
     EntityUpstream,
     FieldMapping,
+    HierarchyType,
     SourceField,
     SourceInfo,
     SystemTag,
@@ -59,7 +62,6 @@
 )
 from metaphor.thought_spot.utils import (
     ThoughtSpot,
-    create_virtual_hierarchy,
     from_list,
     getColumnTransformation,
     mapping_chart_type,
@@ -102,24 +104,7 @@ async def extract(self) -> Collection[ENTITY_TYPES]:
 
         self.fetch_virtual_views()
         self.fetch_dashboards()
-
-        virtual_hierarchies = [
-            create_virtual_hierarchy(
-                name="Answer", path=[ThoughtSpotDashboardType.ANSWER.name]
-            ),
-            create_virtual_hierarchy(
-                name="Liveboard", path=[ThoughtSpotDashboardType.LIVEBOARD.name]
-            ),
-            create_virtual_hierarchy(
-                name="Table", path=[ThoughtSpotDataObjectType.TABLE.name]
-            ),
-            create_virtual_hierarchy(
-                name="View", path=[ThoughtSpotDataObjectType.VIEW.name]
-            ),
-            create_virtual_hierarchy(
-                name="Worksheet", path=[ThoughtSpotDataObjectType.WORKSHEET.name]
-            ),
-        ]
+        virtual_hierarchies = self._create_virtual_hierarchies()
 
         return list(
             chain(
@@ -246,6 +231,24 @@ def populate_virtual_views(
             )
             self._virtual_views[table_id] = view
 
+    @staticmethod
+    def _create_virtual_hierarchies():
+        return [
+            create_hierarchy(
+                name=name,
+                path=[enum_value.value],
+                platform=AssetPlatform.THOUGHT_SPOT,
+                hierarchy_type=HierarchyType.THOUGHT_SPOT_VIRTUAL_HIERARCHY,
+            )
+            for name, enum_value in [
+                ("Answer", ThoughtSpotDashboardType.ANSWER),
+                ("Liveboard", ThoughtSpotDashboardType.LIVEBOARD),
+                ("Table", ThoughtSpotDataObjectType.TABLE),
+                ("View", ThoughtSpotDataObjectType.VIEW),
+                ("Worksheet", ThoughtSpotDataObjectType.WORKSHEET),
+            ]
+        ]
+
     @staticmethod
     def build_column_expr_map(tml: TMLObject):
         def build_formula_map(tml_table):

metaphor/thought_spot/utils.py

@@ -8,12 +8,7 @@
 from metaphor.common.utils import chunks
 from metaphor.models.metadata_change_event import (
     ChartType,
-    DashboardPlatform,
     DataPlatform,
-    Hierarchy,
-    HierarchyInfo,
-    HierarchyLogicalID,
-    HierarchyType,
     ThoughtSpotDataObjectType,
 )
 from metaphor.thought_spot.config import ThoughtSpotRunConfig
@@ -278,14 +273,3 @@ def getColumnTransformation(target_column: Column) -> Optional[str]:
     if target_column.expression is None or target_column.expression.token is None:
         return None
     return str(target_column.expression.token)
-
-
-def create_virtual_hierarchy(name: str, path: List[str]) -> Hierarchy:
-    return Hierarchy(
-        logical_id=HierarchyLogicalID(
-            path=[DashboardPlatform.THOUGHT_SPOT.name] + path
-        ),
-        hierarchy_info=HierarchyInfo(
-            name=name, type=HierarchyType.THOUGHT_SPOT_VIRTUAL_HIERARCHY
-        ),
-    )