Fix kernel async statement telemetry handle - #923
Conversation
410999a to
7deecf2
Compare
7deecf2 to
6629d87
Compare
![peco-review-bot[bot]](https://avatars.githubusercontent.com/in/3801177?s=60&v=4){kind=small}
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the owning-handle-first / attach-by-id-fallback change is correct: the claim is check-and-set under _async_handles_lock, cleanup is consistent across close_command/close_session, and the new unit + e2e tests cover the main paths. One low-severity note on an unclear telemetry-finalization edge when result-set construction fails after a successful await_result().
| stream = handle.await_result() | ||
| except Exception as exc: | ||
| if uses_owning_handle: | ||
| with self._async_handles_lock: |
There was a problem hiding this comment.
🔵 Low — The owning-handle failure path only discards _async_result_stream_started when await_result() raises. If await_result() succeeds (marker stays set) but the subsequent KernelResultSet.__init__ → arrow_schema() raises and is re-wrapped, the guid remains marked as started. A later retry then takes the attach-by-id (no-op telemetry) branch.
Whether this loses the ExecuteStatementAsync telemetry row depends on when the kernel finalizes it: if finalization happens when await_result() returns, this is harmless (telemetry already committed). If finalization only completes once the result stream is drained, the telemetry is lost on this retry because the owning handle is never reused. The PR's own comments ("first in-process result stream", "clear the claimed marker so a retry can still use the telemetry-bearing owning handle") are ambiguous on this point, and the added test_get_execution_result_owning_handle_failure_can_retry_owning_handle only exercises the await_result()-raises case, not the construct-failure-after-await case. Worth confirming the finalization semantics and, if drain-based, discarding the marker on the construction-failure path too.
There was a problem hiding this comment.
Verdict: 1 Low
Solid, well-tested fix. The owning-handle-first / attach-by-id-fallback logic is correctly lock-guarded, the claimed marker is reset on await failure to preserve retry-with-owning-handle, and unit + e2e coverage and lifecycle comments are updated to match. One low-severity edge: the result-set-construction failure path (second except) doesn't reset _async_result_stream_started, so a retry after that failure loses the telemetry-bearing owning handle — see inline note. The is_telemetry_enabled short-circuit for use_kernel (overriding even force_enable_telemetry) is intentional and tested.
| except Exception as exc: | ||
| if uses_owning_handle: | ||
| with self._async_handles_lock: | ||
| self._async_result_stream_started.discard(command_id.guid) |
There was a problem hiding this comment.
🔵 Low — The retry-preservation logic is asymmetric. _async_result_stream_started is added under the lock before await_result() runs, and the first except (await failure) correctly discards the marker so a retry can reuse the telemetry-bearing owning handle. But if await_result() succeeds and the subsequent KernelResultSet(...) construction raises (the arrow_schema() path this second except exists to map), the marker stays set and the owning handle stays in _async_handles. A retry then takes the attach-by-id fallback — the exact telemetry-loss case this PR fixes — because the owning result stream was never actually drained on the failed attempt.
This is narrow (only when arrow_schema()/result-set construction fails after a successful await_result), and resetting here is itself debatable since re-awaiting the same owning handle may not be safe. Worth a comment noting the intentional gap, or handling it consistently with the first except.
(Anchored to the nearest changed line — see the description for the exact location.)
There was a problem hiding this comment.
Verdict: 1 Low
Looks good — the async owning-handle ownership model (owning handle for first in-process status/result, attach-by-id for re-fetch/cross-process, marker cleared on close/session-close) is coherent and well covered by the new unit tests, and the use_kernel telemetry short-circuit is intentional (verified against the added test). One Low note about asymmetric marker cleanup when result-set construction fails after a successful await_result().
| except Exception as exc: | ||
| if uses_owning_handle: | ||
| with self._async_handles_lock: | ||
| self._async_result_stream_started.discard(command_id.guid) |
There was a problem hiding this comment.
🔵 Low — Asymmetric cleanup of the _async_result_stream_started marker on the owning-handle path.
In get_execution_result, when the owning handle is used, the guid is added to _async_result_stream_started under the lock (line 782). If handle.await_result() raises, the marker is correctly discarded (lines 789-790) so a retry can still use the telemetry-bearing owning handle — this is exactly the behavior the new test_get_execution_result_owning_handle_failure_can_retry_owning_handle locks in.
However, if await_result() succeeds but the following KernelResultSet construction / arrow_schema() call raises (the try/except at 794-797 that this PR leaves in place), the marker is not discarded. Per this PR's own reasoning, kernel async telemetry is finalized when the result is drained — which happens during fetch, after the result set is constructed. So on a construction failure the drain never runs, yet the marker is now sticky: a subsequent get_async_execution_result() retry falls back to attach-by-id (no-op telemetry), losing the ExecuteStatementAsync telemetry row this change exists to preserve.
This is a narrow edge (arrow_schema raising after await succeeds) and telemetry is best-effort, hence Low. If the intent is that a post-await construction failure should still allow a telemetry-preserving retry, discard the marker in the 796-797 handler too (symmetric with the await handler).
(Anchored to the nearest changed line — see the description for the exact location.)
There was a problem hiding this comment.
Verdict: 1 Medium
Looks solid overall — the owning-handle-first / attach-by-id fallback logic, locking, and lifecycle cleanup (close_command / close_session) are consistent and well-tested. One medium concern: the result-set-construction failure path in get_execution_result doesn't clear the _async_result_stream_started marker the way the sibling await_result failure path does, so a retry after a construction error would silently drop to attach-by-id and lose the async telemetry this PR is meant to preserve.
| except Exception as exc: | ||
| if uses_owning_handle: | ||
| with self._async_handles_lock: | ||
| self._async_result_stream_started.discard(command_id.guid) |
There was a problem hiding this comment.
🟡 Medium — The result-set construction failure path leaves _async_result_stream_started set, defeating the retry-with-owning-handle guarantee this PR establishes.
In get_execution_result, the owning handle is claimed by adding command_id.guid to _async_result_stream_started (L782) before the result is materialized. There are then two try blocks:
await_result()(L783–791): on failure,if uses_owning_handle:discards the marker (L788–790), so a retry re-selects the owning handle and preserves telemetry._make_result_set(...)→KernelResultSet.__init__→arrow_schema()(L794–797): on failure, the marker is not discarded.
Per the method's own comment, kernel async telemetry is finalized on the owning handle only when the result stream is drained (during later fetch), not at await_result() time. So if await_result() succeeds but arrow_schema() raises during construction, the caller never receives a usable ResultSet, yet the marker stays set (and the owning handle is still in _async_handles, since this method uses .get() not .pop()). A subsequent retry then hits the attach-by-id branch and loses the SEA async-statement telemetry row — exactly the regression this PR aims to prevent, and inconsistent with the sibling await_result handling.
Consider discarding the marker on this path too (mirroring L788–790), e.g. wrap the construction failure with the same if uses_owning_handle: cleanup. Note test_get_execution_result_owning_handle_failure_can_retry_owning_handle only covers the await_result failure path; the construction-failure retry case is untested.
(Anchored to the nearest changed line — see the description for the exact location.)
2 participants
Summary
get_query_stateand the first result fetch viaget_execution_result.Testing
.venv/bin/python -m pytest tests/unit/test_kernel_client.py -q