Merge pull request #75 from preset-io/export-roles

feat: export roles

README.rst

@@ -113,6 +113,7 @@ The following commands are currently available:
 - ``preset-cli superset export-assets``: export resources (databases, datasets, charts, dashboards) into a directory as YAML files.
 - ``preset-cli superset export-ownership``: export resource ownership (UUID -> email) into a YAML file.
 - ``preset-cli superset export-rls``: export RLS rules into a YAML file.
+- ``preset-cli superset export-roles``: export user roles into a YAML file.
 - ``preset-cli superset export-users``: export users (name, username, email, roles) into a YAML file.
 - ``preset-cli superset sync native``: synchronize the workspace from a directory of templated configuration files.
 - ``preset-cli superset sync dbt-core``: synchronize the workspace from a dbt Core project.
@@ -427,6 +428,8 @@ Exporting users
 
 The ``preset-cli superset export-users`` command can be used to export a list of users. These users can then be imported to Preset via the ``preset-cli import-users`` command.
 
+You can also export roles via ``preset-cli superset export-roles``, and import with ``import-roles``.
+
 Exporting RLS rules
 ~~~~~~~~~~~~~~~~~~~
 

src/preset_cli/api/clients/superset.py

@@ -156,6 +156,15 @@ def parse_html_array(value: str) -> List[str]:
     return [part for part in parts if part.strip()]
 
 
+class RoleType(TypedDict):
+    """
+    Schema for a role.
+    """
+
+    name: str
+    permissions: List[str]
+
+
 class RuleType(TypedDict):
     """
     Schema for an RLS rule.
@@ -677,6 +686,51 @@ def _export_users_superset(
                     "role": parse_html_array(tds[6].text.strip()),
                 }
 
+    def export_roles(self) -> Iterator[RoleType]:
+        """
+        Return all roles.
+        """
+        session = self.auth.get_session()
+        headers = self.auth.get_headers()
+        headers["Referer"] = str(self.baseurl)
+
+        page = 0
+        while True:
+            params = {
+                "psize_RoleModelView": MAX_PAGE_SIZE,
+                "page_RoleModelView": page,
+            }
+            url = self.baseurl / "roles/list/"
+            page += 1
+
+            response = session.get(url, params=params, headers=headers)
+            soup = BeautifulSoup(response.text, features="html.parser")
+            table = soup.find_all("table")[1]
+            trs = table.find_all("tr")
+            if len(trs) == 1:
+                break
+
+            for tr in trs[1:]:  # pylint: disable=invalid-name
+                tds = tr.find_all("td")
+
+                role_id = int(tds[0].find("input").attrs["id"])
+                role_url = self.baseurl / "roles/show" / str(role_id)
+
+                response = session.get(role_url, headers=headers)
+                soup = BeautifulSoup(response.text, features="html.parser")
+                table = soup.find_all("table")[-1]
+                keys: List[Tuple[str, Callable[[Any], Any]]] = [
+                    ("name", str),
+                    ("permissions", parse_html_array),
+                ]
+                yield cast(
+                    RoleType,
+                    {
+                        key: parse(tr.find("td").text.strip())
+                        for (key, parse), tr in zip(keys, table.find_all("tr"))
+                    },
+                )
+
     def export_rls(self) -> Iterator[RuleType]:
         """
         Return all RLS rules.
@@ -696,7 +750,10 @@ def export_rls(self) -> Iterator[RuleType]:
 
             response = session.get(url, params=params, headers=headers)
             soup = BeautifulSoup(response.text, features="html.parser")
-            table = soup.find_all("table")[1]
+            try:
+                table = soup.find_all("table")[1]
+            except IndexError:
+                return
             trs = table.find_all("tr")
             if len(trs) == 1:
                 break

src/preset_cli/cli/superset/export.py

@@ -115,6 +115,25 @@ def export_users(ctx: click.core.Context, path: str) -> None:
         yaml.dump(users, output)
 
 
+@click.command()
+@click.argument(
+    "path",
+    type=click.Path(resolve_path=True),
+    default="roles.yaml",
+)
+@click.pass_context
+def export_roles(ctx: click.core.Context, path: str) -> None:
+    """
+    Export roles to a YAML file.
+    """
+    auth = ctx.obj["AUTH"]
+    url = URL(ctx.obj["INSTANCE"])
+    client = SupersetClient(url, auth)
+
+    with open(path, "w", encoding="utf-8") as output:
+        yaml.dump(list(client.export_roles()), output)
+
+
 @click.command()
 @click.argument(
     "path",

src/preset_cli/cli/superset/main.py

@@ -11,6 +11,7 @@
     export_assets,
     export_ownership,
     export_rls,
+    export_roles,
     export_users,
 )
 from preset_cli.cli.superset.import_ import import_ownership, import_rls
