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

Author: boonhapus

cs_tools/api/middlewares/tsload.py

@@ -1,7 +1,8 @@
 from typing import Any, Dict, List, Union
-from io import BufferedIOBase
+from io import BufferedIOBase, TextIOWrapper
 from tempfile import _TemporaryFileWrapper
 import logging
+import time
 
 from pydantic import validate_arguments
 
@@ -38,7 +39,7 @@ def _check_privileges(self) -> None:
     @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def upload(
         self,
-        fd: Union[BufferedIOBase, _TemporaryFileWrapper],
+        fd: Union[BufferedIOBase, TextIOWrapper, _TemporaryFileWrapper],
         *,
         database: str,
         table: str,
@@ -154,21 +155,29 @@ def upload(
         self.ts.api.ts_dataservice.load_commit(cycle_id)
         return cycle_id
 
-    # @validate_arguments
-    # def status(self, cycle_id: str, *, wait_for_complete: bool = False):
-    #     """
-    #     """
-    #     self._check_privileges()
+    @validate_arguments
+    def status(self, cycle_id: str, *, wait_for_complete: bool = False):
+        """
+        """
+        self._check_privileges()
+
+        while True:
+            r = self.ts.api.ts_dataservice.load_status(cycle_id)
+            data = r.json()
+
+            if not wait_for_complete:
+                break
 
-    #     while True:
-    #         r = self.ts.api.ts_dataservice.load_status(cycle_id)
-    #         data = r.json()
+            if data['internal_stage'] == 'COMMITTING':
+                pass
+            elif data['internal_stage'] == 'DONE':
+                break
+            elif data['status']['message'] != 'OK':
+                break
 
-    #         if not wait_for_complete:
-    #             break
+            time.sleep(1)
 
-    #         print(data)
-    #         raise
+        return data
 
     # @validate_arguments
     # def bad_records(self, cycle_id: str) -> List[Dict[str, Any]]:

cs_tools/cli/tools/rtsload/app.py

@@ -37,13 +37,13 @@
 def status(
     ctx: typer.Context,
     cycle_id: str=A_(..., help='data load cycle id'),
-    bad_records: str = O_(
-        None,
-        '--bad_records_file',
-        help='file to use for storing rows that failed to load',
-        metavar='protocol://DEFINITION.toml',
-        callback=lambda ctx, to: SyncerProtocolType().convert(to, ctx=ctx)
-    )
+    # bad_records: str = O_(
+    #     None,
+    #     '--bad_records_file',
+    #     help='file to use for storing rows that failed to load',
+    #     metavar='protocol://DEFINITION.toml',
+    #     callback=lambda ctx, to: SyncerProtocolType().convert(to, ctx=ctx)
+    # )
 ):
     """
     Get the status of a data load.
@@ -59,18 +59,21 @@ def status(
             f'\nIgnored rows: {data["ignored_row_count"]}'
         )
 
-    if data['ignored_row_count'] > 0:
-        now = dt.datetime.now().strftime('%Y-%m-%dT%H_%M_%S')
-        fp = f'BAD_RECORDS_{now}_{cycle_id}'
-        console.print(
-            f'[red]\n\nBad records found...\n\twriting to {bad_records.directory / fp}'
-        )
-        data = ts.tsload.bad_records(cycle_id)
-        bad_records.dump(fp, data=data)
+    # if int(data['ignored_row_count']) > 0:
+    #     now = dt.datetime.now().strftime('%Y-%m-%dT%H_%M_%S')
+    #     fp = f'BAD_RECORDS_{now}_{cycle_id}'
+    #     console.print(
+    #         f'[red]\n\nBad records found...\n\twriting to {bad_records.directory / fp}'
+    #     )
+    #     data = ts.tsload.bad_records(cycle_id)
+    #     bad_records.dump(fp, data=data)
 
     if data['status']['code'] == 'LOAD_FAILED':
         console.print(f'\nFailure reason:\n  [red]{data["status"]["message"]}[/]')
 
+    if data.get('parsing_errors', False):
+        console.print(f'[red]{data["parsing_errors"]}')
+
 
 @app.command(cls=CSToolsCommand)
 @depends(
@@ -96,13 +99,13 @@ def file(
     has_header_row: bool = O_(False, '--has_header_row', show_default=False, help='indicates that the input file contains a header row'),
     escape_character: str = O_('"', '--escape_character', help='specifies the escape character used in the input file'),
     enclosing_character: str = O_('"', '--enclosing_character', help='enclosing character in csv source format'),
-    bad_records: str = O_(
-        None,
-        '--bad_records_file',
-        help='file to use for storing rows that failed to load',
-        metavar='protocol://DEFINITION.toml',
-        callback=lambda ctx, to: SyncerProtocolType().convert(to, ctx=ctx)
-    ),
+    # bad_records: str = O_(
+    #     None,
+    #     '--bad_records_file',
+    #     help='file to use for storing rows that failed to load',
+    #     metavar='protocol://DEFINITION.toml',
+    #     callback=lambda ctx, to: SyncerProtocolType().convert(to, ctx=ctx)
+    # ),
     flexible: bool = O_(False, '--flexible', show_default=False, help='whether input data file exactly matches target schema', hidden=True)
 ):
     """
@@ -120,9 +123,9 @@ def file(
     #
 
     opts = {
-        'target_database': target_database,
-        'target_table': target_table,
-        'target_schema': target_schema,
+        'database': target_database,
+        'table': target_table,
+        'schema_': target_schema,
         'empty_target': empty_target,
         'max_ignored_rows': max_ignored_rows,
         'date_format': date_format,
@@ -138,10 +141,11 @@ def file(
         'enclosing_character': enclosing_character
     }
 
-    with console.status(f'[bold green]Loading {file} to ThoughtSpot..'):
-        cycle_id = ts.tsload.upload(ts, fp=file, **opts, verbose=True)
+    with console.status(f'[bold green]Loading [yellow]{file}[/] to ThoughtSpot..'):
+        with file.open('r', encoding='utf-8', newline='') as fd:
+            cycle_id = ts.tsload.upload(fd, **opts)
 
-    console.log(f'Data load cycle_id: {cycle_id}')
+    console.log(f'Data load cycle_id: [cyan]{cycle_id}')
 
     # if bad_records is None:
     #     return

cs_tools/sync/falcon/syncer.py

@@ -24,12 +24,14 @@ class Falcon:
     database: str = 'cs_tools'
     schema_: str = Field('falcon_default_schema', alias='schema')
     empty_target: bool = True
+    timeout: float = 60.0
 
     # DATABASE ATTRIBUTES
     __is_database__ = True
     metadata = None
 
     def __post_init_post_parse__(self):
+        self.timeout = self.timeout or None
         ctx = click.get_current_context()
         self.engine = sa.engine.create_mock_engine('sqlite://', self.intercept_create_table)
         self.cnxn = self.engine.connect()
@@ -51,7 +53,11 @@ def intercept_create_table(self, sql, *multiparams, **params):
         if 'ts_bi_server' in q:
             return
 
-        self.ts.tql.command(command=f'{q};', database=self.database)
+        self.ts.tql.command(
+            command=f'{q};',
+            database=self.database,
+            http_timeout=self.timeout
+        )
 
     def ensure_setup(self, metadata, cnxn, **kw):
 
@@ -69,8 +75,14 @@ def ensure_setup(self, metadata, cnxn, **kw):
             )
 
         # 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_};')
+        self.ts.tql.command(
+            command=f'CREATE DATABASE {self.database};', http_timeout=self.timeout
+        )
+
+        self.ts.tql.command(
+            command=f'CREATE SCHEMA {self.database}.{self.schema_};',
+            http_timeout=self.timeout
+        )
 
     def capture_metadata(self, metadata, cnxn, **kw):
         self.metadata = metadata
@@ -88,7 +100,11 @@ def load(self, table: str) -> List[Dict[str, Any]]:
         t = self.metadata.tables[table]
         q = t.select().compile(dialect=self.engine.dialect)
         q = str(q).strip()
-        r = self.ts.tql.query(statement=f'{q};', database=self.database)
+        r = self.ts.tql.query(
+            statement=f'{q};',
+            database=self.database,
+            http_timeout=self.timeout
+        )
         d = next(_['data'] for _ in r if 'data' in _)  # there will be only 1 response
         return d
 

docs/syncer/falcon.md

@@ -36,6 +36,10 @@ Falcon is __ThoughtSpot__'s proprietary in-memory database that exists as part o
 <br/>*<span class=fc-mint>default</span>:* `true`
 <br/>*a* `TRUNCATE` *statement will be issued prior to loading any data loads if* `true` *is used*
 
+> __timeout__{ .fc-blue }: <span class=fc-coral>optional</span>, number of seconds to wait for individual commands to Falcon
+<br/>*<span class=fc-mint>default</span>:* `60.0`
+<br/>*if the timeout value is set to* `0` *then calls over to Falcon will never reach a timeout threshold*
+
 
 ??? question "How do I use the Falcon syncer in commands?"
 
@@ -62,6 +66,7 @@ Falcon is __ThoughtSpot__'s proprietary in-memory database that exists as part o
 database = 'cs_tools'
 schema = 'falcon_default_schema'
 empty_target = true
+timeout = 5.0
 ```
 
 [gh-issue25]: https://github.com/thoughtspot/cs_tools/issues/25