Merge branch 'v1.3.0-hot-fixes' into dev

boonhapus · boonhapus · commit 5f35e1bec200 · 2022-05-18T21:46:44.000-05:00
diff --git a/cs_tools/api/middlewares/tql.py b/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
diff --git a/cs_tools/cli/app_config.py b/cs_tools/cli/app_config.py
@@ -41,7 +41,7 @@ def show(
     configs = [f for f in APP_DIR.iterdir() if f.name.startswith('cluster-cfg_')]
     s = 's' if (len(configs) > 1 or len(configs) == 0) else ''
     meta = _meta_config()
-    meta_cfg = None if meta is None else meta['default']['config']
+    meta_cfg = meta['default']['config'] if meta else {}
 
     console.print(f'\nCluster configs located at: {APP_DIR}\n')
 
diff --git a/cs_tools/cli/tools/rtql/app.py b/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)
diff --git a/cs_tools/cli/tools/rtql/interactive.py b/cs_tools/cli/tools/rtql/interactive.py
@@ -3,6 +3,7 @@
 import json
 
 from rich.console import Console
+import rich
 import httpx
 import typer
 
@@ -15,6 +16,16 @@
 log = logging.getLogger(__name__)
 
 
+def _to_table(headers, rows=None):
+    header = [column['name'] for column in headers]
+
+    if rows is None:
+        rows = [{'v': (' ',) * len(header)}]
+
+    data = [dict(zip(header, row['v'])) for row in rows]
+    return data
+
+
 class InteractiveTQL:
     """
     An interactive TQL client.
@@ -150,6 +161,11 @@ def _handle_query(self, lines: List[str]) -> None:
         """
 
         """
+        color_map = {
+            'INFO': '[white]',
+            'WARNING': '[yellow]',
+            'ERROR': '[red]'
+        }
         new_ctx = {}
 
         for line in lines:
@@ -172,12 +188,22 @@ def _handle_query(self, lines: List[str]) -> None:
                 continue
 
             if 'message' in data['result']:
-                msg = self.ts.api.ts_dataservice._parse_api_messages(data['result']['message'])
-                self.print(msg)
+                for message in data['result']['message']:
+                    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]'
+
+                    self.print(c + m + '[/]', end='')
 
             if 'table' in data['result']:
-                msg = self.ts.api.ts_dataservice._parse_tql_query(data['result']['table'])
-                self.print(msg)
+                d = _to_table(**data['result']['table'])
+                t = rich.table.Table(*d[0].keys(), box=rich.box.HORIZONTALS)
+                [t.add_row(*_.values()) for _ in d]
+                self.print(t)
 
         return new_ctx
 
diff --git a/cs_tools/cli/ux.py b/cs_tools/cli/ux.py
@@ -125,15 +125,21 @@ def convert(
         if validate_only:
             return value
 
-        if getattr(syncer, '__is_database__', False):
+        is_database_check = getattr(syncer, '__is_database__', False)
+        is_tools_cmd = 'tools' in sys.argv[1:]
+
+        if is_database_check or not is_tools_cmd:
             if getattr(syncer, 'metadata') is not None:
                 metadata = syncer.metadata
                 [t.to_metadata(metadata) for t in models.SQLModel.metadata.sorted_tables]
             else:
                 metadata = models.SQLModel.metadata
 
             metadata.create_all(syncer.cnxn)
-            metadata.reflect(syncer.cnxn, views=True)
+
+            # DEV NOTE: conditionally expose ability to grab views
+            if syncer.name != 'falcon':
+                metadata.reflect(syncer.cnxn, views=True)
 
         return syncer
 
diff --git a/cs_tools/sync/csv/syncer.py b/cs_tools/sync/csv/syncer.py
@@ -49,8 +49,9 @@ def file_reference(self, file: str, mode: str) -> TextIO:
         Handle open-close on a file, potentially in a zip archive.
         """
         file_opts = {
-            'newline': '',
-            'mode': 'r' if mode == 'r' else 'w'
+            'mode': 'r' if mode == 'r' else 'w',
+            'encoding': 'utf-8',
+            'newline': ''
         }
 
         if self.zipped:
diff --git a/cs_tools/sync/falcon/syncer.py b/cs_tools/sync/falcon/syncer.py
@@ -9,7 +9,7 @@
 import httpx
 import click
 
-from cs_tools.errors import SyncerError, TSLoadServiceUnreachable
+from cs_tools.errors import SyncerError, ThoughtSpotUnreachable, TSLoadServiceUnreachable
 from . import compiler, sanitize
 
 
@@ -33,13 +33,7 @@ def __post_init_post_parse__(self):
         ctx = click.get_current_context()
         self.engine = sa.engine.create_mock_engine('sqlite://', self.intercept_create_table)
         self.cnxn = self.engine.connect()
-        self._thoughtspot = getattr(ctx.obj, 'thoughtspot', '')
-
-        if self._thoughtspot.platform.deployment == 'cloud':
-            raise SyncerError(
-                'Falcon is not available for data load operations on TS Cloud '
-                'deployments'
-            )
+        self._thoughtspot = getattr(ctx.obj, 'thoughtspot', None)
 
         # decorators must be declared here, SQLAlchemy doesn't care about instances
         sa.event.listen(sa.schema.MetaData, 'before_create', self.ensure_setup)
@@ -60,6 +54,20 @@ def intercept_create_table(self, sql, *multiparams, **params):
         self.ts.tql.command(command=f'{q};', database=self.database)
 
     def ensure_setup(self, metadata, cnxn, **kw):
+
+        if self.ts is None:
+            # DEV NOTE:
+            # I think we can realistically only reach here if Falcon is meant to be
+            # active AND we are attempting to run a tools command, so that's not the 
+            # case, @boonhapus has gotta take a better look.
+            raise ThoughtSpotUnreachable('unknown reason')
+
+        if self.ts.platform.deployment == 'cloud':
+            raise SyncerError(
+                'Falcon is not available for data load operations on TS Cloud '
+                'deployments'
+            )
+
         # create the database and schema if it doesn't exist
         self.ts.tql.command(command=f'CREATE DATABASE {self.database};')
         self.ts.tql.command(command=f'CREATE SCHEMA {self.database}.{self.schema_};')
@@ -81,7 +89,8 @@ def load(self, table: str) -> List[Dict[str, Any]]:
         q = t.select().compile(dialect=self.engine.dialect)
         q = str(q).strip()
         r = self.ts.tql.query(statement=f'{q};', database=self.database)
-        return r
+        d = next(_['data'] for _ in r if 'data' in _)  # there will be only 1 response
+        return d
 
     def dump(self, table: str, *, data: List[Dict[str, Any]]) -> None:
         if not data:
