From d71d96191d173228dd9333011250185cd8bdee67 Mon Sep 17 00:00:00 2001 From: JAI Date: Fri, 2 Oct 2026 22:09:56 +0530 Subject: [PATCH 1/2] fix: preserve task identity in threaded query execution _exec_queries_threaded reconstructed which result belonged to which query_id by zipping thread_tasks (submission-order keys) against both an intermediate task iterator and, later, the results list (completion order, from asyncio.as_completed). Those two sequences aren't the same order, so the zip silently mispaired query_ids with results and exceptions whenever tasks didn't complete in submission order -- the common case under real concurrency. Two concrete failure modes from this: - The final concatenated DataFrame's row order ("Sort the results by the order of the tasks") was actually ordered by sorting already- mismatched (wrong_id, result) pairs, which only incidentally matched submission order when alphabetical id order and submission order happened to agree. - Worse, which query_id got logged as "failed" and (with retry=True) actually retried was also determined by this same mismatch, not by which task genuinely failed. A fast, never-failing query could be logged and retried under a different id, while the real failure's result silently never makes it into the output -- a correctness bug on the retry path, not just an ordering annoyance. Fix: tag each task with its own query_id before running it through asyncio.as_completed via a small _run_tagged_task helper, so identity travels with the outcome instead of being reconstructed afterwards by position. Results are now collected into a dict keyed by query_id; final ordering is rebuilt by iterating the original submission-order keys, and retry correctly targets the query_id that actually failed. Added two regression tests to test_async_queries.py: - test_exec_queries_threaded_preserves_identity_out_of_order - test_exec_queries_threaded_retries_the_task_that_actually_failed Confirmed via git stash against unmodified main that the second test fails on the pre-fix code: it retries the wrong query_id (the one that never failed) and silently drops the one that actually did, giving a final result of two copies of the never-failing query's data and zero copies of the one that failed and was "retried". --- .../core/query_provider_connections_mixin.py | 66 +++++++++------ tests/data/test_async_queries.py | 82 ++++++++++++++++++- 2 files changed, 124 insertions(+), 24 deletions(-) diff --git a/msticpy/data/core/query_provider_connections_mixin.py b/msticpy/data/core/query_provider_connections_mixin.py index cc64098cc..e859ff773 100644 --- a/msticpy/data/core/query_provider_connections_mixin.py +++ b/msticpy/data/core/query_provider_connections_mixin.py @@ -412,6 +412,26 @@ 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 _exec_queries_threaded( query_tasks: dict[str, partial], @@ -431,46 +451,46 @@ async def _exec_queries_threaded( query_id: event_loop.run_in_executor(executor, query_func) for query_id, query_func in query_tasks.items() } - results: list[pd.DataFrame] = [] + 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( - asyncio.as_completed(thread_tasks.values()), + task_iter, + total=len(tagged_tasks), unit="sub-queries", desc="Running", ) - 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 + 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.append(result) - except MsticpyDataQueryError as msticpy_dqe: + results[query_id] = result + elif isinstance(exc, MsticpyDataQueryError): 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 + exc.display_exception() failed_tasks_ids.append(query_id) - except Exception as exc: # pylint: disable=broad-exception-caught + else: 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))] + # Results are keyed by query_id, so this restores the original + # submission order regardless of 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 retry and failed_tasks_ids: failed_results: pd.DataFrame = ( @@ -485,9 +505,9 @@ async def _exec_queries_threaded( ) ) if not failed_results.empty: - results.append(failed_results) - if results: - return pd.concat(results, ignore_index=True) + ordered_results.append(failed_results) + if ordered_results: + return pd.concat(ordered_results, ignore_index=True) logger.warning("All queries failed.") return pd.DataFrame() diff --git a/tests/data/test_async_queries.py b/tests/data/test_async_queries.py index 18b8703f8..4e559732e 100644 --- a/tests/data/test_async_queries.py +++ b/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,80 @@ 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) + check.equal(sorted(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) From af83cc50e46a1727c38ef1260cd7f984e38251d5 Mon Sep 17 00:00:00 2001 From: JAI Date: Thu, 15 Oct 2026 08:39:59 +0530 Subject: [PATCH 2/2] fix: preserve submission order when a retried task succeeds ianhelle's review on #952 found a residual ordering gap: ordered_results was built from the first-attempt results dict before the retry pass ran, so a retried task's result was always appended at the end regardless of its real submission position. For tasks submitted as a, b where a fails initially and succeeds on retry while b succeeds immediately, the output was b, a instead of a, b. Extracted the per-batch run loop (submit tagged tasks, await via as_completed, bucket into a query_id-keyed results dict plus a failed-id list) into a new _run_tasks_threaded helper, used for both the first attempt and the retry pass. The retry pass's results are merged back into the same results dict via results.update(retry_results) before the final submission-order reconstruction runs, so a query that only succeeded on retry now lands in its correct position instead of after every query that succeeded on the first attempt. This replaces the previous retry path, which recursively called _exec_queries_threaded(retry=False) and appended its already-concatenated return value to the end of ordered_results -- that's exactly what threw away per-query identity for the retried batch and made correct interleaving impossible without this refactor. Updated test_exec_queries_threaded_retries_the_task_that_actually_failed per review: replaced the sorted(result["query_id"]) assertion (which can't distinguish correct order from the bug) with an exact a, b order assertion. Confirmed this fails on the previous commit's code -- order comes back as b, a -- and passes with this fix. tests/data/test_async_queries.py: 6/6 pass. tests/data/: 383 passed, 5 skipped (pre-existing, unrelated). black/pylint/mypy clean on the changed lines (the file has pre-existing lines that reformat under my locally installed black version due to version drift, same as noted in the original PR -- none of them touched by this change). --- .../core/query_provider_connections_mixin.py | 123 +++++++++++------- tests/data/test_async_queries.py | 8 +- 2 files changed, 82 insertions(+), 49 deletions(-) diff --git a/msticpy/data/core/query_provider_connections_mixin.py b/msticpy/data/core/query_provider_connections_mixin.py index e859ff773..3c0364667 100644 --- a/msticpy/data/core/query_provider_connections_mixin.py +++ b/msticpy/data/core/query_provider_connections_mixin.py @@ -432,6 +432,57 @@ async def _run_tagged_task( 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], @@ -446,66 +497,42 @@ 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: 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", + results, failed_tasks_ids = ( + await QueryProviderConnectionsMixin._run_tasks_threaded( + thread_tasks, progress=progress ) - 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, + ) + + if retry and failed_tasks_ids: + 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] ) - 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, + for query_id in failed_tasks_ids + } + retry_results, _ = ( + await QueryProviderConnectionsMixin._run_tasks_threaded( + retry_thread_tasks, progress=progress ) - failed_tasks_ids.append(query_id) + ) + # 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 the order tasks completed in. + # 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 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, - ) - ) - if not failed_results.empty: - ordered_results.append(failed_results) if ordered_results: return pd.concat(ordered_results, ignore_index=True) diff --git a/tests/data/test_async_queries.py b/tests/data/test_async_queries.py index 4e559732e..e93d058e0 100644 --- a/tests/data/test_async_queries.py +++ b/tests/data/test_async_queries.py @@ -367,7 +367,13 @@ def _tracked(label: str, should_fail: bool) -> pd.DataFrame: # must only have been called once. check.equal(call_counts["a"], 2) check.equal(call_counts["b"], 1) - check.equal(sorted(result["query_id"]), ["a", "b"]) + # "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():