fix(powerbi): fix access token expiry

metadata-ingestion/src/datahub/ingestion/source/powerbi/config.py

@@ -95,6 +95,7 @@ class Constant:
     TITLE = "title"
     EMBED_URL = "embedUrl"
     ACCESS_TOKEN = "access_token"
+    ACCESS_TOKEN_EXPIRY = "expires_in"
     IS_READ_ONLY = "isReadOnly"
     WEB_URL = "webUrl"
     ODATA_COUNT = "@odata.count"

metadata-ingestion/src/datahub/ingestion/source/powerbi/rest_api_wrapper/data_resolver.py

@@ -1,6 +1,7 @@
 import logging
 import math
 from abc import ABC, abstractmethod
+from datetime import datetime, timedelta
 from time import sleep
 from typing import Any, Dict, List, Optional
 
@@ -59,6 +60,7 @@ def __init__(
         tenant_id: str,
     ):
         self.__access_token: Optional[str] = None
+        self.__access_token_expiry_time: Optional[datetime] = None
         self.__tenant_id = tenant_id
         # Test connection by generating access token
         logger.info("Trying to connect to {}".format(self._get_authority_url()))
@@ -128,7 +130,7 @@ def get_authorization_header(self):
         return {Constant.Authorization: self.get_access_token()}
 
     def get_access_token(self):
-        if self.__access_token is not None:
+        if self.__access_token is not None and not self._is_access_token_expired():
             return self.__access_token
 
         logger.info("Generating PowerBi access token")
@@ -150,11 +152,20 @@ def get_access_token(self):
         self.__access_token = "Bearer {}".format(
             auth_response.get(Constant.ACCESS_TOKEN)
         )
+        safety_gap = 300
+        self.__access_token_expiry_time = datetime.now() + timedelta(
+            seconds=(
+                max(auth_response.get(Constant.ACCESS_TOKEN_EXPIRY, 0) - safety_gap, 0)
+            )
+        )
 
         logger.debug(f"{Constant.PBIAccessToken}={self.__access_token}")
 
         return self.__access_token
 
+    def _is_access_token_expired(self) -> bool:
+        return self.__access_token_expiry_time < datetime.now()
+
     def get_dashboards(self, workspace: Workspace) -> List[Dashboard]:
         """
         Get the list of dashboard from PowerBi for the given workspace identifier

metadata-ingestion/tests/integration/powerbi/test_powerbi.py

@@ -2,6 +2,7 @@
 import sys
 from typing import Any, Dict, List, cast
 from unittest import mock
+from unittest.mock import MagicMock
 
 import pytest
 from freezegun import freeze_time
@@ -1061,6 +1062,50 @@ def test_workspace_container(
     )
 
 
+@mock.patch("msal.ConfidentialClientApplication", side_effect=mock_msal_cca)
+def test_access_token_expiry(
+    mock_msal: MagicMock, pytestconfig, tmp_path, mock_time, requests_mock
+):
+    enable_logging()
+
+    register_mock_api(request_mock=requests_mock)
+
+    pipeline = Pipeline.create(
+        {
+            "run_id": "powerbi-test",
+            "source": {
+                "type": "powerbi",
+                "config": {
+                    **default_source_config(),
+                },
+            },
+            "sink": {
+                "type": "file",
+                "config": {
+                    "filename": f"{tmp_path}/powerbi_access_token_mces.json",
+                },
+            },
+        }
+    )
+
+    # for long expiry, the token should only be requested once.
+    mock_msal.acquire_token_for_client = lambda *args, **kwargs: {
+        "access_token": "dummy",
+        "expires_in": 3600,
+    }
+    pipeline.run()
+    mock_msal.return_value.acquire_token_for_client.assert_called_once()
+
+    # for short expiry, the token should be requested when expires.
+    mock_msal.reset_mock()
+    mock_msal.acquire_token_for_client = lambda *args, **kwargs: {
+        "access_token": "dummy",
+        "expires_in": 0,
+    }
+    pipeline.run()
+    assert len(mock_msal.return_value.acquire_token_for_client.mock_calls) > 1
+
+
 def dataset_type_mapping_set_to_all_platform(pipeline: Pipeline) -> None:
     source_config: PowerBiDashboardSourceConfig = cast(
         PowerBiDashboardSource, pipeline.source