From 3633bfbb219834d5d2489c8955bac297d693cee0 Mon Sep 17 00:00:00 2001
From: John Mulligan <jmulligan@redhat.com>
Date: Wed, 2 Jul 2025 14:25:32 -0400
Subject: [PATCH] mgr/smb: move prune function to be a staging store method

The prune function was tightly linked to the staging store. Re-implement
it as more generic operation on the staging store.
Continue shrinking handler.py a bit and preparing to make it simpler to
add future SMBResource types.

Signed-off-by: John Mulligan <jmulligan@redhat.com>
---
 src/pybind/mgr/smb/handler.py | 21 +--------------------
 src/pybind/mgr/smb/staging.py | 23 +++++++++++++++++++++++
 2 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/src/pybind/mgr/smb/handler.py b/src/pybind/mgr/smb/handler.py
index a6221387d787c..5fdf72be1d1e0 100644
--- a/src/pybind/mgr/smb/handler.py
+++ b/src/pybind/mgr/smb/handler.py
@@ -308,7 +308,7 @@ class ClusterConfigHandler:
             )
             with _store_transaction(staging.destination_store):
                 results = staging.save()
-                _prune_linked_entries(staging)
+                staging.prune_linked_entries()
             with _store_transaction(staging.destination_store):
                 self._sync_modified(results)
         return results
@@ -676,25 +676,6 @@ def order_resources(
     return sorted(resource_objs, key=_keyfunc)
 
 
-def _prune_linked_entries(staging: Staging) -> None:
-    cids = set(ClusterEntry.ids(staging))
-    for auth_id in JoinAuthEntry.ids(staging):
-        join_auth = staging.get_join_auth(auth_id)
-        if (
-            join_auth.linked_to_cluster
-            and join_auth.linked_to_cluster not in cids
-        ):
-            JoinAuthEntry.from_store(
-                staging.destination_store, auth_id
-            ).remove()
-    for ug_id in UsersAndGroupsEntry.ids(staging):
-        ug = staging.get_users_and_groups(ug_id)
-        if ug.linked_to_cluster and ug.linked_to_cluster not in cids:
-            UsersAndGroupsEntry.from_store(
-                staging.destination_store, ug_id
-            ).remove()
-
-
 def _generate_share(
     share: resources.Share, resolver: PathResolver, cephx_entity: str
 ) -> Dict[str, Dict[str, str]]:
diff --git a/src/pybind/mgr/smb/staging.py b/src/pybind/mgr/smb/staging.py
index 00c7c084a6e82..658e9a03ff957 100644
--- a/src/pybind/mgr/smb/staging.py
+++ b/src/pybind/mgr/smb/staging.py
@@ -6,6 +6,7 @@ from typing import (
     List,
     Optional,
     Set,
+    Type,
 )
 
 import functools
@@ -27,6 +28,7 @@ from .enums import (
 from .internal import (
     ClusterEntry,
     JoinAuthEntry,
+    ResourceEntry,
     ShareEntry,
     UsersAndGroupsEntry,
     resource_entry,
@@ -141,6 +143,27 @@ class Staging:
         result = Result(resource, success=True, status={'state': state})
         return result
 
+    def _prune(
+        self,
+        cids: Collection[str],
+        ecls: Type[ResourceEntry],
+        rcls: Type[SMBResource],
+    ) -> None:
+        for _id in ecls.ids(self):
+            assert isinstance(_id, str)
+            rentry = ecls.from_store_by_key(self.destination_store, _id)
+            resource = rentry.get_resource_type(rcls)
+            _linked_to_cluster = getattr(resource, 'linked_to_cluster', None)
+            if not _linked_to_cluster:
+                continue
+            if _linked_to_cluster not in cids:
+                rentry.remove()
+
+    def prune_linked_entries(self) -> None:
+        cids = set(ClusterEntry.ids(self))
+        self._prune(cids, JoinAuthEntry, resources.JoinAuth)
+        self._prune(cids, UsersAndGroupsEntry, resources.UsersAndGroups)
+
 
 def auth_refs(cluster: resources.Cluster) -> Collection[str]:
     """Return all IDs for join_auth resources referenced by the supplied
-- 
2.39.5