Merge pull request #227 from preset-io/rls-fix

fix(export-rls): Support Preset and older Superset versions

Author: Vitor-Avila

src/preset_cli/api/clients/superset.py

@@ -191,7 +191,7 @@ class RuleType(TypedDict):
     Schema for an RLS rule.
     """
 
-    name: str
+    name: Optional[str]
     description: Optional[str]
     filter_type: str
     tables: List[str]
@@ -690,6 +690,12 @@ def import_zip(
 
         return payload["message"] == "OK"
 
+    def get_rls(self, **kwargs: str) -> List[Any]:
+        """
+        Return RLS rules, possibly filtered.
+        """
+        return self.get_resources("rowlevelsecurity", **kwargs)
+
     def export_users(self) -> Iterator[UserType]:
         """
         Return all users.
@@ -808,9 +814,9 @@ def export_roles(self) -> Iterator[RoleType]:  # pylint: disable=too-many-locals
                     "users": users,
                 }
 
-    def export_rls(self) -> Iterator[RuleType]:
+    def export_rls_legacy(self) -> Iterator[RuleType]:
         """
-        Return all RLS rules.
+        Return all RLS rules from legacy endpoint.
         """
         page = 0
         while True:
@@ -857,6 +863,12 @@ def export_rls(self) -> Iterator[RuleType]:
                     ("group_key", str),
                     ("clause", str),
                 ]
+
+                # Before Superset 2.1.0, RLS dont have name and description
+                if table.find("th").text.strip() == "Filter Type":
+                    keys.remove(("name", str))
+                    keys.remove(("description", str))
+
                 yield cast(
                     RuleType,
                     {
@@ -865,6 +877,41 @@ def export_rls(self) -> Iterator[RuleType]:
                     },
                 )
 
+    def export_rls(self) -> Iterator[RuleType]:
+        """
+        Return all RLS rules.
+        """
+        url = self.baseurl / "api/v1/rowlevelsecurity/"
+        response = self.session.get(url)
+        if response.status_code == 200:
+            for rule in self.get_rls():
+                keys = [
+                    "name",
+                    "description",
+                    "filter_type",
+                    "tables",
+                    "roles",
+                    "group_key",
+                    "clause",
+                ]
+                data = {}
+                for key in keys:
+                    if key == "tables":
+                        data[key] = [
+                            f"{inner_item['schema']}.{inner_item['table_name']}"
+                            for inner_item in rule.get(key, [])
+                        ]
+                    elif key == "roles":
+                        data[key] = [
+                            inner_item["name"] for inner_item in rule.get(key, [])
+                        ]
+                    else:
+                        data[key] = rule.get(key)
+                yield cast(RuleType, data)
+
+        else:
+            yield from self.export_rls_legacy()
+
     def import_role(self, role: RoleType) -> None:  # pylint: disable=too-many-locals
         """
         Import a given role.

src/preset_cli/cli/superset/export.py

@@ -241,7 +241,7 @@ def export_rls(ctx: click.core.Context, path: str) -> None:
     client = SupersetClient(url, auth)
 
     with open(path, "w", encoding="utf-8") as output:
-        yaml.dump(list(client.export_rls()), output)
+        yaml.dump(list(client.export_rls()), output, sort_keys=False)
 
 
 @click.command()

tests/api/clients/superset_test.py

@@ -1235,6 +1235,18 @@ def test_import_zip_error(requests_mock: Mocker) -> None:
     )
 
 
+def test_get_rls(mocker: MockerFixture) -> None:
+    """
+    Test the ``get_rls`` method.
+    """
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    get_resources = mocker.patch.object(client, "get_resources")
+
+    client.get_rls()
+    get_resources.assert_called_with("rowlevelsecurity")
+
+
 def test_export_users(requests_mock: Mocker) -> None:
     """
     Test ``export_users``.
@@ -1668,9 +1680,9 @@ def test_export_roles_anchor_role_id(
     ]
 
 
