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

arrow/flight/flightsql/client.go

@@ -566,6 +566,7 @@ func (c *Client) LoadPreparedStatementFromResult(result *CreatePreparedStatement
 		handle:        result.PreparedStatementHandle,
 		datasetSchema: dsSchema,
 		paramSchema:   paramSchema,
+		isUpdate:      result.IsUpdate,
 	}, nil
 }
 
@@ -616,6 +617,7 @@ func parsePreparedStatementResponse(c *Client, mem memory.Allocator, results pb.
 		handle:        message.PreparedStatementHandle,
 		datasetSchema: dsSchema,
 		paramSchema:   paramSchema,
+		isUpdate:      message.IsUpdate,
 	}, nil
 }
 
@@ -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
@@ -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) {
+	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 }
 

arrow/flight/flightsql/client_test.go

@@ -423,6 +423,123 @@ func (s *FlightSqlClientSuite) TestPreparedStatementExecute() {
 	s.Equal(&emptyFlightInfo, info)
 }
 
+func (s *FlightSqlClientSuite) TestPreparedStatementExecuteWithIsUpdateFalse() {
+	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"

arrow/flight/flightsql/server.go

@@ -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
 }
 
 type ActionBeginTransactionRequest interface{}
@@ -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())
@@ -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())

arrow/flight/flightsql/server_test.go

@@ -133,6 +133,74 @@ func (*testServer) DoGetStatement(ctx context.Context, ticket flightsql.Statemen
 	return
 }
 
