diff --git a/msticpy/data/core/query_provider_connections_mixin.py b/msticpy/data/core/query_provider_connections_mixin.py index cc64098c..3c036466 100644 --- a/msticpy/data/core/query_provider_connections_mixin.py +++ b/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() diff --git a/tests/data/test_async_queries.py b/tests/data/test_async_queries.py index 18b8703f..e93d058e 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,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)