Add OrphanedObjectsState

ch_tools/chadmin/cli/object_storage_group.py

@@ -6,6 +6,7 @@
 from humanfriendly import format_size
 
 from ch_tools.chadmin.cli.chadmin_group import Chadmin
+from ch_tools.chadmin.internal.object_storage.orphaned_objects_state import OrphanedObjectsState
 from ch_tools.chadmin.internal.zookeeper import (
     check_zk_node,
     create_zk_nodes,
@@ -23,10 +24,6 @@
 # Use big enough timeout for stream HTTP query
 STREAM_TIMEOUT = 10 * 60
 
-# Orphaned objects state fields
-ORPHANED_OBJECTS_SIZE_FIELD = "orphaned_objects_size"
-ORPHANED_OBJECTS_ERROR_MSG_FIELD = "error_msg"
-
 STATE_LOCAL_PATH = "/tmp/object_storage_cleanup_state.json"
 
 
@@ -161,31 +158,29 @@ def clean_command(
     except Exception as e:
         error_msg = str(e)
 
+    state = OrphanedObjectsState(total_size, error_msg)
+
     if store_state_zk_path:
-        _store_state_zk_save(ctx, store_state_zk_path, total_size, error_msg)
+        _store_state_zk_save(ctx, store_state_zk_path, state)
 
     if store_state_local:
-        _store_state_local_save(ctx, total_size, error_msg)
+        _store_state_local_save(ctx, state)
 
     _print_response(ctx, dry_run, deleted, total_size)
 
 
 def _store_state_zk_save(
-    ctx: Context, path: str, total_size: int, error_msg: str
+    ctx: Context, path: str, state: OrphanedObjectsState
 ) -> None:
     if not check_zk_node(ctx, path):
         create_zk_nodes(ctx, [path], make_parents=True)
+    state_data = state.to_json().encode("utf-8")
+    update_zk_nodes(ctx, [path], state_data)
 
-    state_data = json.dumps(
-        create_orphaned_objects_state(total_size, error_msg), indent=4
-    )
-
-    update_zk_nodes(ctx, [path], state_data.encode("utf-8"))
 
-
-def _store_state_local_save(_: Context, total_size: int, error_msg: str) -> None:
+def _store_state_local_save(_: Context, state: OrphanedObjectsState) -> None:
     with open(STATE_LOCAL_PATH, "w", encoding="utf-8") as file:
-        json.dump(create_orphaned_objects_state(total_size, error_msg), file, indent=4)
+        file.write(state.to_json())
 
 
 def _print_response(ctx: Context, dry_run: bool, deleted: int, total_size: int) -> None:
@@ -211,10 +206,3 @@ def _table_formatter(stats):
     print_response(
         ctx, clean_stats, default_format="table", table_formatter=_table_formatter
     )
-
-
-def create_orphaned_objects_state(total_size: int, error_msg: str) -> dict:
-    return {
-        ORPHANED_OBJECTS_SIZE_FIELD: total_size,
-        ORPHANED_OBJECTS_ERROR_MSG_FIELD: error_msg,
-    }

ch_tools/chadmin/internal/object_storage/orphaned_objects_state.py

@@ -0,0 +1,19 @@
+from dataclasses import asdict, dataclass
+import json
+
+
+@dataclass
+class OrphanedObjectsState:
+    orphaned_objects_size: int
+    error_msg: str
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "OrphanedObjectsState":
+        data = json.load(json_str)
+        return cls(
+            orphaned_objects_size=data['orphaned_objects_size'],
+            error_msg=data['error_msg']
+        )
+
+    def to_json(self) -> str:
+        return json.dumps(asdict(self), indent=4)