feat: Remove old users from export ownership (#365)

Vitor-Avila · web-flow · commit 03842c11b215 · 2025-11-06T16:52:46.000-03:00
* feat: Remove old users from export ownership
* Fixing type hint
* Fixing pre-commit compatibility issues
diff --git a/dev-requirements.txt b/dev-requirements.txt
@@ -40,7 +40,7 @@ mccabe==0.7.0
 multidict==6.0.2
 nodeenv==1.7.0
 numpy==1.23.1
-packaging==21.3
+packaging==24.0
 pandas==1.4.3
 pep517==0.13.0
 pip-tools==6.13.0
diff --git a/src/preset_cli/api/clients/superset.py b/src/preset_cli/api/clients/superset.py
@@ -1165,11 +1165,15 @@ def get_role_id(self, role_name: str) -> int:
 
         return id_
 
-    def export_ownership(self, resource_name: str) -> Iterator[OwnershipType]:
+    def export_ownership(
+        self,
+        resource_name: str,
+        users: Dict[int, str],
+        exclude_old_users: bool,
+    ) -> Iterator[OwnershipType]:
         """
         Return information about resource ownership.
         """
-        emails = {user["id"]: user["email"] for user in self.export_users()}
         uuids = self.get_uuids(resource_name)
         name_key = {
             "dataset": "table_name",
@@ -1178,11 +1182,30 @@ def export_ownership(self, resource_name: str) -> Iterator[OwnershipType]:
         }[resource_name]
 
         for resource in self.get_resources(resource_name):
-            yield {
+            info: OwnershipType = {
                 "name": resource[name_key],
                 "uuid": uuids[resource["id"]],
-                "owners": [emails[owner["id"]] for owner in resource.get("owners", [])],
+                "owners": [],
             }
+            for owner in resource.get("owners", []):
+                if user := users.get(owner["id"]):
+                    info["owners"].append(user)
+                else:
+                    if not exclude_old_users:
+                        raise Exception(
+                            f"User {owner['first_name']} {owner['last_name']} owns the "
+                            f"{resource_name} {resource[name_key]} but is not a member in the team."
+                            " You can trigger this command with ``--exclude-old-users`` to "
+                            "automatically remove old members.",
+                        )
+                    _logger.warning(
+                        "User %s %s not found in team. Removing from ownership config for %s %s.",
+                        owner["first_name"],
+                        owner["last_name"],
+                        resource_name,
+                        resource[name_key],
+                    )
+            yield info
 
     def import_ownership(
         self,
diff --git a/src/preset_cli/cli/superset/export.py b/src/preset_cli/cli/superset/export.py
@@ -456,17 +456,30 @@ def export_rls(
     type=click.Path(resolve_path=True),
     default="ownership.yaml",
 )
+@click.option(
+    "--asset-type",
+    help="Asset type",
+    multiple=True,
+)
 @click.option(
     "--force-unix-eol",
     is_flag=True,
     default=False,
     help="Force Unix end-of-line characters, otherwise use system default",
 )
+@click.option(
+    "--exclude-old-users",
+    is_flag=True,
+    default=False,
+    help="Exclude users that are no longer a member in the team from the generated file",
+)
 @click.pass_context
 def export_ownership(
     ctx: click.core.Context,
     path: str,
+    asset_type: Tuple[str, ...],
     force_unix_eol: bool = False,
+    exclude_old_users: bool = False,
 ) -> None:
     """
     Export DBs/datasets/charts/dashboards ownership to a YAML file.
@@ -475,9 +488,21 @@ def export_ownership(
     url = URL(ctx.obj["INSTANCE"])
     client = SupersetClient(url, auth)
 
+    users = {user["id"]: user["email"] for user in client.export_users()}
+    asset_types = set(asset_type) if asset_type else set()
+    resources_to_export = (
+        [r for r in ["dataset", "chart", "dashboard"] if r in asset_types]
+        if asset_types
+        else ["dataset", "chart", "dashboard"]
+    )
+
     ownership = defaultdict(list)
-    for resource_name in ["dataset", "chart", "dashboard"]:
-        for resource in client.export_ownership(resource_name):
+    for resource_name in resources_to_export:
+        for resource in client.export_ownership(
+            resource_name,
+            users,
+            exclude_old_users,
+        ):
             ownership[resource_name].append(
                 {
                     "name": resource["name"],
@@ -486,6 +511,10 @@ def export_ownership(
                 },
             )
 
-    newline = get_newline_char(force_unix_eol)
-    with open(path, "w", encoding="utf-8", newline=newline) as output:
+    with open(
+        path,
+        "w",
+        encoding="utf-8",
+        newline=get_newline_char(force_unix_eol),
+    ) as output:
         yaml.dump(dict(ownership), output)
diff --git a/tests/api/clients/superset_test.py b/tests/api/clients/superset_test.py
@@ -2577,12 +2577,41 @@ def test_export_ownership(mocker: MockerFixture) -> None:
     """
     mocker.patch.object(
         SupersetClient,
-        "export_users",
+        "get_uuids",
+        return_value={
+            1: UUID("e0d20af0-cef9-4bdb-80b4-745827f441bf"),
+        },
+    )
+    mocker.patch.object(
+        SupersetClient,
+        "get_resources",
         return_value=[
-            {"id": 1, "email": "admin@example.com"},
-            {"id": 2, "email": "adoe@example.com"},
+            {
+                "slice_name": "My chart",
+                "id": 1,
+                "owners": [{"id": 1}, {"id": 2}],
+            },
         ],
     )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    users = {1: "admin@example.com", 2: "adoe@example.com"}
+    assert list(client.export_ownership("chart", users, exclude_old_users=False)) == [
+        {
+            "name": "My chart",
+            "owners": ["admin@example.com", "adoe@example.com"],
+            "uuid": UUID("e0d20af0-cef9-4bdb-80b4-745827f441bf"),
+        },
+    ]
+
+
+def test_export_ownership_user_not_found_raises_exception(
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test ``export_ownership`` when user not found and exclude_old_users=False.
+    """
     mocker.patch.object(
         SupersetClient,
         "get_uuids",
@@ -2597,17 +2626,64 @@ def test_export_ownership(mocker: MockerFixture) -> None:
             {
                 "slice_name": "My chart",
                 "id": 1,
-                "owners": [{"id": 1}, {"id": 2}],
+                "owners": [
+                    {"id": 1, "first_name": "John", "last_name": "Doe"},
+                    {"id": 999, "first_name": "Missing", "last_name": "User"},
+                ],
+            },
+        ],
+    )
+
+    auth = Auth()
+    client = SupersetClient("https://superset.example.org/", auth)
+    users = {1: "admin@example.com"}  # User 999 is not in the dict
+
+    with pytest.raises(
+        Exception,
+        match="User Missing User owns the chart My chart but is not a member in the team",
+    ):
+        list(client.export_ownership("chart", users, exclude_old_users=False))
+
+
+def test_export_ownership_user_not_found_with_exclude_flag(
+    mocker: MockerFixture,
+) -> None:
+    """
+    Test ``export_ownership`` when user not found and exclude_old_users=True.
+    """
+    mocker.patch.object(
+        SupersetClient,
+        "get_uuids",
+        return_value={
+            1: UUID("e0d20af0-cef9-4bdb-80b4-745827f441bf"),
+        },
+    )
+    mocker.patch.object(
+        SupersetClient,
+        "get_resources",
+        return_value=[
+            {
+                "slice_name": "My chart",
+                "id": 1,
+                "owners": [
+                    {"id": 1, "first_name": "John", "last_name": "Doe"},
+                    {"id": 999, "first_name": "Missing", "last_name": "User"},
+                ],
             },
         ],
     )
 
     auth = Auth()
     client = SupersetClient("https://superset.example.org/", auth)
-    assert list(client.export_ownership("chart")) == [
+    users = {1: "admin@example.com"}  # User 999 is not in the dict
+
+    # Should not raise exception, and should exclude the missing user
+    result = list(client.export_ownership("chart", users, exclude_old_users=True))
+
+    assert result == [
         {
             "name": "My chart",
-            "owners": ["admin@example.com", "adoe@example.com"],
+            "owners": ["admin@example.com"],  # Only user 1, user 999 excluded
             "uuid": UUID("e0d20af0-cef9-4bdb-80b4-745827f441bf"),
         },
     ]
diff --git a/tests/cli/superset/export_test.py b/tests/cli/superset/export_test.py
