fix: harden ThoughtSpot table import recovery

thoughtspot_deployer.py

@@ -2217,6 +2217,15 @@ class ThoughtSpotDeployer:
                     return result['object']
                 return None
 
+            def _tml_objects_by_name(expected_names, objects):
+                named = {}
+                for idx, obj in enumerate(objects or []):
+                    header = obj.get('response', {}).get('header', {}) if isinstance(obj, dict) else {}
+                    raw_name = header.get('name') or header.get('display_name')
+                    table_name = (raw_name or (expected_names[idx] if idx < len(expected_names) else f"TABLE_{idx}")).upper()
+                    named[table_name] = obj
+                return named
+
             def _import_tml_chunk(phase_label, names, tmls, create_new):
                 payload = {
                     "metadata_tmls": tmls,
@@ -2243,6 +2252,25 @@ class ThoughtSpotDeployer:
                 )
                 elapsed = time.time() - start
 
+                if response.status_code == 401:
+                    log_progress(f"   ⚠️  {phase_label} auth expired; re-authenticating and retrying once")
+                    if _slog:
+                        _slog.log(
+                            "thoughtspot",
+                            "tml import auth expired",
+                            phase=phase_label,
+                            table_names=names,
+                            elapsed_s=round(elapsed, 1),
+                        )
+                    if self.authenticate():
+                        retry_start = time.time()
+                        response = self.session.post(
+                            f"{self.base_url}/api/rest/2.0/metadata/tml/import",
+                            json=payload,
+                            timeout=360,
+                        )
+                        elapsed += time.time() - retry_start
+
                 if response.status_code == 200:
                     objects = _normalize_tml_import_objects(response.json())
                     if objects is None:
@@ -2337,7 +2365,7 @@ class ThoughtSpotDeployer:
                     results['warnings'].append(message)
 
             def _import_tmls_chunked(phase_label, names, tmls, create_new, chunk_size=3, fatal_errors=True):
-                all_objects = []
+                all_objects = {}
                 retriable_statuses = {502, 503, 504}
 
                 def _retry_single_table_after_gateway_timeout(single_name, single_tml, max_attempts=3):
@@ -2375,9 +2403,7 @@ class ThoughtSpotDeployer:
                     if isinstance(result, requests.Response) and result.status_code in retriable_statuses:
                         if create_new:
                             resolved = _resolve_existing_tables_after_timeout(chunk_names)
-                            for name in chunk_names:
-                                if name in resolved:
-                                    all_objects.append(resolved[name])
+                            all_objects.update(resolved)
                             missing = [
                                 (name, tml)
                                 for name, tml in zip(chunk_names, chunk_tmls)
@@ -2416,7 +2442,7 @@ class ThoughtSpotDeployer:
                                     if fatal_errors:
                                         return None
                                     continue
-                                all_objects.extend(single_result)
+                                all_objects.update(_tml_objects_by_name([single_name], single_result))
                             continue
 
                         log_progress(f"   ⚠️  {chunk_label} timed out; retrying once")
@@ -2433,7 +2459,7 @@ class ThoughtSpotDeployer:
                         if create_new and len(chunk_names) == 1:
                             resolved = _resolve_existing_tables_after_timeout(chunk_names)
                             if chunk_names[0] in resolved:
-                                all_objects.append(resolved[chunk_names[0]])
+                                all_objects[chunk_names[0]] = resolved[chunk_names[0]]
                                 continue
                         error = (
                             f"{phase_label} failed: HTTP {result.status_code} - {result.text}"
@@ -2448,7 +2474,7 @@ class ThoughtSpotDeployer:
                         if fatal_errors:
                             return None
                         continue
-                    all_objects.extend(result)
+                    all_objects.update(_tml_objects_by_name(chunk_names, result))
 
                 return all_objects
 
@@ -2466,19 +2492,27 @@ class ThoughtSpotDeployer:
             # Send tables in small chunks. ThoughtSpot can intermittently 504
             # large metadata imports at the gateway even for valid table TML.
             log_progress(f"   Sending table creation requests for {len(table_tmls_batch1)} tables...")
+            table_import_chunk_size = 1 if table_count >= 7 else 3
             objects = _import_tmls_chunked(
                 "Batch 1",
                 table_names_order,
                 table_tmls_batch1,
                 create_new=True,
+                chunk_size=table_import_chunk_size,
             )
 
             if objects is None:
                 return results
 
-            # Process each table result
-            for idx, obj in enumerate(objects):
-                table_name = table_names_order[idx] if idx < len(table_names_order) else f"TABLE_{idx}"
+            # Process each table result by table name. Gateway-timeout recovery can
+            # return a partial set, so positional matching is unsafe here.
+            for table_name in table_names_order:
+                obj = objects.get(table_name)
+                if obj is None:
+                    error = f"Table {table_name} failed: no import response after retry"
+                    print(f"[ThoughtSpot]    ❌ {error}", flush=True)
+                    results['errors'].append(error)
+                    continue
                 
                 if obj.get('response', {}).get('status', {}).get('status_code') == 'OK':
                     table_guid = obj.get('response', {}).get('header', {}).get('id_guid')
@@ -2504,6 +2538,25 @@ class ThoughtSpotDeployer:
                 log_progress("   ❌ No tables were created successfully in Batch 1")
                 return results
 
+            missing_tables = [name for name in table_names_order if name not in table_guids]
+            if missing_tables:
+                error = (
+                    f"Batch 1 incomplete: {len(table_guids)}/{len(table_names_order)} "
+                    f"tables created; missing {', '.join(missing_tables)}"
+                )
+                log_progress(f"   ❌ {error}")
+                results['errors'].append(error)
+                if _slog:
+                    _slog.log(
+                        "thoughtspot",
+                        "table creation incomplete",
+                        error=error,
+                        table_count=len(table_names_order),
+                        created_count=len(table_guids),
+                        missing_tables=missing_tables,
+                    )
+                return results
+
             # Assign tags to tables
             table_guid_list = list(table_guids.values())
             print(f"🔍 DEBUG BEFORE TAG CALL: tag_name='{tag_name}', table_guid_list={table_guid_list}")
@@ -2552,13 +2605,20 @@ class ThoughtSpotDeployer:
                     table_names_order_batch2,
                     table_tmls_batch2,
                     create_new=False,
+                    chunk_size=table_import_chunk_size,
                     fatal_errors=False,
                 )
 
                 if objects is not None:
-                    # Process each result
-                    for idx, obj in enumerate(objects):
-                        table_name = table_names_order_batch2[idx] if idx < len(table_names_order_batch2) else f"TABLE_{idx}"
+                    # Process each result by table name. Partial retry recovery can
+                    # return fewer objects than requested.
+                    for table_name in table_names_order_batch2:
+                        obj = objects.get(table_name)
+                        if obj is None:
+                            warning = f"Joins for {table_name} failed: no import response after retry"
+                            print(f"[ThoughtSpot]    ⚠️  {warning}", flush=True)
+                            results['warnings'].append(warning)
+                            continue
                         
                         if obj.get('response', {}).get('status', {}).get('status_code') == 'OK':
                             print(f"[ThoughtSpot]    ✅ {table_name} joins added", flush=True)