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

Draft
ennuite wants to merge 19 commits into
base: main
Choose a base branch
from
Draft

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
19 commits
Select commit Hold shift + click to select a range
09e063c
Add compiled files from new proto definition
ennuite Mar 25, 2026
1dacbb2
Extend FlightSQL client and server to have new is_update field in Pre…
ennuite Mar 25, 2026
0dc7515
Add client tests for prepared statements with the new field
ennuite Mar 26, 2026
a92965a
Add test for prepared statement query
ennuite Mar 27, 2026
170b38d
Add test for PreparedStatementUpdate
ennuite Mar 27, 2026
3466969
Remove dead code from test
ennuite Jul 3, 2026
2c8c152
Fix state leak
ennuite Jul 3, 2026
8c7286c
Fail fast in client test when server does not set is_update
ennuite Jul 3, 2026
c6f0a54
Refactor test assertions with booleans
ennuite Jul 3, 2026
0288679
Improve comments
ennuite Jul 13, 2026
25bed59
Remove unrelated comment change
ennuite Jul 13, 2026
fcde3da
Allow Substrait to also use the new field
ennuite Jul 26, 2026
a2d9dc5
Fix comment wrap-around
ennuite Jul 26, 2026
eb16088
Add coverage of unset isUpdate field in client unit test
ennuite Jul 29, 2026
8ba207b
Add coverage of unset isUpdate field in TestPreparedStatementLoadFrom…
ennuite Jul 29, 2026
8ce4339
Add coverage of set isUpdate field when loading prepared statement fr…
ennuite Jul 29, 2026
786adc7
Add end-to-end test without isUpdate hint
ennuite Jul 29, 2026
4816b1b
Add server tests for Substrait
ennuite Aug 7, 2026
c1bc900
Fix inconsistent usage of Require()
ennuite Aug 7, 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.

@ennuite ennuite Jul 29, 2026 •
edited
Loading

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.

  1. I addressed this with 2 tests. In TestPreparedStatementLoadFromResult I let things as they are, and tested with the new field unset, see 8ba207b. Note that this test has the same nuance about val that I pointed out in feat(arrow/flight/sql): Add is_update field to ActionCreatePreparedStatementResult #732 (comment)

I added a new test that checks that the field is properly loaded also when set, see a334b86

  1. For that one I purposefully followed the example set by handle, but I will admit I had my own reservations while doing it.

I will not change this in this PR unless you ask me to, but in your opinion is it worth a follow-up PR to have a defensive copy for both handle and isUpdate? If so I can tackle that after this one

}, 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 •
edited
Loading

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 eb16088

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.Require().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.Require().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 @@ -881,9 +1003,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.

@ennuite ennuite Aug 15, 2026 •
edited
Loading

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.

The comment mentions 2 distinct public structs. For clarity, you are asking about release notes for both of them, right? I updated the PR description with the breaking change, let me know if I'm supposed to add it somewhere else.

The type alias refers to flightsql.CreatePreparedStatementResult. This is the raw proto version, what goes in the wire, and I think it's there so that users of the flightsql package don't need to import the proto package directly. In this codebase, this is only used in LoadPreparedStatementFromResult()

But the line about the unkeyed literal is referring to a different struct flightsql.ActionCreatePreparedStatementResult that is unrelated to the alias. It is defined in server.go, and it is a developer-friendly type with the *arrow.Schema objects decoded from wire format. It is used in quite a few places in the codebase.

Thank you for pointing out the compatibility rules, for reference I found them at https://go.dev/doc/go1compat:

For the addition of features in later point releases, it may be necessary to add fields to exported structs in the API. Code that uses unkeyed struct literals (such as pkg.T{3, "x"}) to create values of these types would fail to compile after such a change. However, code that uses keyed literals (pkg.T{A: 3, B: "x"}) will continue to compile after such a change. We will update such data structures in a way that allows keyed struct literals to remain compatible, although unkeyed literals may fail to compile.

}

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