chore: Paginate Manager API calls (#320)

* chore: Paginate Manager API calls * Addressing PR feedbak
preset-io · Oct 10, 2024 · 3e7950a · 3e7950a
1 parent 4deefb7
commit 3e7950a
Show file tree

Hide file tree

Showing 3 changed files with 115 additions and 31 deletions.
diff --git a/src/preset_cli/api/clients/preset.py b/src/preset_cli/api/clients/preset.py
@@ -17,6 +17,8 @@
 
 _logger = logging.getLogger(__name__)
 
+MANAGER_MAX_PAGE_SIZE = 250
+
 
 class Role(int, Enum):
     """
@@ -61,13 +63,22 @@ def get_team_members(self, team_name: str) -> List[Any]:
         """
         Retrieve all users for a given team.
         """
-        url = self.get_base_url() / "teams" / team_name / "memberships"
-        _logger.debug("GET %s", url)
-        response = self.session.get(url)
-        validate_response(response)
+        users = []
+        page_number = 1
+        while True:
+            params = {"page_number": page_number, "page_size": MANAGER_MAX_PAGE_SIZE}
+            url = self.get_base_url() / "teams" / team_name / "memberships" % params
+            _logger.debug("GET %s", url)
+            response = self.session.get(url)
+            validate_response(response)
 
-        payload = response.json()
-        users = payload["payload"]
+            payload = response.json()
+            users.extend(payload["payload"])
+
+            if payload["meta"]["count"] <= page_number * MANAGER_MAX_PAGE_SIZE:
+                break
+
+            page_number += 1
 
         return users
 
@@ -123,27 +134,41 @@ def export_users(self, workspace_url: URL) -> Iterator[UserType]:
         if team_name is None or workspace_id is None:
             raise Exception("Unable to find workspace and/or team")
 
-        url = (
-            self.get_base_url()
-            / "teams"
-            / team_name
-            / "workspaces"
-            / str(workspace_id)
-            / "memberships"
-        )
-        _logger.debug("GET %s", url)
-        response = self.session.get(url)
-        team_members: List[UserType] = [
-            {
-                "id": 0,
-                "username": payload["user"]["username"],
-                "role": [],  # TODO (betodealmeida)
-                "first_name": payload["user"]["first_name"],
-                "last_name": payload["user"]["last_name"],
-                "email": payload["user"]["email"],
-            }
-            for payload in response.json()["payload"]
-        ]
+        workspace_membership: List[UserType] = []
+        page_number = 1
+        while True:
+            params = {"page_number": page_number, "page_size": MANAGER_MAX_PAGE_SIZE}
+            url = (
+                self.get_base_url()
+                / "teams"
+                / team_name
+                / "workspaces"
+                / str(workspace_id)
+                / "memberships"
+                % params
+            )
+            _logger.debug("GET %s", url)
+            response = self.session.get(url)
+            validate_response(response)
+            payload = response.json()
+
+            team_members: List[UserType] = [
+                {
+                    "id": 0,
+                    "username": payload["user"]["username"],
+                    "role": [],  # TODO (betodealmeida)
+                    "first_name": payload["user"]["first_name"],
+                    "last_name": payload["user"]["last_name"],
+                    "email": payload["user"]["email"],
+                }
+                for payload in payload["payload"]
+            ]
+            workspace_membership.extend(team_members)
+
+            if payload["meta"]["count"] <= page_number * MANAGER_MAX_PAGE_SIZE:
+                break
+
+            page_number += 1
 
         # TODO (betodealmeida): improve this
         url = workspace_url / "roles/add"
@@ -156,7 +181,7 @@ def export_users(self, workspace_url: URL) -> Iterator[UserType]:
             for option in select.find_all("option")
         }
 
-        for team_member in team_members:
+        for team_member in workspace_membership:
             # pylint: disable=consider-using-f-string
             full_name = "{first_name} {last_name}".format(**team_member)
             if full_name in ids:

diff --git a/tests/api/clients/preset_test.py b/tests/api/clients/preset_test.py
@@ -91,7 +91,7 @@ def test_preset_client_export_users(requests_mock: Mocker) -> None:
         },
     )
     requests_mock.get(
-        "https://ws.preset.io/v1/teams/team1/workspaces/2/memberships",
+        "https://ws.preset.io/v1/teams/team1/workspaces/2/memberships?page_number=1&page_size=250",
         json={
             "payload": [
                 {
@@ -110,6 +110,16 @@ def test_preset_client_export_users(requests_mock: Mocker) -> None:
                         "email": "[email protected]",
                     },
                 },
+            ],
+            "meta": {
+                "count": 251,
+            },
+        },
+    )
+    requests_mock.get(
+        "https://ws.preset.io/v1/teams/team1/workspaces/2/memberships?page_number=2&page_size=250",
+        json={
+            "payload": [
                 {
                     "user": {
                         "username": "cdoe",
@@ -119,8 +129,12 @@ def test_preset_client_export_users(requests_mock: Mocker) -> None:
                     },
                 },
             ],
+            "meta": {
+                "count": 251,
+            },
         },
     )
+
     requests_mock.get(
         "https://superset.example.org/roles/add",
         text="""
@@ -271,7 +285,7 @@ def test_get_team_members(requests_mock: Mocker) -> None:
     Test the ``get_team_members`` method.
     """
     requests_mock.get(
-        "https://ws.preset.io/v1/teams/botafogo/memberships",
+        "https://ws.preset.io/v1/teams/botafogo/memberships?page_number=1&page_size=250",
         json={
             "payload": [
                 {
@@ -290,6 +304,16 @@ def test_get_team_members(requests_mock: Mocker) -> None:
                         "email": "[email protected]",
                     },
                 },
+            ],
+            "meta": {
+                "count": 251,
+            },
+        },
+    )
+    requests_mock.get(
+        "https://ws.preset.io/v1/teams/botafogo/memberships?page_number=2&page_size=250",
+        json={
+            "payload": [
                 {
                     "user": {
                         "username": "cdoe",
@@ -299,6 +323,9 @@ def test_get_team_members(requests_mock: Mocker) -> None:
                     },
                 },
             ],
+            "meta": {
+                "count": 251,
+            },
         },
     )
 

diff --git a/tests/api/clients/superset_test.py b/tests/api/clients/superset_test.py
@@ -1881,7 +1881,7 @@ def test_export_users_preset(requests_mock: Mocker) -> None:
         },
     )
     requests_mock.get(
-        "https://api.app.preset.io/v1/teams/team1/workspaces/1/memberships",
+        "https://api.app.preset.io/v1/teams/team1/workspaces/1/memberships?page_number=1&page_size=250",
         json={
             "payload": [
                 {
@@ -1901,14 +1901,38 @@ def test_export_users_preset(requests_mock: Mocker) -> None:
                     },
                 },
             ],
+            "meta": {
+                "count": 251,
+            },
         },
     )
+
+    requests_mock.get(
+        "https://api.app.preset.io/v1/teams/team1/workspaces/1/memberships?page_number=2&page_size=250",
+        json={
+            "payload": [
+                {
+                    "user": {
+                        "username": "cdoe",
+                        "first_name": "Clarisse",
+                        "last_name": "Doe",
+                        "email": "[email protected]",
+                    },
+                },
+            ],
+            "meta": {
+                "count": 251,
+            },
+        },
+    )
+
     requests_mock.get(
         "https://superset.example.org/roles/add",
         text="""
 <select id="user">
     <option value="1">Alice Doe</option>
     <option value="2">Bob Doe</option>
+    <option value="3">Clarisse Doe</option>
 </select>
     """,
     )
@@ -1932,6 +1956,14 @@ def test_export_users_preset(requests_mock: Mocker) -> None:
             "email": "[email protected]",
             "role": [],
         },
+        {
+            "id": 3,
+            "first_name": "Clarisse",
+            "last_name": "Doe",
+            "username": "cdoe",
+            "email": "[email protected]",
+            "role": [],
+        },
     ]