Merge pull request #8577 from OpenMined/aziz/storageperm

Add storage permissions to mongo document storage

packages/syft/src/syft/store/mongo_client.py

@@ -236,6 +236,31 @@ def with_collection_permissions(
 
         return Ok(collection_permissions)
 
+    def with_collection_storage_permissions(
+        self, collection_settings: PartitionSettings, store_config: StoreConfig
+    ) -> Result[MongoCollection, Err]:
+        """
+        For each collection, create a corresponding collection
+        that store the permissions to the data in that collection
+        """
+        res = self.with_db(db_name=store_config.db_name)
+        if res.is_err():
+            return res
+        db = res.ok()
+
+        try:
+            collection_storage_permissions_name: str = (
+                collection_settings.name + "_storage_permissions"
+            )
+            storage_permissons_collection = db.get_collection(
+                name=collection_storage_permissions_name,
+                codec_options=SYFT_CODEC_OPTIONS,
+            )
+        except BaseException as e:
+            return Err(str(e))
+
+        return Ok(storage_permissons_collection)
+
     def close(self) -> None:
         self.client.close()
         MongoClientCache.__client_cache__.pop(hash(str(self.config)), None)

packages/syft/src/syft/store/mongo_document_store.py

@@ -22,6 +22,7 @@
 from ..service.action.action_permissions import ActionObjectREAD
 from ..service.action.action_permissions import ActionObjectWRITE
 from ..service.action.action_permissions import ActionPermission
+from ..service.action.action_permissions import StoragePermission
 from ..service.context import AuthedServiceContext
 from ..service.response import SyftSuccess
 from ..types.syft_object import SYFT_OBJECT_VERSION_2
@@ -148,8 +149,18 @@ def init_store(self) -> Result[Ok, Err]:
         if collection_permissions_status.is_err():
             return collection_permissions_status
 
+        collection_storage_permissions_status = (
+            client.with_collection_storage_permissions(
+                collection_settings=self.settings, store_config=self.store_config
+            )
+        )
+
+        if collection_storage_permissions_status.is_err():
+            return collection_storage_permissions_status
+
         self._collection = collection_status.ok()
         self._permissions = collection_permissions_status.ok()
+        self._storage_permissions = collection_storage_permissions_status.ok()
 
         return self._create_update_index()
 
@@ -233,6 +244,15 @@ def permissions(self) -> Result[MongoCollection, Err]:
 
         return Ok(self._permissions)
 
+    @property
+    def storage_permissions(self) -> Result[MongoCollection, Err]:
+        if not hasattr(self, "_storage_permissions"):
+            res = self.init_store()
+            if res.is_err():
+                return res
+
+        return Ok(self._storage_permissions)
+
     def set(self, *args: Any, **kwargs: Any) -> Result[SyftObject, str]:
         return self._set(*args, **kwargs)
 
@@ -291,8 +311,12 @@ def _set(
                 self.add_permissions(add_permissions)
 
             if add_storage_permission:
-                # TODO: add storage permissions to Mongo store
-                pass
+                self.add_storage_permission(
+                    StoragePermission(
+                        uid=obj.id,
+                        node_uid=self.node_uid,
+                    )
+                )
 
             return Ok(obj)
         else:
@@ -533,6 +557,82 @@ def remove_permission(
         else:
             return Err(f"the permission {permission.permission_string} does not exist!")
 
+    def add_storage_permission(self, storage_permission: StoragePermission) -> None:
+        storage_permissions_or_err = self.storage_permissions
+        if storage_permissions_or_err.is_err():
+            return storage_permissions_or_err
+        storage_permissions_collection: MongoCollection = (
+            storage_permissions_or_err.ok()
+        )
+
+        storage_permissions: dict | None = storage_permissions_collection.find_one(
+            {"_id": storage_permission.uid}
+        )
+        if storage_permissions is None:
+            # Permission doesn't exist, add a new one
+            storage_permissions_collection.insert_one(
+                {
+                    "_id": storage_permission.uid,
+                    "node_uids": {storage_permission.node_uid},
+                }
+            )
+        else:
+            # update the permissions with the new permission string
+            node_uids: set = storage_permissions["node_uids"]
+            node_uids.add(storage_permission.node_uid)
+            storage_permissions_collection.update_one(
+                {"_id": storage_permission.uid},
+                {"$set": {"node_uids": node_uids}},
+            )
+
+    def add_storage_permissions(self, permissions: list[StoragePermission]) -> None:
+        for permission in permissions:
+            self.add_storage_permission(permission)
+
+    def has_storage_permission(self, permission: StoragePermission) -> bool:  # type: ignore
+        """Check if the storage_permission is inside the storage_permission collection"""
+        storage_permissions_or_err = self.storage_permissions
+        if storage_permissions_or_err.is_err():
+            return storage_permissions_or_err
+        storage_permissions_collection: MongoCollection = (
+            storage_permissions_or_err.ok()
+        )
+        storage_permissions: dict | None = storage_permissions_collection.find_one(
+            {"_id": permission.uid}
+        )
+
+        if storage_permissions is None or "node_uids" not in storage_permissions:
+            return False
+
+        return permission.node_uid in storage_permissions["node_uids"]
+
+    def remove_storage_permission(
+        self, storage_permission: StoragePermission
+    ) -> Result[None, Err]:
+        storage_permissions_or_err = self.storage_permissions
+        if storage_permissions_or_err.is_err():
+            return storage_permissions_or_err
+        storage_permissions_collection = storage_permissions_or_err.ok()
+
+        storage_permissions: dict | None = storage_permissions_collection.find_one(
+            {"_id": storage_permission.uid}
+        )
+        if storage_permissions is None:
+            return Err(
+                f"storage permission with UID {storage_permission.uid} not found!"
+            )
+        node_uids: set = storage_permissions["node_uids"]
+        if storage_permission.node_uid in node_uids:
+            node_uids.remove(storage_permission.node_uid)
+            storage_permissions_collection.update_one(
+                {"_id": storage_permission.uid},
+                {"$set": {"node_uids": node_uids}},
+            )
+        else:
+            return Err(
+                f"the node_uid {storage_permission.node_uid} does not exist in the storage permission!"
+            )
+
     def take_ownership(
         self, uid: UID, credentials: SyftVerifyKey
     ) -> Result[SyftSuccess, str]:

packages/syft/tests/syft/stores/mongo_document_store_test.py

@@ -11,6 +11,7 @@
 from syft.node.credentials import SyftVerifyKey
 from syft.service.action.action_permissions import ActionObjectPermission
 from syft.service.action.action_permissions import ActionPermission
+from syft.service.action.action_permissions import StoragePermission
 from syft.service.action.action_store import ActionObjectEXECUTE
 from syft.service.action.action_store import ActionObjectOWNER
 from syft.service.action.action_store import ActionObjectREAD
@@ -696,6 +697,41 @@ def test_mongo_store_partition_add_remove_permission(
     assert permissions_collection.count_documents({}) == 1
 
 
+def test_mongo_store_partition_add_remove_storage_permission(
+    root_verify_key: SyftVerifyKey,
+    mongo_store_partition: MongoStorePartition,
+) -> None:
+    """
+    Test the add_storage_permission and remove_storage_permission functions of MongoStorePartition
+    """
+
+    obj = MockSyftObject(data=1)
+
+    storage_permission = StoragePermission(
+        uid=obj.id,
+        node_uid=UID(),
+    )
+    assert not mongo_store_partition.has_storage_permission(storage_permission)
+    mongo_store_partition.add_storage_permission(storage_permission)
+    assert mongo_store_partition.has_storage_permission(storage_permission)
+    mongo_store_partition.remove_storage_permission(storage_permission)
+    assert not mongo_store_partition.has_storage_permission(storage_permission)
+
+    obj2 = MockSyftObject(data=1)
+    mongo_store_partition.set(root_verify_key, obj2, add_storage_permission=False)
+    storage_permission3 = StoragePermission(
+        uid=obj2.id, node_uid=mongo_store_partition.node_uid
+    )
+    assert not mongo_store_partition.has_storage_permission(storage_permission3)
+
+    obj3 = MockSyftObject(data=1)
+    mongo_store_partition.set(root_verify_key, obj3, add_storage_permission=True)
+    storage_permission4 = StoragePermission(
+        uid=obj3.id, node_uid=mongo_store_partition.node_uid
+    )
+    assert mongo_store_partition.has_storage_permission(storage_permission4)
+
+
 def test_mongo_store_partition_add_permissions(
     root_verify_key: SyftVerifyKey,
     guest_verify_key: SyftVerifyKey,