fix(query): record real row count in registry for Arrow path (#333)#376
fix(query): record real row count in registry for Arrow path (#333)#376
Conversation
Complete(queryID, 0) was called synchronously after arrowJSONQueryFunc returned, but the Arrow path streams via SetBodyStreamWriter — the actual row count from streamArrowJSON is only available inside the async callback. This caused the query registry to always record 0 rows for Arrow-path queries. Fix: add onComplete func(int) to executeArrowJSONQuery and the arrowJSONQueryFunc type. Call site 1 passes a closure that calls h.queryRegistry.Complete(queryID, rc) with the real count; call site 2 (measurement handler, no registry) passes nil. The callback invokes onComplete(rc) after streamArrowJSON returns.
The previous fix only covered the success path. The post-implementation review found that error and timeout branches in executeArrowJSONQuery returned handled=true without notifying the registry, leaving queries stuck in "running" state permanently. Changes: - Add onFail func(string) and onTimeout func() alongside onComplete - "no files found" path calls onComplete(0) — it is a success with 0 rows - Timeout path calls onTimeout() - Query error path calls onFail(err.Error()) - Update arrowJSONQueryFunc type and both call sites accordingly - Update RELEASE_NOTES_2026.05.1.md
There was a problem hiding this comment.
Code Review
This pull request updates the Arrow-native query path to accurately report row counts by introducing an onComplete callback. This allows the query registry to be updated with the actual row count from the asynchronous stream writer instead of a placeholder value. A review comment identifies a potential nil pointer dereference in the onComplete closure due to its asynchronous execution and suggests capturing the registry in a local variable to ensure thread safety.
| var onTimeout func() | ||
| if h.queryRegistry != nil && queryID != "" { |
There was a problem hiding this comment.
The onComplete closure captures h.queryRegistry directly. Since this closure is executed asynchronously by SetBodyStreamWriter, there's a potential for h.queryRegistry to become nil between the time the closure is defined and when it's executed, leading to a nil pointer dereference. It's safer to capture the queryRegistry value into a local variable and perform a nil check inside the closure.
| var onTimeout func() | |
| if h.queryRegistry != nil && queryID != "" { | |
| registry := h.queryRegistry | |
| id := queryID | |
| onComplete = func(rc int) { | |
| if registry != nil { | |
| registry.Complete(id, rc) | |
| } | |
| } |
Summary
Complete(queryID, 0)was called synchronously right afterarrowJSONQueryFuncreturned, but the Arrow path streams viaSetBodyStreamWriter— the real row count is only available inside the async callback. The query registry always recorded 0 for Arrow-path queries.onComplete func(int)parameter toexecuteArrowJSONQueryand thearrowJSONQueryFunctype. The callback now callsonComplete(rc)afterstreamArrowJSONcompletes, delivering the real count to the registry.Test plan
GOFLAGS="-tags=duckdb_arrow" go build ./internal/...passesgo test ./internal/api/... -tags=duckdb_arrowpasses/api/v1/queries/{id}—row_countshould match actual rows returned, not 0Closes #333