From f0c4937c97345cb5a0fe77fb2d14373aa9ce4cb2 Mon Sep 17 00:00:00 2001
From: Kushal Deb <Kushal.Deb@ibm.com>
Date: Thu, 5 Dec 2024 16:27:34 +0530
Subject: [PATCH] cephadm: orch upgrade status to return output according to
 the format specified

Signed-off-by: Kushal Deb <Kushal.Deb@ibm.com>
---
 src/pybind/mgr/orchestrator/_interface.py | 33 +++++++++++++++++++++++
 src/pybind/mgr/orchestrator/module.py     | 28 ++++++++++---------
 2 files changed, 49 insertions(+), 12 deletions(-)

diff --git a/src/pybind/mgr/orchestrator/_interface.py b/src/pybind/mgr/orchestrator/_interface.py
index d2002f9fc0e7f..f7c54228ade6f 100644
--- a/src/pybind/mgr/orchestrator/_interface.py
+++ b/src/pybind/mgr/orchestrator/_interface.py
@@ -1069,6 +1069,39 @@ class UpgradeStatusSpec(object):
             'is_paused': self.is_paused,
         }
 
+    def to_dict(self) -> Dict:
+        out: Dict[str, Any] = {}
+        out['in_progress'] = self.in_progress
+        out['target_image'] = self.target_image
+        out['services_complete'] = self.services_complete
+        out['which'] = self.which
+        out['progress'] = self.progress
+        out['message'] = self.message
+        out['is_paused'] = self.is_paused
+        return out
+
+    @classmethod
+    def from_json(cls, data: dict) -> 'UpgradeStatusSpec':
+        if not isinstance(data, dict):
+            raise ValueError(f'Expected a dictionary, but got {type(data)}')
+        instance = cls()
+        instance.in_progress = data.get('in_progress', False)
+        instance.target_image = data.get('target_image', None)
+        instance.services_complete = data.get('services_complete', [])
+        instance.which = data.get('which', '<unknown>')
+        instance.progress = data.get('progress', None)
+        instance.message = data.get('message', "")
+        instance.is_paused = data.get('is_paused', False)
+
+        return instance
+
+    @staticmethod
+    def yaml_representer(dumper: 'yaml.Dumper', data: 'UpgradeStatusSpec') -> yaml.Node:
+        return dumper.represent_dict(cast(Mapping, data.to_json().items()))
+
+
+yaml.add_representer(UpgradeStatusSpec, UpgradeStatusSpec.yaml_representer)
+
 
 def handle_type_error(method: FuncT) -> FuncT:
     @wraps(method)
diff --git a/src/pybind/mgr/orchestrator/module.py b/src/pybind/mgr/orchestrator/module.py
index d3e278c33334b..42385ffca356f 100644
--- a/src/pybind/mgr/orchestrator/module.py
+++ b/src/pybind/mgr/orchestrator/module.py
@@ -49,6 +49,7 @@ from ._interface import (
     MgmtGatewaySpec,
     OAuth2ProxySpec,
     ServiceDescription,
+    UpgradeStatusSpec,
     TunedProfileSpec,
     _cli_read_command,
     _cli_write_command,
@@ -2367,21 +2368,24 @@ Usage:
         return HandleCommandResult(stdout=out)
 
     @_cli_write_command('orch upgrade status')
-    def _upgrade_status(self) -> HandleCommandResult:
+    def _upgrade_status(self, format: Optional[str] = None) -> HandleCommandResult:
         """Check the status of any potential ongoing upgrade operation"""
         completion = self.upgrade_status()
         status = raise_if_exception(completion)
-        r = {
-            'target_image': status.target_image,
-            'in_progress': status.in_progress,
-            'which': status.which,
-            'services_complete': status.services_complete,
-            'progress': status.progress,
-            'message': status.message,
-            'is_paused': status.is_paused,
-        }
-        out = json.dumps(r, indent=4)
-        if r.get('in_progress'):
+        upgrade_status = UpgradeStatusSpec()
+        upgrade_status.target_image = status.target_image
+        upgrade_status.in_progress = status.in_progress
+        upgrade_status.which = status.which
+        upgrade_status.services_complete = status.services_complete
+        upgrade_status.progress = status.progress
+        upgrade_status.message = status.message
+        upgrade_status.is_paused = status.is_paused
+        if format is not None:
+            format = Format(format)
+            out = to_format(upgrade_status, format, many=False, cls=UpgradeStatusSpec)
+        else:
+            out = json.dumps(upgrade_status.to_dict(), indent=4)
+        if upgrade_status.in_progress:
             return HandleCommandResult(stdout=out)
         return HandleCommandResult(stdout="There are no upgrades in progress currently.")
 
-- 
2.39.5