fix(ingest/looker): update for Looker query API breaking change (data…

…hub-project#9865)
FundingCircle · Feb 16, 2024 · cda34b5 · cda34b5
1 parent ae1806f
commit cda34b5
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 10 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_lib_wrapper.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_lib_wrapper.py
@@ -17,6 +17,7 @@
     Look,
     LookmlModel,
     LookmlModelExplore,
+    LookWithQuery,
     Query,
     User,
     WriteQuery,
@@ -64,6 +65,7 @@ class LookerAPIStats(BaseModel):
     all_looks_calls: int = 0
     all_models_calls: int = 0
     get_query_calls: int = 0
+    get_look_calls: int = 0
     search_looks_calls: int = 0
     search_dashboards_calls: int = 0
 
@@ -255,6 +257,14 @@ def get_query(self, query_id: str, fields: Union[str, List[str]]) -> Query:
             transport_options=self.transport_options,
         )
 
+    def get_look(self, look_id: str, fields: Union[str, List[str]]) -> LookWithQuery:
+        self.client_stats.get_look_calls += 1
+        return self.client.look(
+            look_id=look_id,
+            fields=self.__fields_mapper(fields),
+            transport_options=self.transport_options,
+        )
+
     def search_dashboards(
         self, fields: Union[str, List[str]], deleted: str
     ) -> Sequence[Dashboard]:

diff --git a/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py b/metadata-ingestion/src/datahub/ingestion/source/looker/looker_source.py
@@ -1199,7 +1199,18 @@ def extract_independent_looks(self) -> Iterable[MetadataWorkUnit]:
                 logger.info(f"query_id is None for look {look.title}({look.id})")
                 continue
 
-            query: Query = self.looker_api.get_query(look.query_id, query_fields)
+            if look.id is not None:
+                query: Optional[Query] = self.looker_api.get_look(
+                    look.id, fields=["query"]
+                ).query
+                # Only include fields that are in the query_fields list
+                query = Query(
+                    **{
+                        key: getattr(query, key)
+                        for key in query_fields
+                        if hasattr(query, key)
+                    }
+                )
 
             dashboard_element: Optional[
                 LookerDashboardElement

diff --git a/metadata-ingestion/tests/integration/looker/test_looker.py b/metadata-ingestion/tests/integration/looker/test_looker.py
@@ -311,15 +311,17 @@ def setup_mock_look(mocked_client):
         )
     ]
 
-    mocked_client.query.return_value = Query(
-        id="1",
-        view="sales_explore",
-        model="sales_model",
-        fields=[
-            "sales.profit",
-        ],
-        dynamic_fields=None,
-        filters=None,
+    mocked_client.look.return_value = LookWithQuery(
+        query=Query(
+            id="1",
+            view="sales_explore",
+            model="sales_model",
+            fields=[
+                "sales.profit",
+            ],
+            dynamic_fields=None,
+            filters=None,
+        )
     )