This is an automated email from the ASF dual-hosted git repository.
zeroshade pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-go.git
The following commit(s) were added to refs/heads/main by this push:
new da6ec021 fix(flight/flightsql): recover real error when prepared DoPut
stream closes early (#920)
da6ec021 is described below
commit da6ec021bd6eeb7453402c026ea26382506f0716
Author: Fredrik Fornwall <[email protected]>
AuthorDate: Fri Jul 10 18:41:50 2026 +0200
fix(flight/flightsql): recover real error when prepared DoPut stream closes
early (#920)
**DISCLAIMERS**:
- I'm not a user of this repository directly - the only reason this was
encountered was due to flakyness in `arrow-adbc`, see
https://github.com/apache/arrow-adbc/pull/4497
- This PR was generated by AI. I have reviewed the code and is ready to
make adjustments or to look into questions. Perhaps make comments less
verbose, or adopt code to project conventions? Let me know!
When a prepared statement's bound-parameter `DoPut` stream is torn down
by the server before the client finishes sending (e.g. the server
handler returns without draining the request stream), gRPC surfaces the
parameter send as a bare `io.EOF`.
Per gRPC's contract an `io.EOF` from a send only means "the stream ended
-- receive to learn why", but `bindParameters` and `ExecuteUpdate`
forwarded that EOF directly, so the caller saw an uninformative "EOF
(Unknown)" instead of the server's actual status.
Treat a send-side `io.EOF` as a signal to fall through to the receive,
which recovers the real error (or the response the server produced
before closing).
Adds a regression test whose server returns `InvalidArgument` without
reading the stream; a large binding forces the send to block on flow
control and observe the teardown, so the test fails deterministically
without the fix.
Signed-off-by: Fredrik Fornwall <[email protected]>
---
arrow/flight/flightsql/client.go | 37 +++++++++----
arrow/flight/flightsql/server_test.go | 99 +++++++++++++++++++++++++++++++++++
2 files changed, 126 insertions(+), 10 deletions(-)
diff --git a/arrow/flight/flightsql/client.go b/arrow/flight/flightsql/client.go
index 192cec7b..e76b4163 100644
--- a/arrow/flight/flightsql/client.go
+++ b/arrow/flight/flightsql/client.go
@@ -1231,7 +1231,10 @@ func (p *PreparedStatement) ExecuteUpdate(ctx
context.Context, opts ...grpc.Call
}
if p.hasBindParameters() {
wr, err = p.writeBindParametersToStream(pstream, desc)
- if err != nil {
+ // A bare io.EOF from a send means the server has already
closed the
+ // stream; the authoritative status is delivered by the Recv
below, so
+ // don't bail out early here (see bindParameters for the full
rationale).
+ if err != nil && !errors.Is(err, io.EOF) {
return
}
} else {
@@ -1239,17 +1242,21 @@ func (p *PreparedStatement) ExecuteUpdate(ctx
context.Context, opts ...grpc.Call
wr = flight.NewRecordWriter(pstream, ipc.WithSchema(schema))
wr.SetFlightDescriptor(desc)
rec := array.NewRecordBatch(schema, []arrow.Array{}, 0)
- if err = wr.Write(rec); err != nil {
+ if err = wr.Write(rec); err != nil && !errors.Is(err, io.EOF) {
return
}
}
- if err = wr.Close(); err != nil {
- return
- }
- if err = pstream.CloseSend(); err != nil {
- return
+ // wr is nil only when writeBindParametersToStream returned an error
above
+ // (e.g. a send failed); otherwise it is the writer to flush and close.
+ if wr != nil {
+ if err = wr.Close(); err != nil && !errors.Is(err, io.EOF) {
+ return
+ }
}
+ // Ignore CloseSend's error: if the server already tore the stream down
the
+ // send side is gone, and the authoritative status is delivered by Recv.
+ _ = pstream.CloseSend()
if res, err = pstream.Recv(); err != nil {
return
}
@@ -1272,11 +1279,21 @@ func (p *PreparedStatement) bindParameters(ctx
context.Context, desc *pb.FlightD
return nil, err
}
wr, err := p.writeBindParametersToStream(pstream, desc)
- if err != nil {
+ // gRPC hands back a bare io.EOF from a send once the server
has closed
+ // the stream; its contract is that the real outcome -- a
genuine error,
+ // or a response the server produced before closing -- must
then be
+ // recovered with a receive. Fall through to
captureDoPutPreparedStatementHandle
+ // (which performs that receive) rather than surfacing the
uninformative
+ // EOF, which otherwise reaches the caller as "EOF (Unknown)".
+ if err != nil && !errors.Is(err, io.EOF) {
return nil, err
}
- if err = wr.Close(); err != nil {
- return nil, err
+ // wr is nil when writeBindParametersToStream returned an error
above
+ // (e.g. a send failed); otherwise it is the writer to flush
and close.
+ if wr != nil {
+ if err = wr.Close(); err != nil && !errors.Is(err,
io.EOF) {
+ return nil, err
+ }
}
pstream.CloseSend()
if err = p.captureDoPutPreparedStatementHandle(pstream); err !=
nil {
diff --git a/arrow/flight/flightsql/server_test.go
b/arrow/flight/flightsql/server_test.go
index ca07eef6..74dd59f7 100644
--- a/arrow/flight/flightsql/server_test.go
+++ b/arrow/flight/flightsql/server_test.go
@@ -857,6 +857,105 @@ func TestBaseServer(t *testing.T) {
suite.Run(t, &FlightSqlServerSessionSuite{sessionManager:
session.NewStatelessServerSessionManager()})
}
+// earlyCloseServer models a Flight SQL server whose
DoPutPreparedStatementQuery
+// handler returns before consuming the client's bound-parameter stream. gRPC
+// tears the request stream down as soon as the handler returns, so the
client's
+// in-flight parameter send races that teardown and observes a bare io.EOF. Per
+// gRPC's contract that io.EOF only signals "the stream ended -- receive to
learn
+// why", so the client must recover the real server status rather than
surfacing
+// the uninformative EOF (which otherwise reaches callers as "EOF (Unknown)").
+type earlyCloseServer struct {
+ flightsql.BaseServer
+}
+
+func (*earlyCloseServer) CreatePreparedStatement(ctx context.Context, req
flightsql.ActionCreatePreparedStatementRequest)
(flightsql.ActionCreatePreparedStatementResult, error) {
+ return flightsql.ActionCreatePreparedStatementResult{
+ Handle: []byte("early-close"),
+ ParameterSchema: arrow.NewSchema([]arrow.Field{
+ {Name: "n", Type: arrow.PrimitiveTypes.Int32, Nullable:
true},
+ }, nil),
+ }, nil
+}
+
+func (*earlyCloseServer) DoPutPreparedStatementQuery(ctx context.Context, cmd
flightsql.PreparedStatementQuery, reader flight.MessageReader, writer
flight.MetadataWriter) ([]byte, error) {
+ // Deliberately return without draining reader so the request stream is
torn
+ // down while the client is still sending its parameters.
+ return nil, status.Error(codes.InvalidArgument, "boom")
+}
+
+func (*earlyCloseServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd
flightsql.PreparedStatementUpdate, reader flight.MessageReader) (int64, error) {
+ // Same as the query handler: return without draining reader so the
client's
+ // parameter send observes the teardown (the ExecuteUpdate path).
+ return 0, status.Error(codes.InvalidArgument, "boom")
+}
+
+func (*earlyCloseServer) ClosePreparedStatement(ctx context.Context, req
flightsql.ActionClosePreparedStatementRequest) error {
+ return nil
+}
+
+// TestPreparedStatementDoPutEarlyClose is a regression test for a prepared
+// statement whose bound-parameter DoPut is torn down early by the server: the
+// client must surface the server's real error instead of a bare io.EOF. Both
+// the query path (Execute -> bindParameters) and the sibling update path
+// (ExecuteUpdate) carry the same fix, so each is exercised as a subtest.
+func TestPreparedStatementDoPutEarlyClose(t *testing.T) {
+ srv := flight.NewServerWithMiddleware(nil)
+
srv.RegisterFlightService(flightsql.NewFlightServer(&earlyCloseServer{}))
+ require.NoError(t, srv.Init("localhost:0"))
+ go func() { _ = srv.Serve() }()
+ defer srv.Shutdown()
+
+ cl, err := flightsql.NewClient(srv.Addr().String(), nil, nil,
dialOpts...)
+ require.NoError(t, err)
+ defer cl.Close()
+
+ ctx := context.Background()
+
+ // run prepares a statement, binds a deliberately large parameter set
(well
+ // beyond gRPC's flow-control window) and drives exec. Because the
server
+ // returns without reading, the window never opens, so the client's
+ // parameter send blocks and then observes the stream teardown as a bare
+ // io.EOF -- deterministically exercising the path that previously
leaked
+ // "EOF (Unknown)" to the caller. The server's real status must surface.
+ run := func(t *testing.T, exec func(*flightsql.PreparedStatement)
error) {
+ prep, err := cl.Prepare(ctx, "SELECT 1")
+ require.NoError(t, err)
+ defer prep.Close(ctx)
+
+ bldr := array.NewRecordBuilder(memory.DefaultAllocator,
prep.ParameterSchema())
+ defer bldr.Release()
+ fb := bldr.Field(0).(*array.Int32Builder)
+ fb.Reserve(1 << 21)
+ for i := 0; i < 1<<21; i++ {
+ fb.Append(int32(i))
+ }
+ rec := bldr.NewRecordBatch()
+ defer rec.Release()
+ prep.SetParameters(rec)
+
+ err = exec(prep)
+ require.Error(t, err)
+ st := status.Convert(err)
+ require.Equalf(t, codes.InvalidArgument, st.Code(),
+ "expected the server's status to be surfaced, not a
bare EOF; got: %v", err)
+ require.Contains(t, st.Message(), "boom")
+ }
+
+ t.Run("Execute", func(t *testing.T) {
+ run(t, func(p *flightsql.PreparedStatement) error {
+ _, err := p.Execute(ctx)
+ return err
+ })
+ })
+
+ t.Run("ExecuteUpdate", func(t *testing.T) {
+ run(t, func(p *flightsql.PreparedStatement) error {
+ _, err := p.ExecuteUpdate(ctx)
+ return err
+ })
+ })
+}
+
func TestStatefulServerSessionCookies(t *testing.T) {
// Generate session IDs deterministically
sessionIDGenerator := func(ids []string) func() string {
