diff --git a/src/preset_cli/api/clients/preset.py b/src/preset_cli/api/clients/preset.py index aed6992..ff4a899 100644 --- a/src/preset_cli/api/clients/preset.py +++ b/src/preset_cli/api/clients/preset.py @@ -7,7 +7,7 @@ from enum import Enum from typing import Any, Dict, Iterator, List, Optional, Union -from bs4 import BeautifulSoup +import prison from yarl import URL from preset_cli import __version__ @@ -18,6 +18,7 @@ _logger = logging.getLogger(__name__) MANAGER_MAX_PAGE_SIZE = 250 +SUPERSET_MAX_PAGE_SIZE = 100 class Role(int, Enum): @@ -170,22 +171,34 @@ def export_users(self, workspace_url: URL) -> Iterator[UserType]: page_number += 1 - # TODO (betodealmeida): improve this - url = workspace_url / "roles/add" - _logger.debug("GET %s", url) - response = self.session.get(url) - soup = BeautifulSoup(response.text, features="html.parser") - select = soup.find("select", id="user") - ids = { - option.text: int(option.attrs["value"]) - for option in select.find_all("option") - } + ids = {} + page = 0 + while True: + query = prison.dumps( + { + "page": page, + "page_size": SUPERSET_MAX_PAGE_SIZE, + }, + ) + url = workspace_url / "api/v1/chart/related/owners" % {"q": query} + _logger.debug("GET %s", url) + response = self.session.get(url) + + validate_response(response) + payload = response.json() + if not payload["result"]: + break + + ids.update( + {user["extra"]["email"]: user["value"] for user in payload["result"]}, + ) + + page += 1 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: - team_member["id"] = ids[full_name] + if team_member["email"] in ids: + team_member["id"] = ids[team_member["email"]] yield team_member def import_users(self, teams: List[str], users: List[UserType]) -> None: diff --git a/tests/api/clients/preset_test.py b/tests/api/clients/preset_test.py index cfad198..d549d0e 100644 --- a/tests/api/clients/preset_test.py +++ b/tests/api/clients/preset_test.py @@ -136,13 +136,26 @@ def test_preset_client_export_users(requests_mock: Mocker) -> None: ) requests_mock.get( - "https://superset.example.org/roles/add", - text=""" - - """, + "https://superset.example.org/api/v1/chart/related/owners?q=(page:0,page_size:100)", + json={ + "count": 2, + "result": [ + { + "extra": {"active": True, "email": "adoe@example.com"}, + "text": "Alice Doe", + "value": 1, + }, + { + "extra": {"active": True, "email": "bdoe@example.com"}, + "text": "Bob Doe", + "value": 2, + }, + ], + }, + ) + requests_mock.get( + "https://superset.example.org/api/v1/chart/related/owners?q=(page:1,page_size:100)", + json={"count": 2, "result": []}, ) auth = Auth() diff --git a/tests/api/clients/superset_test.py b/tests/api/clients/superset_test.py index 24902a5..8c55b8d 100644 --- a/tests/api/clients/superset_test.py +++ b/tests/api/clients/superset_test.py @@ -1927,14 +1927,40 @@ def test_export_users_preset(requests_mock: Mocker) -> None: ) requests_mock.get( - "https://superset.example.org/roles/add", - text=""" - - """, + "https://superset.example.org/api/v1/chart/related/owners?q=(page:0,page_size:100)", + json={ + "count": 3, + "result": [ + { + "extra": { + "active": True, + "email": "adoe@example.com", + }, + "text": "Alice Doe", + "value": 1, + }, + { + "extra": { + "active": True, + "email": "bdoe@example.com", + }, + "text": "Bob Doe", + "value": 2, + }, + { + "extra": { + "active": True, + "email": "cdoe@example.com", + }, + "text": "Clarisse Doe", + "value": 3, + }, + ], + }, + ) + requests_mock.get( + "https://superset.example.org/api/v1/chart/related/owners?q=(page:1,page_size:100)", + json={"count": 3, "result": []}, # empty result to break the loop ) auth = Auth()