-def test_export_rls(requests_mock: Mocker) -> None:
+def test_export_rls_legacy(requests_mock: Mocker) -> None:
     """
-    Test ``export_rls``.
+    Test ``export_rls_legacy``.
     """
     requests_mock.get(
         (
@@ -1766,7 +1778,7 @@ def test_export_rls(requests_mock: Mocker) -> None:
 
     auth = Auth()
     client = SupersetClient("https://superset.example.org/", auth)
-    assert list(client.export_rls()) == [
+    assert list(client.export_rls_legacy()) == [
         {
             "name": "My Rule",
             "description": "This is a rule. There are many others like it, but this one is mine.",
@@ -1779,9 +1791,173 @@ def test_export_rls(requests_mock: Mocker) -> None:
     ]
 
 
-def test_export_rls_no_rules(requests_mock: Mocker) -> None:
+def test_export_rls_legacy_older_superset(requests_mock: Mocker) -> None:
+    """
+    Test ``export_rls_legacy`` with older Superset version.
+    """
+    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>
+    <table>
+      <tr>
+        <th></th>
+        <th>Filter Type</th>
+        <th>Tables</th>
+        <th>Roles</th>
+        <th>Clause</th>
+        <th>Creator</th>
+        <th>Modified</th>
+      </tr>
+      <tr>
+        <td><input id="1" /></td>
+        <td>Regular</td>
+        <td>[main.test_table]</td>
+        <td>client_id = 9</td>
+        <td>admin admin</td>
+        <td>35 minutes ago</td>
+      </tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+    requests_mock.get(
+        (
+            "https://superset.example.org/rowlevelsecurityfiltersmodelview/list/?"
+            "psize_RowLevelSecurityFiltersModelView=100&"
+            "page_RowLevelSecurityFiltersModelView=1"
+        ),
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table></table>
+    <table>
+      <tr>
+        <th></th>
+        <th>Filter Type</th>
+        <th>Tables</th>
+        <th>Roles</th>
+        <th>Clause</th>
+        <th>Creator</th>
+        <th>Modified</th>
+      </tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+    requests_mock.get(
+        "https://superset.example.org/rowlevelsecurityfiltersmodelview/show/1",
+        text="""
+<!DOCTYPE html>
+<html lang="en">
+  <head>
+    <meta charset="utf-8">
+  </head>
+  <body>
+    <table>
+      <tr><th>Filter Type</th><td>Regular</td></tr>
+      <tr><th>Tables</th><td>[main.test_table]</td></tr>
+      <tr><th>Roles</th><td>[Gamma]</td></tr>
+      <tr><th>Group Key</th><td>department</td></tr>
+      <tr><th>Clause</th><td>client_id = 9</td></tr>
+    </table>
+  </body>
+</html>
+        """,
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    assert list(client.export_rls_legacy()) == [
+        {
+            "filter_type": "Regular",
+            "tables": ["main.test_table"],
+            "roles": ["Gamma"],
+            "group_key": "department",
+            "clause": "client_id = 9",
+        },
+    ]
+
+
+def test_export_rls_legacy_route(requests_mock: Mocker, mocker: MockerFixture) -> None:
+    """
+    Test ``export_rls`` going through the legacy route
+    """
+    requests_mock.get(
+        "https://superset.example.org/api/v1/rowlevelsecurity/",
+        status_code=404,
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+
+    export_rls_legacy = mocker.patch.object(client, "export_rls_legacy")
+    list(client.export_rls())
+    export_rls_legacy.assert_called_once()
+
+
+def test_export_rls(requests_mock: Mocker, mocker: MockerFixture) -> None:
+    """
+    Test ``export_rls``
+    """
+    requests_mock.get(
+        "https://superset.example.org/api/v1/rowlevelsecurity/",
+        status_code=200,
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+
+    get_rls = mocker.patch.object(client, "get_rls")
+    get_rls.return_value = [
+        {
+            "changed_on_delta_humanized": "2 days ago",
+            "clause": "client_id = 9",
+            "description": "This is a rule. There are many others like it, but this one is mine.",
+            "filter_type": "Regular",
+            "group_key": "department",
+            "id": 9,
+            "name": "My Rule",
+            "roles": [{"id": 1, "name": "Admin"}, {"id": 2, "name": "Gamma"}],
+            "tables": [
+                {"id": 18, "schema": "main", "table_name": "test_table"},
+                {"id": 20, "schema": "main", "table_name": "second_test"},
+            ],
+        },
+    ]
+
+    assert list(client.export_rls()) == [
+        {
+            "name": "My Rule",
+            "description": "This is a rule. There are many others like it, but this one is mine.",
+            "filter_type": "Regular",
+            "tables": ["main.test_table", "main.second_test"],
+            "roles": ["Admin", "Gamma"],
+            "group_key": "department",
+            "clause": "client_id = 9",
+        },
+    ]
+
+
+def test_export_rls_legacy_no_rules(requests_mock: Mocker) -> None:
     """
-    Test ``export_rls``.
+    Test ``export_rls_legacy`` with no rows returned.
     """
     requests_mock.get(
         (
@@ -1805,7 +1981,7 @@ def test_export_rls_no_rules(requests_mock: Mocker) -> None:
 
     auth = Auth()
     client = SupersetClient("https://superset.example.org/", auth)
-    assert list(client.export_rls()) == []
+    assert list(client.export_rls_legacy()) == []
 
 
 def test_export_ownership(mocker: MockerFixture) -> None: