CLI: Add support for max_rows, display all columns wraparound

Signed-off-by: Dinesh Dutt <[email protected]>

Author: ddutt

suzieq/cli/sqcmds/command.py

@@ -274,6 +274,7 @@ def _pager_print(self, df: pd.DataFrame) -> None:
         else:
             print(df)
 
+    # pylint: disable=too-many-statements
     def _gen_output(self, df: pd.DataFrame, json_orient: str = "records",
                     dont_strip_cols: bool = False, sort: bool = True):
 
@@ -291,6 +292,9 @@ def _gen_output(self, df: pd.DataFrame, json_orient: str = "records",
                 cols = df.columns.tolist()
             is_error = False
 
+        max_rows = self.ctxt.max_rows
+        all_columns = self.ctxt.all_columns
+
         if dont_strip_cols or not all(item in df.columns for item in cols):
             cols = df.columns.tolist()
 
@@ -317,9 +321,11 @@ def _gen_output(self, df: pd.DataFrame, json_orient: str = "records",
                 'timestamp' not in self.columns and not dont_strip_cols and
                     'timestamp' in df.columns and 'timestamp' in cols):
                 cols.remove('timestamp')
-            with pd.option_context('precision', 3,
-                                   'display.max_colwidth', max_colwidth,
-                                   'display.max_rows', 256):
+            with pd.option_context(
+                    'precision', 3,
+                    'display.max_colwidth', max_colwidth,
+                    'display.max_rows', max_rows,
+                    'display.expand_frame_repr', not all_columns):
                 df = self.sqobj.humanize_fields(df)
                 if df.empty:
                     print(df)

suzieq/cli/sqcmds/context_commands.py

@@ -22,6 +22,11 @@
 @argument("pager", description="Enable pagination prompt on longer outputs",
           choices=['on', 'off'])
 @argument('col_width', description='Max Width of each column in table display')
+@argument('max_rows',
+          description='Max rows to display before ellipsis is shown')
+@argument('all_columns',
+          description='Display all columns wrapping around if necessary',
+          choices=['yes', 'no'])
 @argument(
     "engine",
     choices=SUPPORTED_ENGINES,
@@ -57,6 +62,8 @@ def set_ctxt(
         datadir: str = "",
         debug: str = '',
         col_width: int = 50,
+        max_rows: int = None,
+        all_columns: str = None,
         rest_server_ip: str = "",
         rest_server_port: str = "",
         rest_api_key: str = "",
@@ -93,6 +100,15 @@ def set_ctxt(
     if col_width:
         ctxt.col_width = int(col_width)
 
+    if max_rows is not None:
+        if max_rows == 0:
+            ctxt.max_rows = None
+        else:
+            ctxt.max_rows = max_rows
+
+    if all_columns is not None:
+        ctxt.all_columns = all_columns == "yes"
+
     if pager == 'on':
         ctxt.pager = True
     elif pager == 'off':

suzieq/shared/context.py

@@ -18,6 +18,8 @@ class SqContext:
     exec_time: str = ''
     engine: str = None
     col_width: int = 50
+    max_rows: int = 256
+    all_columns: bool = False
     debug: bool = False
     sort_fields: List[str] = field(default_factory=list)
     view: str = None