fix(ingest/sigma): handle members api paginated response (#11535)

datahub-project · Oct 4, 2024 · 364496c · 364496c
1 parent 81151db
commit 364496c
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 25 deletions.
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py b/metadata-ingestion/src/datahub/ingestion/source/sigma/sigma_api.py
@@ -141,12 +141,19 @@ def _get_users(self) -> Dict[str, str]:
         logger.debug("Fetching all accessible users metadata.")
         try:
             users: Dict[str, str] = {}
-            response = self._get_api_call(f"{self.config.api_url}/members")
-            response.raise_for_status()
-            for user_dict in response.json():
-                users[
-                    user_dict[Constant.MEMBERID]
-                ] = f"{user_dict[Constant.FIRSTNAME]}_{user_dict[Constant.LASTNAME]}"
+            members_url = url = f"{self.config.api_url}/members?limit=50"
+            while True:
+                response = self._get_api_call(url)
+                response.raise_for_status()
+                response_dict = response.json()
+                for user_dict in response_dict[Constant.ENTRIES]:
+                    users[
+                        user_dict[Constant.MEMBERID]
+                    ] = f"{user_dict[Constant.FIRSTNAME]}_{user_dict[Constant.LASTNAME]}"
+                if response_dict[Constant.NEXTPAGE]:
+                    url = f"{members_url}&page={response_dict[Constant.NEXTPAGE]}"
+                else:
+                    break
             return users
         except Exception as e:
             self._log_http_error(

diff --git a/metadata-ingestion/tests/integration/sigma/test_sigma.py b/metadata-ingestion/tests/integration/sigma/test_sigma.py
@@ -381,25 +381,29 @@ def register_mock_api(request_mock: Any, override_data: dict = {}) -> None:
         "https://aws-api.sigmacomputing.com/v2/members": {
             "method": "GET",
             "status_code": 200,
-            "json": [
-                {
-                    "organizationId": "b94da709-176c-4242-bea6-6760f34c9228",
-                    "memberId": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
-                    "memberType": "admin",
-                    "firstName": "Shubham",
-                    "lastName": "Jagtap",
-                    "email": "[email protected]",
-                    "profileImgUrl": None,
-                    "createdBy": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
-                    "updatedBy": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
-                    "createdAt": "2023-11-28T10:59:20.957Z",
-                    "updatedAt": "2024-03-12T21:21:17.996Z",
-                    "homeFolderId": "9bb94df1-e8af-49eb-9c37-2bd40b0efb2e",
-                    "userKind": "internal",
-                    "isArchived": False,
-                    "isInactive": False,
-                },
-            ],
+            "json": {
+                "entries": [
+                    {
+                        "organizationId": "b94da709-176c-4242-bea6-6760f34c9228",
+                        "memberId": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
+                        "memberType": "admin",
+                        "firstName": "Shubham",
+                        "lastName": "Jagtap",
+                        "email": "[email protected]",
+                        "profileImgUrl": None,
+                        "createdBy": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
+                        "updatedBy": "CPbEdA26GNQ2cM2Ra2BeO0fa5Awz1",
+                        "createdAt": "2023-11-28T10:59:20.957Z",
+                        "updatedAt": "2024-03-12T21:21:17.996Z",
+                        "homeFolderId": "9bb94df1-e8af-49eb-9c37-2bd40b0efb2e",
+                        "userKind": "internal",
+                        "isArchived": False,
+                        "isInactive": False,
+                    },
+                ],
+                "total": 1,
+                "nextPage": None,
+            },
         },
     }