@@ -56,6 +57,7 @@ def superset_cli(
 superset_cli.add_command(export_assets, name="export")  # for backwards compatibility
 superset_cli.add_command(export_users)
 superset_cli.add_command(export_rls)
+superset_cli.add_command(export_roles)
 superset_cli.add_command(export_ownership)
 superset_cli.add_command(import_rls)
 superset_cli.add_command(import_ownership)

tests/api/clients/superset_test.py

@@ -1387,6 +1387,115 @@ def test_export_users_preset(requests_mock: Mocker) -> None:
     ]
 
 
+def test_export_roles(requests_mock: Mocker) -> None:
+    """
+    Test ``export_roles``.
+    """
+    requests_mock.get(
+        (
+            "https://superset.example.org/roles/list/?"
+            "psize_RoleModelView=100&"
+            "page_RoleModelView=0"
+        ),
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table></table>
+    <table>
+      <tr>
+        <th></th>
+        <th>Name</th>
+      </tr>
+      <tr>
+        <td><input id="1" /></td>
+        <td>Admin</td>
+      </tr>
+      <tr>
+        <td><input id="2" /></td>
+        <td>Public</td>
+      </tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+    requests_mock.get(
+        (
+            "https://superset.example.org/roles/list/?"
+            "psize_RoleModelView=100&"
+            "page_RoleModelView=1"
+        ),
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table></table>
+    <table>
+      <tr>
+        <th></th>
+        <th>Name</th>
+      </tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+    requests_mock.get(
+        "https://superset.example.org/roles/show/1",
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table>
+      <tr><th>Name</th><td>Admin</td></tr>
+      <tr><th>Permissions</th><td>[can this, can that]</td></tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+    requests_mock.get(
+        "https://superset.example.org/roles/show/2",
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table>
+      <tr><th>Name</th><td>Public</td></tr>
+      <tr><th>Permissions</th><td>[]</td></tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    assert list(client.export_roles()) == [
+        {
+            "name": "Admin",
+            "permissions": ["can this", "can that"],
+        },
+        {
+            "name": "Public",
+            "permissions": [],
+        },
+    ]
+
+
 def test_export_rls(requests_mock: Mocker) -> None:
     """
     Test ``export_rls``.
@@ -1498,6 +1607,35 @@ def test_export_rls(requests_mock: Mocker) -> None:
     ]
 
 
+def test_export_rls_no_rules(requests_mock: Mocker) -> None:
+    """
+    Test ``export_rls``.
+    """
+    requests_mock.get(
+        (
+            "https://superset.example.org/rowlevelsecurityfiltersmodelview/list/?"
+            "psize_RowLevelSecurityFiltersModelView=100&"
+            "page_RowLevelSecurityFiltersModelView=0"
+        ),
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table></table>
+    No records found
+  </body>
+</html>
+        """,
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    assert list(client.export_rls()) == []
+
+
 def test_export_ownership(mocker: MockerFixture) -> None:
     """
     Test ``export_ownership``.

tests/cli/superset/export_test.py

@@ -218,9 +218,41 @@ def test_export_users(mocker: MockerFixture, fs: FakeFilesystem) -> None:
     ]
 
 
+def test_export_roles(mocker: MockerFixture, fs: FakeFilesystem) -> None:
+    """
+    Test the ``export_roles`` command.
+    """
+    mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth")
+    SupersetClient = mocker.patch("preset_cli.cli.superset.export.SupersetClient")
+    client = SupersetClient()
+    client.export_roles.return_value = [
+        {
+            "name": "Public",
+            "permissions": [],
+        },
+    ]
+
+    runner = CliRunner()
+    result = runner.invoke(
+        superset_cli,
+        ["https://superset.example.org/", "export-roles", "roles.yaml"],
+        catch_exceptions=False,
+    )
+    assert result.exit_code == 0
+
+    with open("roles.yaml", encoding="utf-8") as input_:
+        contents = yaml.load(input_, Loader=yaml.SafeLoader)
+    assert contents == [
+        {
+            "name": "Public",
+            "permissions": [],
+        },
+    ]
+
+
 def test_export_rls(mocker: MockerFixture, fs: FakeFilesystem) -> None:
     """
-    Test the ``export_users`` command.
+    Test the ``export_rls`` command.
     """
     mocker.patch("preset_cli.cli.superset.main.UsernamePasswordAuth")
     SupersetClient = mocker.patch("preset_cli.cli.superset.export.SupersetClient")

tests/cli/superset/main_test.py

@@ -116,6 +116,7 @@ def test_superset() -> None:
   export-assets
   export-ownership
   export-rls
+  export-roles
   export-users
   import-assets
   import-ownership