From 9bcd08c3753d1a0d55a78550d5e74a26af30b4b3 Mon Sep 17 00:00:00 2001
From: Kefu Chai <kchai@redhat.com>
Date: Sun, 17 Jan 2021 09:29:20 +0800
Subject: [PATCH] mgr/iostat: use CLICommand to define command

Signed-off-by: Kefu Chai <kchai@redhat.com>
---
 src/pybind/mgr/iostat/module.py | 79 ++++++++++++++-------------------
 1 file changed, 33 insertions(+), 46 deletions(-)

diff --git a/src/pybind/mgr/iostat/module.py b/src/pybind/mgr/iostat/module.py
index 215187377c2..5a765a151cd 100644
--- a/src/pybind/mgr/iostat/module.py
+++ b/src/pybind/mgr/iostat/module.py
@@ -1,18 +1,7 @@
-
-from mgr_module import MgrModule
+from mgr_module import CLIReadCommand, HandleCommandResult, MgrModule
 
 
 class Module(MgrModule):
-    COMMANDS = [
-        {
-            "cmd": "iostat",
-            "desc": "Get IO rates",
-            "perm": "r",
-            "poll": True
-        },
-    ]
-
-
     def __init__(self, *args, **kwargs):
         super(Module, self).__init__(*args, **kwargs)
 
@@ -26,7 +15,11 @@ class Module(MgrModule):
         assert('num_write' in r['pg_stats_delta']['stat_sum'])
         assert('num_read' in r['pg_stats_delta']['stat_sum'])
 
-    def handle_command(self, inbuf, command):
+    @CLIReadCommand('iostat', poll=True)
+    def iostat(self, width: int = 80, print_header: bool = False):
+        """
+        Get IO rates
+        """
         rd = 0
         wr = 0
         total = 0
@@ -35,39 +28,33 @@ class Module(MgrModule):
         total_ops = 0
         ret = ''
 
-        if command['prefix'] == 'iostat':
-            r = self.get('io_rate')
-
-            stamp_delta = float(r['pg_stats_delta']['stamp_delta'])
-            if (stamp_delta > 0):
-                rd = int(r['pg_stats_delta']['stat_sum']['num_read_kb']) / stamp_delta
-                wr = int(r['pg_stats_delta']['stat_sum']['num_write_kb']) / stamp_delta
-                # The values are in kB, but to_pretty_iec() requires them to be in bytes
-                rd = int(rd) << 10
-                wr = int(wr) << 10
-                total = rd + wr
-
-                rd_ops = int(r['pg_stats_delta']['stat_sum']['num_read']) / stamp_delta
-                wr_ops = int(r['pg_stats_delta']['stat_sum']['num_write']) / stamp_delta
-                total_ops = rd_ops + wr_ops
-
-            if 'width' in command:
-                width = command['width']
-            else:
-                width = 80
-
-            if command.get('print_header', False):
-                elems = ['Read', 'Write', 'Total', 'Read IOPS', 'Write IOPS', 'Total IOPS']
-                ret += self.get_pretty_header(elems, width)
+        r = self.get('io_rate')
 
-            elems = [
-                self.to_pretty_iec(rd) + 'B/s',
-                self.to_pretty_iec(wr) + 'B/s',
-                self.to_pretty_iec(total) + 'B/s',
-                int(rd_ops),
-                int(wr_ops),
-                int(total_ops)
+        stamp_delta = float(r['pg_stats_delta']['stamp_delta'])
+        if (stamp_delta > 0):
+            rd = int(r['pg_stats_delta']['stat_sum']['num_read_kb']) / stamp_delta
+            wr = int(r['pg_stats_delta']['stat_sum']['num_write_kb']) / stamp_delta
+            # The values are in kB, but to_pretty_iec() requires them to be in bytes
+            rd = int(rd) << 10
+            wr = int(wr) << 10
+            total = rd + wr
+
+            rd_ops = int(r['pg_stats_delta']['stat_sum']['num_read']) / stamp_delta
+            wr_ops = int(r['pg_stats_delta']['stat_sum']['num_write']) / stamp_delta
+            total_ops = rd_ops + wr_ops
+
+        if print_header:
+            elems = ['Read', 'Write', 'Total', 'Read IOPS', 'Write IOPS', 'Total IOPS']
+            ret += self.get_pretty_header(elems, width)
+
+        elems = [
+            self.to_pretty_iec(rd) + 'B/s',
+            self.to_pretty_iec(wr) + 'B/s',
+            self.to_pretty_iec(total) + 'B/s',
+            int(rd_ops),
+            int(wr_ops),
+            int(total_ops)
             ]
-            ret += self.get_pretty_row(elems, width)
+        ret += self.get_pretty_row(elems, width)
 
-        return 0, '', ret
+        return HandleCommandResult(stdout=ret)
-- 
2.39.5