refactor: BI-5758 add USM resave conditionally after post_save_hook (#…

…596)

lib/dl_connector_bundle_chs3/dl_connector_bundle_chs3/chs3_base/core/lifecycle.py

@@ -8,6 +8,7 @@
 
 from dl_api_commons.base_models import RequestContextInfo
 from dl_core.connectors.base.lifecycle import ConnectionLifecycleManager
+from dl_core.lifecycle.base import PostSaveHookResult
 from dl_file_uploader_task_interface.tasks import (
     DeleteFileTask,
     SaveSourceTask,
@@ -94,14 +95,16 @@ def get_task_processor(self, req_id: Optional[str] = None) -> TaskProcessor:
         task_processor = task_processor_factory.make(req_id)
         return task_processor
 
-    def post_save_hook(self) -> None:
+    def post_save_hook(self) -> PostSaveHookResult:
         super().post_save_hook()
 
         rci = self._us_manager.bi_context
         task_processor = self.get_task_processor(rci.request_id)
         scheduler = FileConnTaskScheduler(task_processor=task_processor, rci=rci)
         scheduler.schedule_sources_update(self.entry)
 
+        return PostSaveHookResult()
+
     def post_delete_hook(self) -> None:
         super().post_delete_hook()
 

lib/dl_core/dl_core/lifecycle/base.py

@@ -22,6 +22,16 @@
 _US_ENTRY_TV = TypeVar("_US_ENTRY_TV", bound=USEntry)
 
 
+@attr.s
+class HookResult:
+    ...
+
+
+@attr.s
+class PostSaveHookResult(HookResult):
+    additional_save_needed: bool = False
+
+
 @attr.s
 class EntryLifecycleManager(abc.ABC, Generic[_US_ENTRY_TV]):
     ENTRY_CLS: ClassVar[Type[_US_ENTRY_TV]]  # type: ignore  # 2024-01-24 # TODO: ClassVar cannot contain type variables  [misc]
@@ -38,8 +48,8 @@ def entry(self) -> _US_ENTRY_TV:
     def pre_save_hook(self) -> None:
         pass
 
-    def post_save_hook(self) -> None:
-        pass
+    def post_save_hook(self) -> PostSaveHookResult:
+        return PostSaveHookResult()
 
     def post_copy_hook(self) -> None:
         pass

lib/dl_core/dl_core/us_entry.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import logging
+import typing
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -50,9 +51,9 @@ class USEntry:
     hidden: bool
     links: Optional[dict] = None
 
-    _stored_in_db = False
+    _stored_in_db: bool = False
     _us_resp: Optional[dict] = None
-    _lock = None
+    _lock: typing.Optional[str] = None
     _us_manager: Optional[USManagerBase]
 
     @classmethod
@@ -141,6 +142,22 @@ def us_manager(self) -> USManagerBase:
         assert self._us_manager is not None
         return self._us_manager
 
+    @property
+    def stored_in_db(self) -> bool:
+        return self._stored_in_db
+
+    @stored_in_db.setter
+    def stored_in_db(self, value: bool) -> None:
+        self._stored_in_db = value
+
+    @property
+    def lock(self) -> typing.Optional[str]:
+        return self._lock
+
+    @lock.setter
+    def lock(self, value: typing.Optional[str]) -> None:
+        self._lock = value
+
     @property
     def _context(self) -> RequestContextInfo:
         return self._us_manager.bi_context  # type: ignore  # TODO: fix

lib/dl_core/dl_core/us_manager/us_manager.py

@@ -354,7 +354,7 @@ def _entry_dict_to_obj(self, us_resp: dict, expected_type: Optional[Type[USEntry
                 decrypted_data = self._crypto_controller.decrypt(json.loads(old_data)) if old_data is not None else None
                 serializer.set_data_attr(entry, key, decrypted_data)
 
-        entry._stored_in_db = True
+        entry.stored_in_db = True
         entry._us_resp = us_resp
 
         return entry
@@ -411,7 +411,7 @@ def _get_entry_save_params(self, entry: USEntry) -> dict:
 
         save_params: dict[str, Any] = {}
 
-        if not entry._stored_in_db:
+        if not entry.stored_in_db:
             if entry.permissions_mode is not None:
                 save_params.update(permissionsMode=entry.permissions_mode)
             if entry.initial_permissions is not None:
@@ -449,6 +449,17 @@ def _get_entry_save_params(self, entry: USEntry) -> dict:
 
         return save_params
 
+    def _prepare_update_entry_params(self, entry: USEntry, update_revision: Optional[bool] = None) -> dict:
+        assert entry.uuid is not None
+        save_params = self._get_entry_save_params(entry)
+        assert "data" in save_params and "unversioned_data" in save_params
+
+        save_params.pop("scope")
+        save_params.pop("type")
+        save_params["update_revision"] = update_revision
+
+        return save_params
+
     def copy_entry(self, source: _ENTRY_TV, key: Optional[EntryLocation] = None) -> _ENTRY_TV:
         if not isinstance(source, Dataset):
             raise ValueError("Only dataset can be copied at this time")

lib/dl_core/dl_core/us_manager/us_manager_async.py

@@ -132,15 +132,15 @@ async def get_by_id(self, entry_id: str, expected_type: Type[_ENTRY_TV] = None)
         return obj
 
     async def save(self, entry: USEntry, update_revision: Optional[bool] = None) -> None:
-        self.get_lifecycle_manager(entry=entry).pre_save_hook()
+        lifecycle_manager = self.get_lifecycle_manager(entry=entry)
+        lifecycle_manager.pre_save_hook()
 
         save_params = self._get_entry_save_params(entry)
         us_scope = save_params.pop("scope")
         us_type = save_params.pop("type")
         assert "data" in save_params and "unversioned_data" in save_params
 
-        # noinspection PyProtectedMember
-        if not entry._stored_in_db:
+        if not entry.stored_in_db:
             entry_loc = entry.entry_key
             assert entry_loc is not None, "Entry location must be set before saving US entry"
 
@@ -151,7 +151,7 @@ async def save(self, entry: USEntry, update_revision: Optional[bool] = None) ->
                 **save_params,
             )
             entry.uuid = resp["entryId"]
-            entry._stored_in_db = True
+            entry.stored_in_db = True
         else:
             # noinspection PyProtectedMember
             save_params["update_revision"] = update_revision
@@ -169,7 +169,7 @@ async def delete(self, entry: USEntry) -> None:
 
         # noinspection PyProtectedMember
         await self._us_client.delete_entry(entry.uuid, lock=entry._lock)
-        entry._stored_in_db = False
+        entry.stored_in_db = False
 
         # noinspection PyBroadException
         try:

lib/dl_core/dl_core/us_manager/us_manager_sync.py

@@ -129,32 +129,34 @@ def save(self, entry: USEntry, update_revision: Optional[bool] = None) -> None:
         assert "data" in save_params and "unversioned_data" in save_params
         assert entry_loc is not None, "Can not save entry without key/workbook data"
 
-        # noinspection PyProtectedMember
-        if not entry._stored_in_db:
+        if not entry.stored_in_db:
             resp = self._us_client.create_entry(
                 key=entry_loc,
                 scope=us_scope,
                 type_=us_type,
                 **save_params,
             )
             entry.uuid = resp["entryId"]
-            entry._stored_in_db = True
+            entry.stored_in_db = True
         else:
-            # noinspection PyProtectedMember
             save_params["update_revision"] = update_revision
-            resp = self._us_client.update_entry(
-                entry.uuid, lock=entry._lock, **save_params  # type: ignore  # TODO: fix
-            )
+            assert entry.uuid is not None
+            resp = self._us_client.update_entry(entry.uuid, lock=entry.lock, **save_params)
 
         entry._us_resp = resp  # type: ignore  # TODO: fix
-        lifecycle_manager.post_save_hook()
+
+        post_save_result = lifecycle_manager.post_save_hook()
+        if post_save_result.additional_save_needed:
+            save_params = self._prepare_update_entry_params(entry, False)
+            assert entry.uuid is not None
+            entry._us_resp = self._us_client.update_entry(entry.uuid, lock=entry.lock, **save_params)
 
     def delete(self, entry: USEntry) -> None:
         lifecycle_manager = self.get_lifecycle_manager(entry=entry)
         lifecycle_manager.pre_delete_hook()
 
-        self._us_client.delete_entry(entry.uuid, lock=entry._lock)
-        entry._stored_in_db = False
+        self._us_client.delete_entry(entry.uuid, lock=entry.lock)
+        entry.stored_in_db = False
         # noinspection PyBroadException
         try:
             LOGGER.info("Executing post-delete hook %s", entry.uuid)
@@ -266,14 +268,13 @@ def reload_data(self, entry: USEntry) -> None:
     #
     def acquire_lock(self, entry: USEntry, duration: Optional[int] = None, wait_timeout: Optional[int] = None) -> str:
         lock_token = self._us_client.acquire_lock(entry.uuid, duration, wait_timeout)
-        entry._lock = lock_token
+        entry.lock = lock_token
         return lock_token
 
     def release_lock(self, entry: USEntry) -> None:
-        # noinspection PyProtectedMember
-        if entry._lock:
-            self._us_client.release_lock(entry.uuid, entry._lock)
-            entry._lock = None
+        if entry.lock:
+            self._us_client.release_lock(entry.uuid, entry.lock)
+            entry.lock = None
 
     @contextlib.contextmanager  # type: ignore  # TODO: fix
     def locked_cm(  # type: ignore  # TODO: fix
@@ -297,12 +298,12 @@ def get_locked_entry_cm(  # type: ignore  # TODO: fix
         entry = None
         try:
             entry = self.get_by_id(entry_id, expected_type=expected_type)
-            entry._lock = lock_token
+            entry.lock = lock_token
             yield entry
         finally:
             self._us_client.release_lock(entry_id, lock_token)
             if entry is not None:
-                entry._lock = None
+                entry.lock = None
 
     # Dependencies
     #