Fix rate limiting issues with Power BI admin APIs

metaphor/power_bi/extractor.py

@@ -63,6 +63,9 @@
 from metaphor.power_bi.graph_api_client import GraphApiClient
 from metaphor.power_bi.models import (
     PowerBIApp,
+    PowerBIDashboard,
+    PowerBIDataset,
+    PowerBIReport,
     PowerBISubscription,
     PowerBiSubscriptionUser,
     WorkspaceInfo,
@@ -113,13 +116,17 @@ def __init__(self, config: PowerBIRunConfig):
     async def extract(self) -> Collection[ENTITY_TYPES]:
         logger.info(f"Fetching metadata from Power BI tenant ID: {self._tenant_id}")
 
+        dataset_map = {d.id: d for d in self._client.get_datasets()}
+        dashboard_map = {d.id: d for d in self._client.get_dashboards()}
+        report_map = {r.id: r for r in self._client.get_reports()}
+        app_map = {app.id: app for app in self._client.get_apps()}
+
         if len(self._workspaces) == 0:
             self._workspaces = [w.id for w in self._client.get_groups()]
 
-        logger.info(f"Process {len(self._workspaces)} workspaces: {self._workspaces}")
-
-        apps = self._client.get_apps()
-        app_map = {app.id: app for app in apps}
+        logger.info(
+            f"Processing {len(self._workspaces)} workspaces: {self._workspaces}"
+        )
 
         workspaces: List[WorkspaceInfo] = []
 
@@ -135,11 +142,11 @@ async def extract(self) -> Collection[ENTITY_TYPES]:
         # As there may be cross-workspace reference in dashboards & reports,
         # we must process the datasets across all workspaces first
         for workspace in workspaces:
-            await self.map_wi_datasets_to_virtual_views(workspace)
+            await self.map_wi_datasets_to_virtual_views(workspace, dataset_map)
 
         for workspace in workspaces:
-            await self.map_wi_reports_to_dashboard(workspace, app_map)
-            await self.map_wi_dashboards_to_dashboard(workspace, app_map)
+            await self.map_wi_reports_to_dashboard(workspace, report_map, app_map)
+            await self.map_wi_dashboards_to_dashboard(workspace, dashboard_map, app_map)
 
         self.extract_subscriptions(workspaces)
 
@@ -327,9 +334,9 @@ def _extract_pipeline_info(
                 pipeline_mapping=pipeline_mappings
             )
 
-    async def map_wi_datasets_to_virtual_views(self, workspace: WorkspaceInfo) -> None:
-        dataset_map = {d.id: d for d in self._client.get_datasets(workspace.id)}
-
+    async def map_wi_datasets_to_virtual_views(
+        self, workspace: WorkspaceInfo, dataset_map: Dict[str, PowerBIDataset]
+    ) -> None:
         for wds in workspace.datasets:
             ds = dataset_map.get(wds.id, None)
             if ds is None:
@@ -402,10 +409,11 @@ async def map_wi_datasets_to_virtual_views(self, workspace: WorkspaceInfo) -> No
             self._virtual_views[wds.id] = virtual_view
 
     async def map_wi_reports_to_dashboard(
-        self, workspace: WorkspaceInfo, app_map: Dict[str, PowerBIApp]
+        self,
+        workspace: WorkspaceInfo,
+        report_map: Dict[str, PowerBIReport],
+        app_map: Dict[str, PowerBIApp],
     ) -> None:
-        report_map = {r.id: r for r in self._client.get_reports(workspace.id)}
-
         for wi_report in workspace.reports:
             if wi_report.datasetId is None:
                 logger.warning(f"Skipping report without datasetId: {wi_report.id}")
@@ -465,12 +473,13 @@ async def map_wi_reports_to_dashboard(
             self._dashboards[wi_report.id] = dashboard
 
     async def map_wi_dashboards_to_dashboard(
-        self, workspace: WorkspaceInfo, app_map: Dict[str, PowerBIApp]
+        self,
+        workspace: WorkspaceInfo,
+        dashboard_map: Dict[str, PowerBIDashboard],
+        app_map: Dict[str, PowerBIApp],
     ) -> None:
-        dashboard_map = {d.id: d for d in self._client.get_dashboards(workspace.id)}
-
         for wi_dashboard in workspace.dashboards:
-            tiles = self._client.get_tiles(wi_dashboard.id)
+            tiles = self._client.get_dashboard_tiles(workspace.id, wi_dashboard.id)
             upstream = []
             for tile in tiles:
                 dataset_id = tile.datasetId

metaphor/power_bi/power_bi_client.py

@@ -105,24 +105,9 @@ def get_apps(self) -> List[PowerBIApp]:
             url, List[PowerBIApp], transform_response=lambda r: r.json()["value"]
         )
 
-    def get_tiles(self, dashboard_id: str) -> List[PowerBITile]:
-        # https://docs.microsoft.com/en-us/rest/api/power-bi/admin/dashboards-get-tiles-as-admin
-        url = f"{self.API_ENDPOINT}/admin/dashboards/{dashboard_id}/tiles"
-
-        try:
-            return self._call_get(
-                url, List[PowerBITile], transform_response=lambda r: r.json()["value"]
-            )
-        except EntityNotFoundError:
-            logger.error(
-                f"Unable to find dashboard {dashboard_id}."
-                f"Please add the service principal as a viewer to the workspace"
-            )
-            return []
-
-    def get_datasets(self, group_id: str) -> List[PowerBIDataset]:
-        # https://docs.microsoft.com/en-us/rest/api/power-bi/admin/datasets-get-datasets-in-group-as-admin
-        url = f"{self.API_ENDPOINT}/admin/groups/{group_id}/datasets"
+    def get_datasets(self) -> List[PowerBIDataset]:
+        # https://learn.microsoft.com/en-us/rest/api/power-bi/admin/datasets-get-datasets-as-admin
+        url = f"{self.API_ENDPOINT}/admin/datasets"
         return self._call_get(
             url, List[PowerBIDataset], transform_response=lambda r: r.json()["value"]
         )
@@ -163,16 +148,32 @@ def get_dataset_datasources(
             )
             return []
 
-    def get_dashboards(self, group_id: str) -> List[PowerBIDashboard]:
-        # https://docs.microsoft.com/en-us/rest/api/power-bi/admin/dashboards-get-dashboards-in-group-as-admin
-        url = f"{self.API_ENDPOINT}/admin/groups/{group_id}/dashboards"
+    def get_dashboards(self) -> List[PowerBIDashboard]:
+        # https://learn.microsoft.com/en-us/rest/api/power-bi/admin/dashboards-get-dashboards-as-admin
+        url = f"{self.API_ENDPOINT}/admin/dashboards"
         return self._call_get(
             url, List[PowerBIDashboard], transform_response=lambda r: r.json()["value"]
         )
 
-    def get_reports(self, group_id: str) -> List[PowerBIReport]:
-        # https://docs.microsoft.com/en-us/rest/api/power-bi/admin/reports-get-reports-in-group-as-admin
-        url = f"{self.API_ENDPOINT}/admin/groups/{group_id}/reports"
+    def get_dashboard_tiles(
+        self, group_id: str, dashboard_id: str
+    ) -> List[PowerBITile]:
+        # https://learn.microsoft.com/en-us/rest/api/power-bi/dashboards/get-tiles-in-group
+        url = f"{self.API_ENDPOINT}/groups/{group_id}/dashboards/{dashboard_id}/tiles"
+
+        try:
+            return self._call_get(
+                url, List[PowerBITile], transform_response=lambda r: r.json()["value"]
+            )
+        except Exception as e:
+            logger.error(
+                f"Failed to get tiles from dashboard {dashboard_id} in workspace {group_id}: {e}"
+            )
+            return []
+
+    def get_reports(self) -> List[PowerBIReport]:
+        # https://learn.microsoft.com/en-us/rest/api/power-bi/admin/reports-get-reports-as-admin
+        url = f"{self.API_ENDPOINT}/admin/reports"
         return self._call_get(
             url, List[PowerBIReport], transform_response=lambda r: r.json()["value"]
         )

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "metaphor-connectors"
-version = "0.14.157"
+version = "0.14.158"
 license = "Apache-2.0"
 description = "A collection of Python-based 'connectors' that extract metadata from various sources to ingest into the Metaphor app."
 authors = ["Metaphor <[email protected]>"]

tests/power_bi/test_extractor.py

@@ -177,7 +177,7 @@
 dataflow_id2 = "00000000-0000-0000-0002-00000000000A"
 
 
-def fake_get_datasets(workspace_id: str) -> List[PowerBIDataset]:
+def fake_get_datasets() -> List[PowerBIDataset]:
     return [dataset1, dataset2, dataset3]
 
 
@@ -194,15 +194,15 @@ def fake_get_dataset_parameters(
     ]
 
 
-def fake_get_reports(workspace_id: str) -> List[PowerBIReport]:
+def fake_get_reports() -> List[PowerBIReport]:
     return [report1, report2, report1_app]
 
 
-def fake_get_dashboards(workspace_id: str) -> List[PowerBIDashboard]:
+def fake_get_dashboards() -> List[PowerBIDashboard]:
     return [dashboard1, dashboard2, dashboard1_app]
 
 
-def fake_get_tiles(dashboard_id: str) -> List[PowerBITile]:
+def fake_get_dashboard_tiles(workspace_id: str, dashboard_id: str) -> List[PowerBITile]:
     return tiles[dashboard_id]
 
 
@@ -555,7 +555,9 @@ def fake_export_dataflow(workspace_id: str, df_id: str) -> dict:
     )
     mocked_pbi_client_instance.get_reports.side_effect = fake_get_reports
     mocked_pbi_client_instance.get_dashboards.side_effect = fake_get_dashboards
-    mocked_pbi_client_instance.get_tiles.side_effect = fake_get_tiles
+    mocked_pbi_client_instance.get_dashboard_tiles.side_effect = (
+        fake_get_dashboard_tiles
+    )
     mocked_pbi_client_instance.get_pages.side_effect = fake_get_pages
     mocked_pbi_client_instance.get_refreshes.side_effect = fake_get_refreshes
     mocked_pbi_client_instance.get_apps.side_effect = fake_get_apps