From 495ad0b24a477331cc9f28f7976f33903a3aa202 Mon Sep 17 00:00:00 2001 From: Vitor Avila Date: Fri, 28 Feb 2025 01:35:42 -0300 Subject: [PATCH 1/2] fix(Preset): Export users --- src/preset_cli/api/clients/preset.py | 41 +++++++++++++++++---------- tests/api/clients/preset_test.py | 29 +++++++++++++------ tests/api/clients/superset_test.py | 42 ++++++++++++++++++++++------ 3 files changed, 82 insertions(+), 30 deletions(-) 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..8ea7e85 100644 --- a/tests/api/clients/preset_test.py +++ b/tests/api/clients/preset_test.py @@ -136,17 +136,30 @@ 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() - client = PresetClient("https://ws.preset.io/", auth) + client = PresetClient("https://api.app.preset.io/", auth) assert list(client.export_users(URL("https://superset.example.org/"))) == [ { "id": 1, 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() From 5a73633a62f02cdf4ac9066334dace2937237208 Mon Sep 17 00:00:00 2001 From: Vitor Avila Date: Fri, 28 Feb 2025 01:59:02 -0300 Subject: [PATCH 2/2] fixing test --- tests/api/clients/preset_test.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/api/clients/preset_test.py b/tests/api/clients/preset_test.py index 8ea7e85..d549d0e 100644 --- a/tests/api/clients/preset_test.py +++ b/tests/api/clients/preset_test.py @@ -159,7 +159,7 @@ def test_preset_client_export_users(requests_mock: Mocker) -> None: ) auth = Auth() - client = PresetClient("https://api.app.preset.io/", auth) + client = PresetClient("https://ws.preset.io/", auth) assert list(client.export_users(URL("https://superset.example.org/"))) == [ { "id": 1,