Merge branch 'v1.6.2' into dev

Author: boonhapus

cs_tools/api/workflows/tsload.py

@@ -166,7 +166,7 @@ async def upload_data(
         http._redirected_url_due_to_tsload_load_balancer = httpx.URL(host=redirect["host"], port=redirect["port"])  # type: ignore[attr-defined]
         # DEV NOTE: @boonhapus, 2025/01/28
         # Technically speaking, this endpoint just delegates to the AUTH SERVICE on each node, so any persistent login
-        # API method would work here, it just needs to point at the redirected node. CS Tools offers multiple login
+        # API method should work here, it just needs to point at the redirected node. CS Tools offers multiple login
         # methods, but it's a pretty safe bet that the customer on Falcon will have a BASIC auth context.
         r = await http.v1_dataservice_dataload_session(**auth_info)
         r.raise_for_status()
@@ -182,19 +182,30 @@ async def upload_data(
     return data["cycle_id"]
 
 
-async def wait_for_dataload_completion(cycle_id: _types.GUID, *, http: RESTAPIClient) -> _types.APIResult:
+async def wait_for_dataload_completion(
+    cycle_id: _types.GUID,
+    *,
+    timeout: int = 300,
+    http: RESTAPIClient,
+) -> _types.APIResult:
     """Wait for dataload to complete."""
+    start = dt.datetime.now(tz=dt.timezone.utc)
+
     while True:
         _LOG.info(f"Checking status of dataload {cycle_id}...")
         r = await http.v1_dataservice_dataload_status(cycle_id=cycle_id)
 
         status_data = r.json()
 
-        if "code" in status_data["status"]:
+        if "code" in status_data.get("status", {}):
+            break
+
+        if (dt.datetime.now(tz=dt.timezone.utc) - start).seconds >= timeout:
+            _LOG.warning(f"Reached the {timeout / 60:.1f} minute CS Tools timeout, giving up on cycle_id {cycle_id}")
             break
 
         _LOG.debug(status_data)
-        await asyncio.sleep(1)
+        await asyncio.sleep(5)
 
     _LOG.info(
         f"Cycle ID: {status_data['cycle_id']} ({status_data['status']['code']})"

cs_tools/cli/tools/scriptability/utils.py

@@ -4,7 +4,9 @@
 import datetime as dt
 import json
 import logging
+import os
 import pathlib
+import re
 
 import pydantic
 import rich
@@ -16,6 +18,9 @@
 
 _LOG = logging.getLogger(__name__)
 
+RE_ENVVAR_STRUCTURE = re.compile(r"\$\{\{\s*env\.(?P<envvar>[A-Za-z0-9_]+)\s*\}\}", flags=re.MULTILINE)
+"""Variable structure looks like ${{ env.MY_VAR_NAME }} and can be inline with other text."""
+
 
 def is_allowed_object(
     metadata_object: _types.APIResult,
@@ -197,6 +202,43 @@ def checkpoint(
 
     def disambiguate(self, tml: _types.TMLObject, delete_unmapped_guids: bool = True) -> _types.TMLObject:
         """Disambiguate an incoming TML object."""
+        # ADDITIONAL MAPPING NEEDS TO COME FIRST, IN CASE WE DELETE A GUID DURING DISAMBIGUATION.
+        if self.additional_mapping:
+            try:
+                text = original = tml.dumps(format_type="YAML")
+
+                for to_find, to_replace in self.additional_mapping.items():
+                    # PRE-PROCESS TO FETCH ENVIRONMENT VARIABLES, WITH FALLBACK TO THE ORIGINAL VALUE.
+                    if match := RE_ENVVAR_STRUCTURE.search(to_replace):
+                        # the whole matched string ............... ${{ env.MY_VAR_NAME }}
+                        envvar_template = match.group(0)
+                        # just the value name .................... MY_VAR_NAME
+                        envvar_name = match.group("envvar")
+                        # fetch the value from os.environ ........ anything~
+                        envvar_value = os.getenv(envvar_name, envvar_template)
+
+                        to_replace = to_replace.replace(envvar_template, envvar_value)
+
+                    # PERFORM STRING REPLACEMENTS VIA EXACT MATCH (case sensitive)
+                    if to_find in text:
+                        _LOG.info(f"Replacing '{to_find}' with '{to_replace}' in {tml.name} ({tml.tml_type_name})")
+                        text = text.replace(to_find, to_replace)
+
+                if text != original:
+                    tml = tml.loads(text)
+
+            except (IndexError, thoughtspot_tml.exceptions.TMLDecodeError) as e:
+                # JUST FALL BACK TO TML WITH THE VARIABLES IN IT, WHICH SHOULD FAIL ON TML IMPORT ANYWAY.
+                _LOG.warning(f"Could not variablize '{tml.name}' ({tml.tml_type_name}), see logs for details..")
+                _LOG.debug(e, exc_info=True)
+
+                if hasattr(e, "parent_exc"):
+                    _LOG.debug(f"due to..\n{e}", exc_info=True)
+
+                _LOG.info(f"TML MODIFIED STATE:\n\n{text}")
+                raise
+
+        # STANDARD DISMABIGUATION
         tml = thoughtspot_tml.utils.disambiguate(
             tml=tml,
             guid_mapping=self.mapping,

cs_tools/sync/falcon/syncer.py

@@ -121,6 +121,7 @@ def dump(self, tablename: str, *, data: _types.TableRowsFormat) -> None:
                 "schema": self.schema_,
                 "table": tablename,
                 "has_header_row": True,
+                "date_time_format": utils.FMT_TSLOAD_DATETIME,
                 "ignore_node_redirect": self.ignore_load_balancer_redirect,
             }