Skip to content

feat(arrow/flight/sql): Add is_update field to ActionCreatePreparedStatementResult #732

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ennuite wants to merge 16 commits into
base: main
Choose a base branch
from
+1,064 −1,638
Open

feat(arrow/flight/sql): Add is_update field to ActionCreatePreparedStatementResult #732

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
16 commits
Select commit Hold shift + click to select a range
b3ce6c4
Add compiled files from new proto definition
ennuite Mar 25, 2026
cf15f96
Extend FlightSQL client and server to have new is_update field in Pre…
ennuite Mar 25, 2026
8651cbd
Add client tests for prepared statements with the new field
ennuite Mar 26, 2026
d3595fd
Add test for prepared statement query
ennuite Mar 27, 2026
c378e8e
Add test for PreparedStatementUpdate
ennuite Mar 27, 2026
c38caff
Remove dead code from test
ennuite Jul 3, 2026
140b575
Fix state leak
ennuite Jul 3, 2026
21ab0bb
Fail fast in client test when server does not set is_update
ennuite Jul 3, 2026
bd61a0e
Refactor test assertions with booleans
ennuite Jul 3, 2026
24e2862
Improve comments
ennuite Jul 13, 2026
7437f87
Remove unrelated comment change
ennuite Jul 13, 2026
f6a8588
Allow Substrait to also use the new field
ennuite Jul 26, 2026
9670cca
Fix comment wrap-around
ennuite Jul 26, 2026
a324e27
Add coverage of unset isUpdate field in client unit test
ennuite Jul 29, 2026
514d3c6
Add coverage of unset isUpdate field in TestPreparedStatementLoadFrom…
ennuite Jul 29, 2026
a334b86
Add coverage of set isUpdate field when loading prepared statement fr…
ennuite Jul 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 17 additions & 3 deletions arrow/flight/flightsql/client.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -566,6 +566,7 @@ func (c *Client) LoadPreparedStatementFromResult(result *CreatePreparedStatement
handle: result.PreparedStatementHandle,
datasetSchema: dsSchema,
paramSchema: paramSchema,
isUpdate: result.IsUpdate,

@zeroshade zeroshade Jul 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two small things on this propagation site.

  1. Untested. TestPreparedStatementLoadFromResult (client_test.go:971) is the only coverage for LoadPreparedStatementFromResult, and it doesn't assert IsUpdate. Since that test already constructs the result and checks the other two propagated fields (ParameterSchema, DatasetSchema), adding an IsUpdate assertion there would be consistent and nearly free. (Couldn't leave this as an inline comment on that line — it's outside the diff.)

  2. Nit, no change required: this is the one place that stores a *bool owned by the caller, so a caller mutating result.IsUpdate after this call would change the statement's view. Prepare/PrepareSubstrait don't have this property since they point at their own unmarshaled local. Not worth a defensive copy unless you'd want it for handle too — just noting it so the asymmetry is a conscious choice.

}, nil
}

Expand Down Expand Up @@ -616,6 +617,7 @@ func parsePreparedStatementResponse(c *Client, mem memory.Allocator, results pb.
handle: message.PreparedStatementHandle,
datasetSchema: dsSchema,
paramSchema: paramSchema,
isUpdate: message.IsUpdate,
}, nil
}

