Fix kernel async statement telemetry handle

[jay-xiao446]

Summary

• Use the retained owning kernel async handle for the first full in-process async flow: status polling via get_query_state and the first result fetch via get_execution_result.
• Mark result streaming once the owning handle is claimed for fetch, then fall back to attach-by-ID for later state/result calls.
• Preserve attach-by-ID for missing-owner cases such as re-fetch, fresh cursor, or cross-process lookup.
• Update unit coverage and lifecycle comments for the owning-handle versus attach-by-ID behavior.

Testing

• .venv/bin/python -m pytest tests/unit/test_kernel_client.py -q

[peco-review-bot]

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().

[peco-review-bot]

🔵 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.

[peco-review-bot]

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.

[peco-review-bot]

🔵 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.)

[peco-review-bot]

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().

[peco-review-bot]

🔵 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.)

[peco-review-bot]

✅ No issues identified by the review bot.

[peco-review-bot]

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.

[peco-review-bot]

🟡 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.)

[peco-review-bot]

Verdict: 2 Low

Looks good overall — the owning-handle-first / attach-by-id-fallback logic, marker lifecycle (add on claim, discard in close_command, clear in close_session), and the await_result rollback are coherent and well-covered by the new unit tests. Two low-severity notes: an asymmetric marker rollback on result-set construction failure that can silently lose the async telemetry the PR aims to preserve, and newly-introduced concurrent sharing of a single owning-handle object across get_query_state/get_execution_result.

[peco-review-bot]

🔵 Low — The claimed-marker rollback is asymmetric between the two failure points, which partially defeats this PR's telemetry goal.

When the owning handle is claimed, _async_result_stream_started is add()ed before await_result(). The first try (786) correctly discards the marker if await_result() fails, so a retry can reuse the telemetry-bearing owning handle (covered by test_get_execution_result_owning_handle_failure_can_retry_owning_handle).

But the second try (795) — _make_result_set(stream, ...), which calls arrow_schema() and can raise (see the comment at 792-593) — does not discard the marker on failure. If construction fails after the owning handle was claimed, the marker stays set, so any retry of get_execution_result for this id falls into the attach-by-id path and drains the result on a no-op-telemetry attached handle — losing the SEA async statement row this PR exists to preserve. The owning handle is still live in _async_handles, so it could be reused.

If leaving the marker set here is intentional (i.e. await_result() succeeding is considered the point telemetry is finalized on the owning handle, making a later attach acceptable), a one-line comment on the second try would make that explicit; otherwise the discard should mirror the first block.

(Anchored to the nearest changed line — see the description for the exact location.)

[peco-review-bot]

🔵 Low — New cross-call sharing of a single owning-handle object. Before this change, both get_query_state and get_execution_result always attached a fresh handle by id, so they never operated on the same kernel handle object concurrently. Now, until the fetch marker is set, both methods return the same object from self._async_handles.get(...) and call handle.status() / handle.await_result() on it outside _async_handles_lock.

For the normal per-cursor sequential CUJ (poll then fetch) this is fine, but _async_handles is keyed per-connection by guid, so a second cursor that adopts the same command id (or a background poll racing the fetch) can drive status() on the owning handle while await_result() runs on the same object. Whether the kernel handle tolerates concurrent method calls isn't visible from the connector; if it doesn't, this is a latent data race that the old attach-by-id path avoided. Worth a note confirming the kernel handle is safe for concurrent status()/await_result(), or gating the shared use.

(Anchored to the nearest changed line — see the description for the exact location.)