diff --git a/src/preset_cli/api/clients/preset.py b/src/preset_cli/api/clients/preset.py index 41308af..aed6992 100644 --- 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 index c686160..cfad198 100644 --- 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": "bdoe@example.com", }, }, + ], + "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": "bdoe@example.com", }, }, + ], + "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 index 1e024d1..84af9f2 100644 --- 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": "cdoe@example.com", + }, + }, + ], + "meta": { + "count": 251, + }, + }, + ) + requests_mock.get( "https://superset.example.org/roles/add", text=""" """, ) @@ -1932,6 +1956,14 @@ def test_export_users_preset(requests_mock: Mocker) -> None: "email": "bdoe@example.com", "role": [], }, + { + "id": 3, + "first_name": "Clarisse", + "last_name": "Doe", + "username": "cdoe", + "email": "cdoe@example.com", + "role": [], + }, ]