Expand Down Expand Up @@ -1104,14 +1106,15 @@ func (tx *Txn) RollbackSavepoint(ctx context.Context, sp Savepoint, opts ...grpc
// and maintains a reference to the Client that created it along with the
// prepared statement handle.
//
// If the server returned the Dataset Schema or Parameter Binding schemas
// at creation, they will also be accessible from this object. Close
// should be called when no longer needed.
// If the server returned the Dataset Schema, Parameter Binding schemas,
// or an IsUpdate hint at creation, they will also be accessible from this object.
// Close should be called when no longer needed.
type PreparedStatement struct {
client *Client
handle []byte
datasetSchema *arrow.Schema
paramSchema *arrow.Schema
isUpdate *bool
paramBinding arrow.RecordBatch
streamBinding array.RecordReader
closed bool
Expand Down Expand Up @@ -1365,6 +1368,17 @@ func (p *PreparedStatement) DatasetSchema() *arrow.Schema { return p.datasetSche
// the prepared statement.
func (p *PreparedStatement) ParameterSchema() *arrow.Schema { return p.paramSchema }

// IsUpdate reports the server's hint for how to execute this prepared statement.
// val is true if the server indicated an update, false if it indicated a query.
// ok is false if the server did not provide a hint, in which case the client
// can choose how to execute the statement.
func (p *PreparedStatement) IsUpdate() (val bool, ok bool) {
Comment thread
ennuite marked this conversation as resolved.
if p.isUpdate == nil {
return false, false
}
return *p.isUpdate, true
}

// The handle associated with this PreparedStatement
func (p *PreparedStatement) Handle() []byte { return p.handle }

Expand Down
144 changes: 144 additions & 0 deletions arrow/flight/flightsql/client_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -406,6 +406,8 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecute() {
defer prepared.Close(context.TODO(), s.callOpts...)

s.Equal(string(prepared.Handle()), "query")
_, ok := prepared.IsUpdate()
s.False(ok)

info, err := prepared.Execute(context.TODO(), s.callOpts...)
s.NoError(err)
Expand All @@ -416,13 +418,133 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecute() {

secondPrepare := flightsql.NewPreparedStatement(&s.sqlClient, prepared.Handle())
s.Equal(string(secondPrepare.Handle()), "query")
_, ok = secondPrepare.IsUpdate()
s.False(ok)

defer secondPrepare.Close(context.TODO(), s.callOpts...)

info, err = secondPrepare.Execute(context.TODO(), s.callOpts...)
s.NoError(err)
s.Equal(&emptyFlightInfo, info)
}

func (s *FlightSqlClientSuite) TestPreparedStatementExecuteWithIsUpdateFalse() {

@zeroshade zeroshade Jul 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add a test for the unset/nil case. This is my only blocking request.

Both new client tests here, and both new server tests, assert ok == true. Nothing anywhere asserts ok == false. That's the path taken by every server that doesn't implement the new field — i.e. the backward-compatibility contract — and it's currently the only untested branch of the getter.

TestPreparedStatementExecute above (line 368) already builds an ActionCreatePreparedStatementResult with no IsUpdate, so this can be as small as adding two lines there:

val, ok := prepared.IsUpdate()
s.False(ok)
s.False(val)

Worth locking down explicitly because a future refactor that switched isUpdate *bool to a plain bool would silently turn "no hint" into "is a query" — which is a real behavior change for clients that branch on this — and no existing test would catch it.

@ennuite ennuite Jul 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do we want the s.False(val) check? If s.False(ok), then val is meaningless.

@ennuite ennuite Jul 29, 2026

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I addressed this in a324e27

I do not think it makes sense to have the check in val: it being False seems like an implementation detail, and not a part of the contract of the IsUpdate() method. Let me know if you disagree: I can change it to look exactly like your suggestion.

const query = "query"

cmd := &pb.ActionCreatePreparedStatementRequest{Query: query}
action := getAction(cmd)
action.Type = flightsql.CreatePreparedStatementActionType
closeAct := getAction(&pb.ActionClosePreparedStatementRequest{PreparedStatementHandle: []byte(query)})
closeAct.Type = flightsql.ClosePreparedStatementActionType

isUpdate := false
result := &pb.ActionCreatePreparedStatementResult{
PreparedStatementHandle: []byte(query), // handle is the query string for simplicity
IsUpdate: &isUpdate,
}
var out anypb.Any
out.MarshalFrom(result)
data, _ := proto.Marshal(&out)

createRsp := &mockDoActionClient{}
defer createRsp.AssertExpectations(s.T())
createRsp.On("Recv").Return(&pb.Result{Body: data}, nil).Once()
createRsp.On("Recv").Return(&pb.Result{}, io.EOF).Once()
createRsp.On("CloseSend").Return(nil).Once()

closeRsp := &mockDoActionClient{}
defer closeRsp.AssertExpectations(s.T())
closeRsp.On("Recv").Return(&pb.Result{}, io.EOF)
closeRsp.On("CloseSend").Return(nil)

s.mockClient.On("DoAction", flightsql.CreatePreparedStatementActionType, action.Body, s.callOpts).
Return(createRsp, nil).Once()
s.mockClient.On("DoAction", flightsql.ClosePreparedStatementActionType, closeAct.Body, s.callOpts).
Return(closeRsp, nil)

infoCmd := &pb.CommandPreparedStatementQuery{PreparedStatementHandle: []byte(query)}
desc := getDesc(infoCmd)
s.mockClient.On("GetFlightInfo", desc.Type, desc.Cmd, s.callOpts).Return(&emptyFlightInfo, nil).Once()

prepared, err := s.sqlClient.Prepare(context.TODO(), query, s.callOpts...)
s.NoError(err)
defer prepared.Close(context.TODO(), s.callOpts...)

s.Equal(string(prepared.Handle()), query)
val, ok := prepared.IsUpdate()
s.Require().True(ok)
s.False(val)

info, err := prepared.Execute(context.TODO(), s.callOpts...)
s.NoError(err)
s.Equal(&emptyFlightInfo, info)
}

func (s *FlightSqlClientSuite) TestPreparedStatementExecuteUpdateWithIsUpdateTrue() {
const query = "DML query"

cmd := &pb.ActionCreatePreparedStatementRequest{Query: query}
action := getAction(cmd)
action.Type = flightsql.CreatePreparedStatementActionType
closeAct := getAction(&pb.ActionClosePreparedStatementRequest{PreparedStatementHandle: []byte(query)})
closeAct.Type = flightsql.ClosePreparedStatementActionType

// Set is_update to true
isUpdate := true
result := &pb.ActionCreatePreparedStatementResult{
PreparedStatementHandle: []byte(query),
IsUpdate: &isUpdate,
}
var out anypb.Any
out.MarshalFrom(result)
data, _ := proto.Marshal(&out)

createRsp := &mockDoActionClient{}
defer createRsp.AssertExpectations(s.T())
createRsp.On("Recv").Return(&pb.Result{Body: data}, nil).Once()
createRsp.On("Recv").Return(&pb.Result{}, io.EOF).Once()
createRsp.On("CloseSend").Return(nil).Once()

closeRsp := &mockDoActionClient{}
defer closeRsp.AssertExpectations(s.T())
closeRsp.On("Recv").Return(&pb.Result{}, io.EOF)
closeRsp.On("CloseSend").Return(nil)

s.mockClient.On("DoAction", flightsql.CreatePreparedStatementActionType, action.Body, s.callOpts).
Return(createRsp, nil).Once()
s.mockClient.On("DoAction", flightsql.ClosePreparedStatementActionType, closeAct.Body, s.callOpts).
Return(closeRsp, nil)

// Mock DoPut for ExecuteUpdate
updateCmd := &pb.CommandPreparedStatementUpdate{PreparedStatementHandle: []byte(query)}
updateDesc := getDesc(updateCmd)
updateResult := &pb.DoPutUpdateResult{RecordCount: 1}
resdata, _ := proto.Marshal(updateResult)

mockedPut := &mockDoPutClient{}
defer mockedPut.AssertExpectations(s.T())
mockedPut.On("Send", mock.MatchedBy(func(fd *flight.FlightData) bool {
return proto.Equal(updateDesc, fd.FlightDescriptor)
})).Return(nil)
mockedPut.On("CloseSend").Return(nil)
mockedPut.On("Recv").Return(&pb.PutResult{AppMetadata: resdata}, nil)
s.mockClient.On("DoPut", s.callOpts).Return(mockedPut, nil)

prepared, err := s.sqlClient.Prepare(context.TODO(), query, s.callOpts...)
s.NoError(err)
defer prepared.Close(context.TODO(), s.callOpts...)

s.Equal(string(prepared.Handle()), query)
val, ok := prepared.IsUpdate()
s.Require().True(ok)
s.True(val)

// Execute as update
num, err := prepared.ExecuteUpdate(context.TODO(), s.callOpts...)
s.NoError(err)
s.EqualValues(1, num)
}

func (s *FlightSqlClientSuite) TestPreparedStatementExecuteParamBinding() {
const query = "query"
const handle = "handle"
Expand Down Expand Up @@ -878,9 +1000,31 @@ func (s *FlightSqlClientSuite) TestPreparedStatementLoadFromResult() {
s.NoError(err)
defer datasetRec.Release()

_, ok := prepared.IsUpdate()
s.False(ok)

s.Equal(string(prepared.Handle()), "query")
}

func (s *FlightSqlClientSuite) TestPreparedStatementLoadFromResultWithIsUpdate() {
const query = "query"

isUpdate := true
result := &pb.ActionCreatePreparedStatementResult{
PreparedStatementHandle: []byte(query),
IsUpdate: &isUpdate,
}

prepared, err := s.sqlClient.LoadPreparedStatementFromResult(result)
s.NoError(err)

s.Equal(string(prepared.Handle()), "query")

val, ok := prepared.IsUpdate()
s.True(ok)
s.True(val)
}

func TestFlightSqlClient(t *testing.T) {
suite.Run(t, new(FlightSqlClientSuite))
}
5 changes: 5 additions & 0 deletions arrow/flight/flightsql/server.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -155,6 +155,9 @@ type ActionCreatePreparedStatementResult struct {
Handle []byte
DatasetSchema *arrow.Schema
ParameterSchema *arrow.Schema
// IsUpdate indicates whether the prepared statement should be executed
// as an update (true) or query (false). If nil, the client can choose.
IsUpdate *bool

@zeroshade zeroshade Jul 27, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not a defect, but this needs a release note, because the public API impact is wider than this struct.

types.go:882 has:

type CreatePreparedStatementResult = pb.ActionCreatePreparedStatementResult

Being a type alias, the regenerated proto field lands on a public flightsql symbol too. So this PR adds an exported field to two publicly-reachable structs. Keyed composite literals are unaffected (and that's what I found in the downstream usage I spot-checked), but any downstream unkeyed literal like:

return flightsql.ActionCreatePreparedStatementResult{handle, ds, ps}, nil

will stop compiling. That's permitted under the Go 1 compatibility rules — unkeyed literals of imported struct types are explicitly not covered — and go vet's composites check flags them, so real-world exposure should be low. I don't think it should block the PR, but it's worth a line in the changelog rather than letting someone discover it at upgrade time.

The unconditional result.IsUpdate = output.IsUpdate on both the statement and Substrait branches is correct, by the way — both sides are *bool, so nil flows through and no presence check is needed. Good that the Substrait path wasn't forgotten.

}

type ActionBeginTransactionRequest interface{}
Expand Down Expand Up @@ -1251,6 +1254,7 @@ func (f *flightSqlServer) DoAction(cmd *flight.Action, stream flight.FlightServi
if output.ParameterSchema != nil {
result.ParameterSchema = flight.SerializeSchema(output.ParameterSchema, f.mem)
}
result.IsUpdate = output.IsUpdate

if err := anycmd.MarshalFrom(&result); err != nil {
return status.Errorf(codes.Internal, "unable to marshal final response: %s", err.Error())
Expand Down Expand Up @@ -1286,6 +1290,7 @@ func (f *flightSqlServer) DoAction(cmd *flight.Action, stream flight.FlightServi
if output.ParameterSchema != nil {
result.ParameterSchema = flight.SerializeSchema(output.ParameterSchema, f.mem)
}
result.IsUpdate = output.IsUpdate

if err := anycmd.MarshalFrom(&result); err != nil {
return status.Errorf(codes.Internal, "unable to marshal final response: %s", err.Error())
Expand Down
Loading