fix: preserve task identity in threaded query execution

msticpy/data/core/query_provider_connections_mixin.py

@@ -412,6 +412,77 @@ def _create_split_queries(
         logger.info("Split query into %s chunks", len(split_queries))
         return split_queries
 
+    @staticmethod
+    async def _run_tagged_task(
+        query_id: str, future: asyncio.Future[pd.DataFrame]
+    ) -> tuple[str, pd.DataFrame | None, Exception | None]:
+        """Await `future`, returning its query_id alongside the outcome.
+
+        `asyncio.as_completed` yields tasks in completion order, not the
+        order they were submitted in, and does not preserve any link back
+        to the key `future` was stored under. Tagging the outcome with its
+        own query_id here means identity survives being raced through
+        `as_completed`, instead of being reconstructed afterwards by
+        zipping the completion-ordered results against the (differently
+        ordered) submission keys.
+        """
+        try:
+            result: pd.DataFrame = await future
+        except Exception as exc:  # pylint: disable=broad-exception-caught
+            return query_id, None, exc
+        return query_id, result, None
+
+    @staticmethod
+    async def _run_tasks_threaded(
+        thread_tasks: dict[str, asyncio.Future[pd.DataFrame]],
+        *,
+        progress: bool,
+    ) -> tuple[dict[str, pd.DataFrame], list[str]]:
+        """Await already-submitted futures, returning results by query_id.
+
+        Shared by the first attempt and the retry pass in
+        `_exec_queries_threaded` so both merge into the same
+        query_id-keyed results dict before final ordering is computed --
+        that is what lets a query that only succeeds on retry still land
+        in its correct submission-order position rather than being
+        appended after every query that succeeded on the first attempt.
+        """
+        results: dict[str, pd.DataFrame] = {}
+        failed_tasks_ids: list[str] = []
+        tagged_tasks = [
+            QueryProviderConnectionsMixin._run_tagged_task(query_id, future)
+            for query_id, future in thread_tasks.items()
+        ]
+        task_iter = asyncio.as_completed(tagged_tasks)
+        if progress:
+            task_iter = tqdm(
+                task_iter,
+                total=len(tagged_tasks),
+                unit="sub-queries",
+                desc="Running",
+            )
+        for coro in task_iter:
+            query_id, result, exc = await coro
+            if exc is None:
+                logger.info("Query task '%s' completed successfully.", query_id)
+                results[query_id] = result
+            elif isinstance(exc, MsticpyDataQueryError):
+                logger.warning(
+                    "Query task '%s' failed",
+                    query_id,
+                )
+                exc.display_exception()
+                failed_tasks_ids.append(query_id)
+            else:
+                logger.warning(
+                    "Query task '%s' failed with exception %s - %s",
+                    query_id,
+                    exc.__class__.__name__,
+                    exc.args,
+                )
+                failed_tasks_ids.append(query_id)
+        return results, failed_tasks_ids
+
     @staticmethod
     async def _exec_queries_threaded(
         query_tasks: dict[str, partial],
@@ -426,68 +497,44 @@ async def _exec_queries_threaded(
         event_loop: asyncio.AbstractEventLoop = _get_event_loop()
 
         with ThreadPoolExecutor(max_workers=max_workers) as executor:
-            # add the additional connections
             thread_tasks: dict[str, asyncio.Future[pd.DataFrame]] = {
                 query_id: event_loop.run_in_executor(executor, query_func)
                 for query_id, query_func in query_tasks.items()
             }
-            results: list[pd.DataFrame] = []
-            failed_tasks_ids: list[str] = []
-            if progress:
-                task_iter = tqdm(
-                    asyncio.as_completed(thread_tasks.values()),
-                    unit="sub-queries",
-                    desc="Running",
+            results, failed_tasks_ids = (
+                await QueryProviderConnectionsMixin._run_tasks_threaded(
+                    thread_tasks, progress=progress
                 )
-            else:
-                task_iter = asyncio.as_completed(thread_tasks.values())
-            ids_and_tasks = dict(zip(thread_tasks, task_iter, strict=False))
-            for query_id, thread_task in ids_and_tasks.items():
-                try:
-                    result: pd.DataFrame | str | None = await thread_task
-                    logger.info("Query task '%s' completed successfully.", query_id)
-                    results.append(result)
-                except MsticpyDataQueryError as msticpy_dqe:
-                    logger.warning(
-                        "Query task '%s' failed",
-                        query_id,
-                    )
-                    msticpy_dqe.display_exception()
-                    # Reusing thread task would result in:
-                    # RuntimeError: cannot reuse already awaited coroutine
-                    # A new task should be queued
-                    failed_tasks_ids.append(query_id)
-                except Exception as exc:  # pylint: disable=broad-exception-caught
-                    logger.warning(
-                        "Query task '%s' failed with exception %s - %s",
-                        query_id,
-                        exc.__class__.__name__,
-                        exc.args,
-                    )
-                    # Reusing thread task would result in:
-                    # RuntimeError: cannot reuse already awaited coroutine
-                    # A new task should be queued
-                    failed_tasks_ids.append(query_id)
-
-        # Sort the results by the order of the tasks
-        results = [result for _, result in sorted(zip(thread_tasks, results, strict=False))]
+            )
 
         if retry and failed_tasks_ids:
-            failed_results: pd.DataFrame = (
-                await QueryProviderConnectionsMixin._exec_queries_threaded(
-                    {
-                        failed_tasks_id: query_tasks[failed_tasks_id]
-                        for failed_tasks_id in failed_tasks_ids
-                    },
-                    progress=progress,
-                    retry=False,
-                    max_workers=max_workers,
+            with ThreadPoolExecutor(max_workers=max_workers) as retry_executor:
+                retry_thread_tasks: dict[str, asyncio.Future[pd.DataFrame]] = {
+                    query_id: event_loop.run_in_executor(
+                        retry_executor, query_tasks[query_id]
+                    )
+                    for query_id in failed_tasks_ids
+                }
+                retry_results, _ = (
+                    await QueryProviderConnectionsMixin._run_tasks_threaded(
+                        retry_thread_tasks, progress=progress
+                    )
                 )
-            )
-            if not failed_results.empty:
-                results.append(failed_results)
-        if results:
-            return pd.concat(results, ignore_index=True)
+            # Merge back into the same query_id-keyed dict the first
+            # attempt populated, so a query that only succeeded on retry
+            # still lands in its correct submission-order position below
+            # instead of being appended after every query that succeeded
+            # on the first attempt.
+            results.update(retry_results)
+
+        # Results are keyed by query_id, so this restores the original
+        # submission order regardless of which attempt each one
+        # succeeded on, or the order tasks completed in.
+        ordered_results: list[pd.DataFrame] = [
+            results[query_id] for query_id in thread_tasks if query_id in results
+        ]
+        if ordered_results:
+            return pd.concat(ordered_results, ignore_index=True)
 
         logger.warning("All queries failed.")
         return pd.DataFrame()

tests/data/test_async_queries.py

@@ -5,7 +5,10 @@
 # --------------------------------------------------------------------------
 """Test async connections and split queries."""
 
+import asyncio
+import time
 from datetime import datetime, timedelta, timezone
+from functools import partial
 from unittest.mock import patch
 
 import pandas as pd
@@ -14,7 +17,10 @@
 from msticpy.common.exceptions import MsticpyDataQueryError
 from msticpy.data.core.data_providers import QueryProvider
 from msticpy.data.drivers.local_data_driver import LocalDataDriver
-from msticpy.data.core.query_provider_connections_mixin import _calc_split_ranges
+from msticpy.data.core.query_provider_connections_mixin import (
+    QueryProviderConnectionsMixin,
+    _calc_split_ranges,
+)
 from msticpy.data.drivers.driver_base import DriverProps
 
 from ..unit_test_lib import get_test_data_path
@@ -290,6 +296,86 @@ def test_multiple_connections_threaded():
         )
 
 
+def test_exec_queries_threaded_preserves_identity_out_of_order():
+    """Results must stay attached to the query_id that actually produced them.
+
+    asyncio.as_completed() yields tasks in completion order, not the order
+    they were submitted in. A correct implementation cannot recover which
+    result belongs to which query_id by position (e.g. zipping the
+    submission-ordered keys against the completion-ordered results) --
+    that silently mislabels results whenever tasks don't finish in
+    submission order, which is the common case under real concurrency.
+    """
+
+    def _tagged_result(label: str, delay: float) -> pd.DataFrame:
+        time.sleep(delay)
+        return pd.DataFrame({"query_id": [label]})
+
+    # "slow" is submitted first but finishes last.
+    query_tasks = {
+        "slow": partial(_tagged_result, "slow", 0.2),
+        "fast": partial(_tagged_result, "fast", 0.0),
+    }
+
+    result: pd.DataFrame = asyncio.run(
+        QueryProviderConnectionsMixin._exec_queries_threaded(
+            query_tasks, progress=False
+        )
+    )
+
+    # Each row must carry the query_id of the task that actually produced
+    # it, and the final frame must preserve submission order (slow before
+    # fast) regardless of which one completed first.
+    check.equal(list(result["query_id"]), ["slow", "fast"])
+
+
+def test_exec_queries_threaded_retries_the_task_that_actually_failed():
+    """The retry mechanism must retry the query_id that failed, not
+    whichever query_id happened to occupy the same position once results
+    were (mis)reconstructed."""
+
+    def _fast_ok(label: str) -> pd.DataFrame:
+        return pd.DataFrame({"query_id": [label]})
+
+    def _slow_fail(label: str) -> pd.DataFrame:
+        time.sleep(0.2)
+        raise MsticpyDataQueryError(label)
+
+    call_counts: dict[str, int] = {"a": 0, "b": 0}
+
+    def _tracked(label: str, should_fail: bool) -> pd.DataFrame:
+        call_counts[label] += 1
+        if should_fail and call_counts[label] == 1:
+            return _slow_fail(label)
+        return _fast_ok(label)
+
+    # "a" is submitted first and fails on its first attempt; "b" is
+    # submitted second, is quick, and always succeeds. "a" finishes after
+    # "b" since it is the one with the artificial delay.
+    query_tasks = {
+        "a": partial(_tracked, "a", True),
+        "b": partial(_tracked, "b", False),
+    }
+
+    result: pd.DataFrame = asyncio.run(
+        QueryProviderConnectionsMixin._exec_queries_threaded(
+            query_tasks, progress=False, retry=True
+        )
+    )
+
+    # "a" must have been retried (called twice); "b" never failed, so it
+    # must only have been called once.
+    check.equal(call_counts["a"], 2)
+    check.equal(call_counts["b"], 1)
+    # "a" was submitted first, so its row must come first in the result
+    # even though it only succeeded on retry (after "b" had already
+    # completed on the first attempt). Asserting the exact order --
+    # not just the sorted set of query_ids -- is what catches the
+    # retried-results-appended-at-the-end bug: a naive fix passes the
+    # sorted comparison but still produces ["b", "a"] here.
+    check.equal(list(result["query_id"]), ["a", "b"])
+
+
 def test_split_queries_sync():
     """Test queries split into time segments."""
     prov_args = dict(query_paths=_LOCAL_DATA_PATHS, data_paths=_LOCAL_DATA_PATHS)