Harden incremental results rendering for mixed types

Author: Maxteabag

sqlit/domains/query/ui/mixins/query_results.py

@@ -151,6 +151,7 @@ def _schedule_results_render(
         rows: list[tuple],
         *,
         escape: bool,
+        coerce_to_str_columns: set[int] | None = None,
         start_index: int,
         row_limit: int,
         render_token: int,
@@ -171,6 +172,17 @@ def add_batch() -> None:
             if end <= index:
                 return
             batch = rows[index:end]
+            if coerce_to_str_columns:
+                coerced: list[tuple] = []
+                for row in batch:
+                    new_row = []
+                    for col_idx, value in enumerate(row):
+                        if col_idx in coerce_to_str_columns and value is not None:
+                            new_row.append(str(value))
+                        else:
+                            new_row.append(value)
+                    coerced.append(tuple(new_row))
+                batch = coerced
             try:
                 table.add_rows(batch)
             except Exception as exc:
@@ -218,26 +230,46 @@ def _render_results_table_incremental(
         has_decimal_in_initial = any(
             isinstance(value, Decimal) for row in initial_rows for value in row
         )
-        if has_decimal_in_initial:
-            decimal_types = self._get_decimal_column_types(rows)
-            if decimal_types:
-                from textual_fastdatatable.backend import ArrowBackend
-                import pyarrow as pa
-
-                arrays = []
-                for idx, _name in enumerate(columns):
-                    values = [row[idx] for row in initial_rows] if initial_rows else []
-                    col_type = decimal_types.get(idx)
-                    if col_type is not None:
-                        arrays.append(pa.array(values, type=col_type))
-                    else:
-                        arrays.append(pa.array(values))
-                table_backend = ArrowBackend(pa.Table.from_arrays(arrays, names=columns))
-                table = self._build_results_table(columns, initial_rows, escape=escape, backend=table_backend)
+        coerce_to_str_columns: set[int] | None = None
+        try:
+            if has_decimal_in_initial:
+                decimal_types = self._get_decimal_column_types(rows)
+                if decimal_types:
+                    from textual_fastdatatable.backend import ArrowBackend
+                    import pyarrow as pa
+
+                    arrays = []
+                    coerce_to_str_columns = set()
+                    for idx, _name in enumerate(columns):
+                        values = [row[idx] for row in initial_rows] if initial_rows else []
+                        col_type = decimal_types.get(idx)
+                        try:
+                            if col_type is not None:
+                                arrays.append(pa.array(values, type=col_type))
+                            else:
+                                arrays.append(pa.array(values))
+                        except (TypeError, ValueError, pa.ArrowInvalid, pa.ArrowTypeError):
+                            coerce_to_str_columns.add(idx)
+                            arrays.append(
+                                pa.array(
+                                    [str(value) if value is not None else None for value in values],
+                                    type=pa.string(),
+                                )
+                            )
+                    table_backend = ArrowBackend(pa.Table.from_arrays(arrays, names=columns))
+                    table = self._build_results_table(columns, initial_rows, escape=escape, backend=table_backend)
+                else:
+                    table = self._build_results_table(columns, initial_rows, escape=escape)
             else:
                 table = self._build_results_table(columns, initial_rows, escape=escape)
-        else:
-            table = self._build_results_table(columns, initial_rows, escape=escape)
+        except Exception as exc:
+            try:
+                self.log.error(f"Results table build failed; falling back to full render: {exc}")
+            except Exception:
+                pass
+            if render_token == getattr(self, "_results_render_token", 0):
+                self._replace_results_table_with_data(columns, rows, escape=escape)
+            return
         if render_token != getattr(self, "_results_render_token", 0):
             return
         self._replace_results_table_with_table(table)
@@ -246,6 +278,7 @@ def _render_results_table_incremental(
             columns,
             rows,
             escape=escape,
+            coerce_to_str_columns=coerce_to_str_columns,
             start_index=initial_count,
             row_limit=row_limit,
             render_token=render_token,

tests/ui/test_results_incremental_rendering.py

@@ -152,3 +152,56 @@ def _wrapped_replace(self, columns, rows, *, escape):
         assert fallback_called["value"] is True
 
     monkeypatch.setattr(SqlitDataTable, "add_rows", original_add_rows)
+
+
+class _WeirdType:
+    def __init__(self, value: str) -> None:
+        self.value = value
+
+    def __str__(self) -> str:
+        return f"Weird({self.value})"
+
+
+@pytest.mark.asyncio
+async def test_incremental_rendering_coerces_unsupported_initial_types():
+    """Unsupported types in initial rows should be coerced without forcing fallback."""
+    connections = [create_test_connection("test-db", "sqlite")]
+    mock_connections = MockConnectionStore(connections)
+    mock_settings = MockSettingsStore({"theme": "tokyo-night"})
+
+    services = build_test_services(
+        connection_store=mock_connections,
+        settings_store=mock_settings,
+    )
+    app = SSMSTUI(services=services)
+
+    columns = ["id", "amount", "range"]
+    rows = []
+    for i in range(201):
+        rows.append((i + 1, Decimal(f"{i + 1}.25"), _WeirdType(str(i))))
+
+    async with app.run_test(size=(120, 40)) as pilot:
+        await pilot.pause()
+
+        fallback_called = {"value": False}
+        original_replace = app._replace_results_table_with_data
+
+        def _wrapped_replace(self, columns, rows, *, escape):
+            fallback_called["value"] = True
+            return original_replace(columns, rows, escape=escape)
+
+        app._replace_results_table_with_data = MethodType(_wrapped_replace, app)
+
+        await app._display_query_results(
+            columns=columns,
+            rows=rows,
+            row_count=len(rows),
+            truncated=False,
+            elapsed_ms=0,
+        )
+
+        for _ in range(3):
+            await pilot.pause(0.05)
+
+        assert app.results_table.row_count == len(rows)
+        assert fallback_called["value"] is False