Merge branch 'dev'

Author: boonhapus

cs_tools/_version.py

@@ -1 +1 @@
-__version__ = '1.3.1'
+__version__ = '1.3.2'

cs_tools/api/_rest_api_v1.py

@@ -34,7 +34,7 @@ def __init__(self, config, ts):
         self._http = httpx.Client(
             headers={'X-Requested-By': 'ThoughtSpot'},
             verify=not config.thoughtspot.disable_ssl,
-            timeout=180.0,
+            timeout=5 * 60.0,
             base_url=config.thoughtspot.fullpath
         )
 

cs_tools/api/middlewares/answer.py

@@ -73,11 +73,19 @@ def all(
             answers.extend(to_extend)
 
             if not to_extend and not answers:
-                rzn  = f"'{category.value}' category ("
-                rzn += 'excluding ' if exclude_system_content else 'including '
-                rzn += 'admin-generated answers)'
-                rzn += '' if tags is None else ' and tags: ' + ', '.join(tags)
-                raise ContentDoesNotExist(type='ANSWER', reason=rzn)
+                info = {
+                    "incl": "exclude" if exclude_system_content else "include",
+                    "category": category,
+                    "tags": ", ".join(tags),
+                    "reason": (
+                        "Zero {type} matched the following filters"
+                        "\n"
+                        "\n  - [blue]{category.value}[/] {type}"
+                        "\n  - [blue]{incl}[/] admin-generated {type}"
+                        "\n  - with tags [blue]{tags}"
+                    )
+                }
+                raise ContentDoesNotExist(type="answers", **info)
 
             if data['isLastBatch']:
                 break

cs_tools/api/middlewares/metadata.py

@@ -97,11 +97,19 @@ def all(
                     break
 
         if not content:
-            rzn  = f"'{category.value}' category ("
-            rzn += 'excluding ' if exclude_system_content else 'including '
-            rzn += 'admin-generated content)'
-            rzn += '' if tags is None else ' and tags: ' + ', '.join(tags)
-            raise ContentDoesNotExist(type=content, reason=rzn)
+            info = {
+                "incl": "exclude" if exclude_system_content else "include",
+                "category": category,
+                "tags": ", ".join(tags),
+                "reason": (
+                    "Zero {type} matched the following filters"
+                    "\n"
+                    "\n  - [blue]{category.value}[/] {type}"
+                    "\n  - [blue]{incl}[/] admin-generated {type}"
+                    "\n  - with tags [blue]{tags}"
+                )
+            }
+            raise ContentDoesNotExist(type="content", **info)
 
         return content
 

cs_tools/api/middlewares/pinboard.py

@@ -64,11 +64,19 @@ def all(
             offset += len(data['headers'])
 
             if not data['headers'] and not pinboards:
-                rzn  = f"'{category.value}' category ("
-                rzn += 'excluding ' if exclude_system_content else 'including '
-                rzn += 'admin-generated pinboards)'
-                rzn += '' if tags is None else ' and tags: ' + ', '.join(tags)
-                raise ContentDoesNotExist(type='PINBOARD', reason=rzn)
+                info = {
+                    "incl": "exclude" if exclude_system_content else "include",
+                    "category": category,
+                    "tags": ", ".join(tags),
+                    "reason": (
+                        "Zero {type} matched the following filters"
+                        "\n"
+                        "\n  - [blue]{category.value}[/] {type}"
+                        "\n  - [blue]{incl}[/] admin-generated {type}"
+                        "\n  - with tags [blue]{tags}"
+                    )
+                }
+                raise ContentDoesNotExist(type="pinboards", **info)
 
             if data['isLastBatch']:
                 break

cs_tools/api/middlewares/search.py

@@ -118,12 +118,15 @@ def __call__(
                    ).json()
 
             if not d['headers']:
-                raise ContentDoesNotExist(type='LOGICAL_TABLE', name=guid)
+                raise ContentDoesNotExist(
+                    type='LOGICAL_TABLE',
+                    reason="No table or worksheet found with the name [blue]{name}"
+                )
 
             d = [_ for _ in d['headers'] if _['name'].casefold() == guid.casefold()]
 
             if len(d) > 1:
-                raise AmbiguousContentError(name=guid, type='LOGICAL_TABLE')
+                raise AmbiguousContentError(type='LOGICAL_TABLE', name=guid)
 
             guid = d[0]['id']
 
@@ -151,9 +154,9 @@ def __call__(
 
             if offset % 500_000 == 0:
                 log.warning(
-                    f'using the Search API to extract >= {offset / 1000: >6,.0f}K rows '
-                    f'is not a scalable practice, please consider adding a filter or '
-                    f'extracting records directly from the underlying data source '
+                    f'using the Search API to extract >= {offset / 1_000_000: >3,.1f}M '
+                    f'rows is not a scalable practice, please consider adding a filter '
+                    f'or extracting records directly from the underlying data source '
                     f'instead!'
                 )
 

cs_tools/api/middlewares/tag.py

@@ -116,6 +116,9 @@ def get(self, tag_name: str, *, create_if_not_exists: bool = False) -> Dict[str,
             tag = self.create(tag_name)
 
         if tag is None:
-            raise ContentDoesNotExist(type='TAG', name=tag_name)
+            raise ContentDoesNotExist(
+                type='tag',
+                reason=f"No tag found with the name [blue]{tag_name}"
+            )
 
         return tag

cs_tools/api/middlewares/tql.py

@@ -7,7 +7,7 @@
 from pydantic import validate_arguments, Field
 
 from cs_tools.data.enums import Privilege
-from cs_tools.errors import InsufficientPrivileges, TableAlreadyExists
+from cs_tools.errors import InsufficientPrivileges
 
 
 log = logging.getLogger(__name__)
@@ -39,7 +39,7 @@ def _check_privileges(self) -> None:
             raise InsufficientPrivileges(
                 user=self.ts.me,
                 service='remote TQL',
-                required_privileges=REQUIRED_PRIVILEGES
+                required_privileges=', '.join(REQUIRED_PRIVILEGES)
             )
 
     @validate_arguments

cs_tools/api/middlewares/tsload.py

@@ -1,15 +1,18 @@
 from typing import Any, Dict, List, Union
 from io import BufferedIOBase, TextIOWrapper
 from tempfile import _TemporaryFileWrapper
+import datetime as dt
 import logging
+import json
 import time
 
 from pydantic import validate_arguments
 
 from cs_tools.data.enums import Privilege
 from cs_tools.errors import InsufficientPrivileges, TSLoadServiceUnreachable
 from cs_tools.const import (
-    FMT_TSLOAD_DATETIME, FMT_TSLOAD_DATE, FMT_TSLOAD_TIME, FMT_TSLOAD_TRUE_FALSE
+    FMT_TSLOAD_DATETIME, FMT_TSLOAD_DATE, FMT_TSLOAD_TIME, FMT_TSLOAD_TRUE_FALSE,
+    APP_DIR
 )
 
 
@@ -25,22 +28,96 @@ class TSLoadMiddleware:
     """
     def __init__(self, ts):
         self.ts = ts
+        # The load server resides on a different port compared to standard ThoughtSpot
+        # services. This is because the service tends to carry heavy file-load
+        # operations, and having a separate web server creates the needed isolation
+        # between standard ThoughtSpot services and tsload operations. By default, this
+        # service runs on all nodes of a ThoughtSpot cluster. This provides load
+        # distribution to address possible simultaneous loads. The tsload server uses
+        # its own load balancer. If an external load balancer is used, the tsload
+        # requests must be sticky, and the tsload load balancer should be disabled.
+        #
+        # To turn off the load balancer, issue the following tscli commands
+        #   tscli --adv service add-gflag etl_http_server.etl_http_server etl_server_enable_load_balancer false
+        #   tscli --adv service add-gflag etl_http_server.etl_http_server etl_server_always_expose_node_ip true
+        #
+        #   DEV NOTE
+        #     On each public method in this middleware, a keyword argument called
+        #     `ignore_node_redirect` which will remove the redirection logic from
+        #     further calls to the tsload service api. Since this is handled on a
+        #     client-by-client basis with no input from the API itself, we expose it as
+        #     a kwarg.
+        #
+        # Further reading:
+        #   https://docs.thoughtspot.com/latest/admin/loading/load-with-tsload.html
+        #
+        self._cache_fp = APP_DIR / '.cache/tsload-node-redirect-by-cycle-id.json'
+
+    def _cache_node_redirect(self, cycle_id: str, *, node_info: Dict = None) -> Dict[str, Dict]:
+        """
+        Method is a total hack.
+        """
+        try:
+            with self._cache_fp.open(mode='r') as j:
+                cache = json.load(j)
+        except FileNotFoundError:
+            cache = {}
+
+        # nothing to write, or we should be reading
+        if node_info is None:
+            return cache
+
+        # write to cache
+        now = dt.datetime.utcnow().timestamp()
+        cache[cycle_id] = {**node_info, 'load_datetime': now}
+
+        # keep only recent data
+        cache = {
+            cycle: details
+            for cycle, details in cache.items()
+            if (now - details['load_datetime']) <= (10 * 86400)  # 10 days
+        }
+
+        with self._cache_fp.open(mode='w') as j:
+            json.dump(cache, j, indent=4, sort_keys=True)
+
+        return cache
+
+    def _check_for_redirect_auth(self, cycle_id: str) -> None:
+        """
+        Attempt a login.
+
+        By default, the tsload service API sits behind a load balancer. When we first
+        init a new load cycle, the balancer will respond with the proper node (if 
+        applicable) to submit file uploads to. If that node is not the main node, then
+        we will be required to authorize again.
+        """
+        cache = self._cache_node_redirect(cycle_id)
+
+        if cycle_id in cache:
+            ds = self.ts.api.ts_dataservice
+            ds._tsload_node = cache[cycle_id]['host']
+            ds._tsload_port = cache[cycle_id]['port']
+            log.debug(f'redirecting to: {ds.etl_server_fullpath}')
+            ds.load_auth()
 
     def _check_privileges(self) -> None:
         """
+        Determine if the user has necessary Data Manager privileges.
         """
         if not set(self.ts.me.privileges).intersection(REQUIRED_PRIVILEGES):
             raise InsufficientPrivileges(
                 user=self.ts.me,
                 service='remote TQL',
-                required_privileges=REQUIRED_PRIVILEGES
+                required_privileges=', '.join(REQUIRED_PRIVILEGES)
             )
 
     @validate_arguments(config=dict(arbitrary_types_allowed=True))
     def upload(
         self,
         fd: Union[BufferedIOBase, TextIOWrapper, _TemporaryFileWrapper],
         *,
+        ignore_node_redirect: bool = False,
         database: str,
         table: str,
         schema_: str = 'falcon_default_schema',
@@ -85,6 +162,11 @@ def upload(
         fp : pathlib.Path
           file to load to thoughtspot
 
+        ignore_node_redirect : bool  [default: False]
+          whether or not to ignore node redirection
+
+        **tsload_options
+
         Returns
         -------
         cycle_id
@@ -150,17 +232,43 @@ def upload(
                 http_error=e
             )
 
-        cycle_id = r.json()['cycle_id']
-        self.ts.api.ts_dataservice.load_start(cycle_id, fd=fd)
-        self.ts.api.ts_dataservice.load_commit(cycle_id)
-        return cycle_id
+        data = r.json()
+        self._cache_node_redirect(data['cycle_id'], node_info=data.get('node_address', None))
+
+        if not ignore_node_redirect:
+            self._check_for_redirect_auth(data['cycle_id'])
+
+        self.ts.api.ts_dataservice.load_start(data['cycle_id'], fd=fd)
+        self.ts.api.ts_dataservice.load_commit(data['cycle_id'])
+        return data['cycle_id']
 
     @validate_arguments
-    def status(self, cycle_id: str, *, wait_for_complete: bool = False):
+    def status(
+        self,
+        cycle_id: str,
+        *,
+        ignore_node_redirect: bool = False,
+        wait_for_complete: bool = False
+    ) -> Dict[str, Any]:
         """
+        Get the status of a previously started data load.
+
+        Parameters
+        ----------
+        cycle_id : str
+          data load to check on
+
+        ignore_node_redirect : bool  [default: False]
+          whether or not to ignore node redirection
+
+        wait_for_complete: bool  [default: False]
+          poll the load server until it responds with OK or ERROR
         """
         self._check_privileges()
 
+        if not ignore_node_redirect:
+            self._check_for_redirect_auth(cycle_id=cycle_id)
+
         while True:
             r = self.ts.api.ts_dataservice.load_status(cycle_id)
             data = r.json()

cs_tools/api/middlewares/user.py

@@ -86,11 +86,14 @@ def get(
         user = r.json()['headers']
 
         if not user:
-            raise ContentDoesNotExist(type='USER', name=principal)
+            raise ContentDoesNotExist(
+                type='USER',
+                reason=f"No user found with the name [blue]{principal}"
+            )
 
         if error_if_ambiguous:
             if len(user) > 1:
-                raise AmbiguousContentError(type='USER', name=principal)
+                raise AmbiguousContentError(type='user', name=principal)
             user = user[0]
 
         return user