Merge branch 'main' into notion-date-validator

MetaphorData · Sep 25, 2024 · 438e60c · 438e60c
2 parents d35f0da + f5c453b
commit 438e60c
Show file tree

Hide file tree

Showing 7 changed files with 1,868 additions and 455 deletions.
diff --git a/metaphor/looker/lookml_parser.py b/metaphor/looker/lookml_parser.py
@@ -29,12 +29,12 @@
     LookerView,
     LookerViewDimension,
     LookerViewMeasure,
-    LookerViewQuery,
     SystemTag,
     SystemTags,
     SystemTagSource,
     VirtualView,
     VirtualViewLogicalID,
+    VirtualViewQuery,
     VirtualViewType,
 )
 
@@ -132,7 +132,7 @@ def _get_upstream_and_query(
     view_name: str,
     raw_model: RawModel,
     connection: LookerConnectionConfig,
-) -> Tuple[Set[EntityId], Optional[LookerViewQuery]]:
+) -> Tuple[Set[EntityId], Optional[VirtualViewQuery]]:
     raw_views = raw_model.raw_views
     raw_view = raw_views.get(view_name)
     if raw_view is None:
@@ -150,7 +150,7 @@ def _get_upstream_and_query(
     derived_table: Optional[Dict] = raw_view.get("derived_table")
     if derived_table is not None:
         if "sql" in derived_table:
-            query = LookerViewQuery(
+            query = VirtualViewQuery(
                 query=derived_table["sql"],
                 source_platform=connection.platform,
                 source_dataset_account=connection.account,

diff --git a/metaphor/quick_sight/extractor.py b/metaphor/quick_sight/extractor.py
@@ -29,11 +29,7 @@
     DATA_SET_DIRECTORIES,
     create_top_level_folders,
 )
-from metaphor.quick_sight.lineage import (
-    extract_virtual_view_schema,
-    extract_virtual_view_upstream,
-    process_dataset_lineage,
-)
+from metaphor.quick_sight.lineage import LineageProcessor
 from metaphor.quick_sight.models import Dashboard, DataSet, ResourceType
 
 logger = get_logger()
@@ -80,15 +76,16 @@ def _extract_virtual_views(self):
                 continue
 
             view = self._init_virtual_view(data_set.Arn, data_set)
-
-            columns, source_entities = process_dataset_lineage(
-                self._resources, data_set
-            )
-
-            view.schema = extract_virtual_view_schema(data_set, columns)
-            view.entity_upstream = extract_virtual_view_upstream(
-                columns, source_entities
-            )
+            output_logical_table_id = LineageProcessor(
+                self._resources, self._virtual_views, data_set
+            ).run()
+
+            # The last logical_table is the output table for the dataset
+            output = self._virtual_views.get(output_logical_table_id)
+            if output:
+                view.schema = output.schema
+                view.entity_upstream = output.entity_upstream
+                self._virtual_views.pop(output_logical_table_id)
 
     def _extract_dashboards(self) -> None:
         for dashboard in self._resources.values():