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()