| 1860 | } |
| 1861 | |
| 1862 | func (s *MockServer) DoPutPreparedStatementQuery(ctx context.Context, qry flightsql.PreparedStatementQuery, r flight.MessageReader, w flight.MetadataWriter) ([]byte, error) { |
| 1863 | if s.ExpectedPreparedStatementSchema != nil { |
| 1864 | if !s.ExpectedPreparedStatementSchema.Equal(r.Schema()) { |
| 1865 | return nil, errors.New("parameter schema: unexpected") |
| 1866 | } |
| 1867 | return qry.GetPreparedStatementHandle(), nil |
| 1868 | } |
| 1869 | |
| 1870 | if s.PreparedStatementParameterSchema != nil && !s.PreparedStatementParameterSchema.Equal(r.Schema()) { |
| 1871 | return nil, fmt.Errorf("parameter schema: %w", arrow.ErrInvalid) |
| 1872 | } |
| 1873 | |
| 1874 | // GH-35328: it's rare, but this function can complete execution and return |
| 1875 | // closing the reader *after* the schema is written but *before* the parameter batch |
| 1876 | // is written (race condition based on goroutine scheduling). In that situation, |
| 1877 | // the client call to Write the parameter record batch will return an io.EOF because |
| 1878 | // this end of the connection will have closed before it attempted to send the batch. |
| 1879 | // This created a flaky test situation that was difficult to reproduce (1-4 failures |
| 1880 | // in 5000 runs). We can avoid this flakiness by simply *explicitly* draining the |
| 1881 | // record batch messages from the reader before returning. |
| 1882 | for r.Next() { |
| 1883 | } |
| 1884 | |
| 1885 | return qry.GetPreparedStatementHandle(), nil |
| 1886 | } |
| 1887 | |
| 1888 | func (s *MockServer) DoGetStatement(ctx context.Context, ticket flightsql.StatementQueryTicket) (*arrow.Schema, <-chan flight.StreamChunk, error) { |
| 1889 | record, _, err := array.RecordFromJSON(memory.DefaultAllocator, s.DataSchema, strings.NewReader(s.Data)) |