+func (*testServer) CreatePreparedStatement(ctx context.Context, req flightsql.ActionCreatePreparedStatementRequest) (result flightsql.ActionCreatePreparedStatementResult, err error) {
+	query := req.GetQuery()
+	result.Handle = []byte(query)
+	switch query {
+	case "prepared query":
+		isUpdate := false
+		result.IsUpdate = &isUpdate
+	case "prepared update":
+		isUpdate := true
+		result.IsUpdate = &isUpdate
+	default:
+		err = fmt.Errorf("unknown query: %s", query)
+	}
+	return
+}
+
+func (*testServer) GetFlightInfoPreparedStatement(ctx context.Context, q flightsql.PreparedStatementQuery, fd *flight.FlightDescriptor) (*flight.FlightInfo, error) {
+	return &flight.FlightInfo{
+		FlightDescriptor: fd,
+		Endpoint: []*flight.FlightEndpoint{{
+			Ticket: &flight.Ticket{Ticket: fd.Cmd},
+		}},
+	}, nil
+}
+
+func (*testServer) DoGetPreparedStatement(ctx context.Context, q flightsql.PreparedStatementQuery) (sc *arrow.Schema, cc <-chan flight.StreamChunk, err error) {
+	handle := string(q.GetPreparedStatementHandle())
+	switch handle {
+	case "prepared query":
+		b := array.NewInt16Builder(memory.DefaultAllocator)
+		sc = arrow.NewSchema([]arrow.Field{{
+			Name:     "t1",
+			Type:     b.Type(),
+			Nullable: true,
+		}}, nil)
+		b.AppendNull()
+		c := make(chan flight.StreamChunk, 2)
+		c <- flight.StreamChunk{
+			Data: array.NewRecordBatch(sc, []arrow.Array{b.NewArray()}, 1),
+		}
+		b.Append(1)
+		c <- flight.StreamChunk{
+			Data: array.NewRecordBatch(sc, []arrow.Array{b.NewArray()}, 1),
+		}
+		close(c)
+		cc = c
+	default:
+		err = fmt.Errorf("unknown prepared statement handle: %s", handle)
+	}
+	return
+}
+
+func (*testServer) ClosePreparedStatement(ctx context.Context, req flightsql.ActionClosePreparedStatementRequest) error {
+	return nil
+}
+
+func (*testServer) DoPutPreparedStatementUpdate(ctx context.Context, cmd flightsql.PreparedStatementUpdate, rdr flight.MessageReader) (int64, error) {
+	handle := string(cmd.GetPreparedStatementHandle())
+	switch handle {
+	case "prepared update":
+		// Simulates executing a DML query
+		// that affects 2 rows
+		return 2, nil
+	default:
+		return 0, fmt.Errorf("unknown prepared statement handle: %s", handle)
+	}
+}
+
 func (*testServer) SetSessionOptions(ctx context.Context, req *flight.SetSessionOptionsRequest) (*flight.SetSessionOptionsResult, error) {
 	session, err := session.GetSessionFromContext(ctx)
 	if err != nil {
@@ -260,9 +328,9 @@ func (s *FlightSqlServerSuite) TestExecute() {
 		for i := 0; i < data.Len(); i++ {
 			switch n {
 			case 0:
-				s.Assert().Equal(true, data.IsNull(i))
+				s.True(data.IsNull(i))
 			case 1:
-				s.Assert().Equal(false, data.IsNull(i))
+				s.False(data.IsNull(i))
 				s.Assert().Equal(int16(1), data.Value(i))
 			}
 			n++
@@ -287,6 +355,69 @@ func (s *FlightSqlServerSuite) TestExecuteChunkError() {
 	}
 }
 
+func (s *FlightSqlServerSuite) TestExecutePreparedStatementQuery() {
+	prep, err := s.cl.Prepare(context.TODO(), "prepared query")
+	s.Require().NoError(err)
+	defer prep.Close(context.TODO())
+
+	val, ok := prep.IsUpdate()
+	s.Require().True(ok)
+	s.False(val)
+
+	fi, err := prep.Execute(context.TODO())
+	s.Require().NoError(err)
+	ep := fi.GetEndpoint()
+	s.Require().Len(ep, 1)
+	fr, err := s.cl.DoGet(context.TODO(), ep[0].GetTicket())
+	s.Require().NoError(err)
+	var recs []arrow.RecordBatch
+	for fr.Next() {
+		rec := fr.RecordBatch()
+		rec.Retain()
+		defer rec.Release()
+		recs = append(recs, rec)
+	}
+	s.Require().NoError(fr.Err())
+	tbl := array.NewTableFromRecords(fr.Schema(), recs)
+	defer tbl.Release()
+	s.Assert().Equal(int64(2), tbl.NumRows())
+	s.Assert().Equal(int64(1), tbl.NumCols())
+	col := tbl.Column(0)
+	s.Assert().Equal("t1", col.Name())
+	s.Assert().Equal(2, col.Len())
+	s.Assert().Equal(1, col.NullN())
+	s.Assert().Equal(arrow.INT16, col.DataType().ID())
+	var n int
+	for _, arr := range col.Data().Chunks() {
+		data := array.NewInt16Data(arr.Data())
+		defer data.Release()
+		for i := 0; i < data.Len(); i++ {
+			switch n {
+			case 0:
+				s.True(data.IsNull(i))
+			case 1:
+				s.False(data.IsNull(i))
+				s.Assert().Equal(int16(1), data.Value(i))
+			}
+			n++
+		}
+	}
+}
+
+func (s *FlightSqlServerSuite) TestExecutePreparedStatementUpdate() {
+	prep, err := s.cl.Prepare(context.TODO(), "prepared update")
+	s.Require().NoError(err)
+	defer prep.Close(context.TODO())
+
+	val, ok := prep.IsUpdate()
+	s.Require().True(ok)
+	s.True(val)
+
+	nrecords, err := prep.ExecuteUpdate(context.TODO())
+	s.Require().NoError(err)
+	s.Assert().Equal(int64(2), nrecords)
+}
+
 func (s *FlightSqlServerSuite) TestExecutePoll() {
 	poll, err := s.cl.ExecutePoll(context.TODO(), "1", nil)
 	s.NoError(err)