🐛 implement proper data response

Author: boonhapus

cs_tools/api/middlewares/tql.py

@@ -17,6 +17,15 @@
 ])
 
 
+def _to_table(headers, rows=None):
+    if rows is None:
+        rows = []
+
+    header = [column['name'] for column in headers]
+    data = [dict(zip(header, row['v'])) for row in rows]
+    return data
+
+
 class TQLMiddleware:
     """
     """
@@ -73,25 +82,36 @@ def query(
         }
 
         r = self.ts.api.ts_dataservice.query(data, timeout=http_timeout)
-        d = [json.loads(_) for _ in r.iter_lines() if _]
-        print(d)
-        raise
-        return d['result']['message']
+        i = [json.loads(_) for _ in r.iter_lines() if _]
+
+        out = []
+
+        for row in i:
+            if 'table' in row['result']:
+                out.append({'data': _to_table(**row['result']['table'])})
+
+            if 'message' in row['result']:
+                out.append({'messages': row['result']['message']})
+
+        return out
 
     @validate_arguments
     def command(
         self,
         command: str,
         *,
         database: str = None,
-        schema_: Annotated[str, Field(alias='schema')] = 'falcon_default_schema',
+        schema_: str = 'falcon_default_schema',
         raise_errors: bool = False,
         http_timeout: int = 5.0
     ) -> List[Dict[str, Any]]:
         """
         """
         self._check_privileges()
 
+        if not command.strip().endswith(';'):
+            command = f'{command.strip()};'
+
         data = {
             'context': {
                 'database': database,
@@ -104,11 +124,18 @@ def command(
         }
 
         r = self.ts.api.ts_dataservice.query(data, timeout=http_timeout)
+        i = [json.loads(_) for _ in r.iter_lines() if _]
+
+        out = []
+
+        for row in i:
+            if 'table' in row['result']:
+                out.append({'data': _to_table(**row['result']['table'])})
+
+            if 'message' in row['result']:
+                out.append({'messages': row['result']['message']})
 
-        d = [json.loads(_) for _ in r.iter_lines() if _]
-        d = r.json()
-        log.debug(d)
-        return d['result']['message']
+        return out
 
     @validate_arguments
     def script(
@@ -133,18 +160,15 @@ def script(
             }
 
         r = self.ts.api.ts_dataservice.script(data, timeout=http_timeout)
-        d = [json.loads(_) for _ in r.iter_lines() if _]
-
-        for _ in d:
-            if 'message' in _['result']:
-                m = self.ts.api.ts_dataservice._parse_api_messages(_['result']['message'])
-            if 'table' in _['result']:
-                m = self.ts.api.ts_dataservice._parse_tql_query(_['result']['table'])
-            
-            log.debug(m)
-
-            if raise_errors and 'returned error' in m:
-                if 'create table' in m.lower():
-                    raise TableAlreadyExists()
-                else:
-                    raise ValueError(m)
+        i = [json.loads(_) for _ in r.iter_lines() if _]
+
+        out = []
+
+        for row in i:
+            if 'table' in row['result']:
+                out.append({'data': _to_table(**row['result']['table'])})
+
+            if 'message' in row['result']:
+                out.append({'messages': row['result']['message']})
+
+        return out

cs_tools/cli/tools/rtql/app.py

@@ -3,6 +3,7 @@
 
 from typer import Argument as A_, Option as O_
 import typer
+import rich
 
 from cs_tools.cli.dependency import depends
 from cs_tools.cli.options import CONFIG_OPT, VERBOSE_OPT, TEMP_DIR_OPT
@@ -78,7 +79,30 @@ def file(
             cat create-schema.sql | tql
     """
     ts = ctx.obj.thoughtspot
-    ts.tql.script(file)
+    r = ts.tql.script(file)
+
+    color_map = {
+        'INFO': '[white]',
+        'ERROR': '[red]'
+    }
+
+    for response in r:
+        if 'messages' in response:
+            for message in response['messages']:
+                c = color_map.get(message['type'], '[yellow]')
+                m = message['value']
+
+                if m.strip() == 'Statement executed successfully.':
+                    c = '[bold green]'
+                if m.strip().endswith(';'):
+                    c = '[cyan]'
+
+                console.print(c + m, end='')
+
+        if 'data' in response:
+            t = rich.table.Table(*response['data'][0].keys(), box=rich.box.HORIZONTALS)
+            [t.add_row(*_.values()) for _ in response['data']]
+            console.print('\n', t)
 
 
 @app.command(cls=CSToolsCommand)
@@ -116,4 +140,27 @@ def command(
         console.print('[red]no valid input given to rtql command')
         raise typer.Exit()
 
-    ts.tql.command(command, schema=schema)
+    r = ts.tql.command(command, schema_=schema)
+
+    color_map = {
+        'INFO': '[white]',
+        'ERROR': '[red]'
+    }
+
+    for response in r:
+        if 'messages' in response:
+            for message in response['messages']:
+                c = color_map.get(message['type'], '[yellow]')
+                m = message['value']
+
+                if m.strip() == 'Statement executed successfully.':
+                    c = '[bold green]'
+                if m.strip().endswith(';'):
+                    c = '[cyan]'
+
+                console.print(c + m, end='')
+
+        if 'data' in response:
+            t = rich.table.Table(*response['data'][0].keys(), box=rich.box.HORIZONTALS)
+            [t.add_row(*_.values()) for _ in response['data']]
+            console.print('\n', t)