lidavidm commented on code in PR #980:
URL: https://github.com/apache/arrow-adbc/pull/980#discussion_r1301630844
##########
go/adbc/driver/flightsql/flightsql_adbc_server_test.go:
##########
@@ -627,3 +632,89 @@ func (suite *DataTypeTests) TestListInt() {
func (suite *DataTypeTests) TestMapIntInt() {
suite.DoTestCase("map[int]int", SchemaMapIntInt)
}
+
+// ---- Multi Table Tests --------------------
+
+type MultiTableTestServer struct {
+ flightsql.BaseServer
+}
+
+func (server *MultiTableTestServer) GetFlightInfoStatement(ctx
context.Context, cmd flightsql.StatementQuery, desc *flight.FlightDescriptor)
(*flight.FlightInfo, error) {
+ query := cmd.GetQuery()
+ tkt, err := flightsql.CreateStatementQueryTicket([]byte(query))
+ if err != nil {
+ return nil, err
+ }
+
+ return &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{{Ticket:
&flight.Ticket{Ticket: tkt}}},
+ FlightDescriptor: desc,
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }, nil
+}
+
+func (server *MultiTableTestServer) GetFlightInfoTables(ctx context.Context,
cmd flightsql.GetTables, desc *flight.FlightDescriptor) (*flight.FlightInfo,
error) {
+ schema := schema_ref.Tables
+ if cmd.GetIncludeSchema() {
+ schema = schema_ref.TablesWithIncludedSchema
+ }
+ server.Alloc = memory.NewCheckedAllocator(memory.DefaultAllocator)
+ info := &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{
+ {Ticket: &flight.Ticket{Ticket: desc.Cmd}},
+ },
+ FlightDescriptor: desc,
+ Schema: flight.SerializeSchema(schema, server.Alloc),
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }
+
+ return info, nil
+}
+
+func (server *MultiTableTestServer) DoGetTables(ctx context.Context, cmd
flightsql.GetTables) (*arrow.Schema, <-chan flight.StreamChunk, error) {
+ bldr := array.NewRecordBuilder(server.Alloc, adbc.GetTableSchemaSchema)
+
+ bldr.Field(0).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(1).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(2).(*array.StringBuilder).AppendValues([]string{"tbl1",
"tbl2"}, nil)
+ bldr.Field(3).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+
+ sc1 := arrow.NewSchema([]arrow.Field{{Name: "a", Type:
arrow.PrimitiveTypes.Int32, Nullable: true}}, nil)
+ sc2 := arrow.NewSchema([]arrow.Field{{Name: "b", Type:
arrow.PrimitiveTypes.Int32, Nullable: true}}, nil)
+ buf1 := flight.SerializeSchema(sc1, server.Alloc)
+ buf2 := flight.SerializeSchema(sc2, server.Alloc)
+
+ bldr.Field(4).(*array.BinaryBuilder).AppendValues([][]byte{buf1, buf2},
nil)
+ defer bldr.Release()
+
+ rec := bldr.NewRecord()
+
+ ch := make(chan flight.StreamChunk)
+ go func() {
+ defer close(ch)
+ ch <- flight.StreamChunk{
+ Data: rec,
+ Desc: nil,
+ Err: nil,
+ }
+ }()
+ return adbc.GetTableSchemaSchema, ch, nil
+}
+
+type MultiTableTests struct {
+ ServerBasedTests
+}
+
+func (suite *MultiTableTests) SetupSuite() {
+ suite.DoSetupSuite(&MultiTableTestServer{}, nil, map[string]string{})
+}
+
+func (suite *MultiTableTests) TestGetTableSchema() {
Review Comment:
```suggestion
// Regression test for https://github.com/apache/arrow-adbc/issues/934
func (suite *MultiTableTests) TestGetTableSchema() {
```
##########
go/adbc/driver/flightsql/flightsql_adbc_server_test.go:
##########
@@ -627,3 +632,89 @@ func (suite *DataTypeTests) TestListInt() {
func (suite *DataTypeTests) TestMapIntInt() {
suite.DoTestCase("map[int]int", SchemaMapIntInt)
}
+
+// ---- Multi Table Tests --------------------
+
+type MultiTableTestServer struct {
+ flightsql.BaseServer
+}
+
+func (server *MultiTableTestServer) GetFlightInfoStatement(ctx
context.Context, cmd flightsql.StatementQuery, desc *flight.FlightDescriptor)
(*flight.FlightInfo, error) {
+ query := cmd.GetQuery()
+ tkt, err := flightsql.CreateStatementQueryTicket([]byte(query))
+ if err != nil {
+ return nil, err
+ }
+
+ return &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{{Ticket:
&flight.Ticket{Ticket: tkt}}},
+ FlightDescriptor: desc,
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }, nil
+}
+
+func (server *MultiTableTestServer) GetFlightInfoTables(ctx context.Context,
cmd flightsql.GetTables, desc *flight.FlightDescriptor) (*flight.FlightInfo,
error) {
+ schema := schema_ref.Tables
+ if cmd.GetIncludeSchema() {
+ schema = schema_ref.TablesWithIncludedSchema
+ }
+ server.Alloc = memory.NewCheckedAllocator(memory.DefaultAllocator)
+ info := &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{
+ {Ticket: &flight.Ticket{Ticket: desc.Cmd}},
+ },
+ FlightDescriptor: desc,
+ Schema: flight.SerializeSchema(schema, server.Alloc),
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }
+
+ return info, nil
+}
+
+func (server *MultiTableTestServer) DoGetTables(ctx context.Context, cmd
flightsql.GetTables) (*arrow.Schema, <-chan flight.StreamChunk, error) {
+ bldr := array.NewRecordBuilder(server.Alloc, adbc.GetTableSchemaSchema)
+
+ bldr.Field(0).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(1).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(2).(*array.StringBuilder).AppendValues([]string{"tbl1",
"tbl2"}, nil)
Review Comment:
I don't think this tests the original bug, though, where the issue is that
if you requested a table ("foo") and the server returned multiple tables where
the requested table was a prefix of one of the requested tables ("foobar"),
then the client would pick that table.
##########
go/adbc/driver/flightsql/flightsql_adbc_server_test.go:
##########
@@ -627,3 +632,89 @@ func (suite *DataTypeTests) TestListInt() {
func (suite *DataTypeTests) TestMapIntInt() {
suite.DoTestCase("map[int]int", SchemaMapIntInt)
}
+
+// ---- Multi Table Tests --------------------
+
+type MultiTableTestServer struct {
+ flightsql.BaseServer
+}
+
+func (server *MultiTableTestServer) GetFlightInfoStatement(ctx
context.Context, cmd flightsql.StatementQuery, desc *flight.FlightDescriptor)
(*flight.FlightInfo, error) {
+ query := cmd.GetQuery()
+ tkt, err := flightsql.CreateStatementQueryTicket([]byte(query))
+ if err != nil {
+ return nil, err
+ }
+
+ return &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{{Ticket:
&flight.Ticket{Ticket: tkt}}},
+ FlightDescriptor: desc,
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }, nil
+}
+
+func (server *MultiTableTestServer) GetFlightInfoTables(ctx context.Context,
cmd flightsql.GetTables, desc *flight.FlightDescriptor) (*flight.FlightInfo,
error) {
+ schema := schema_ref.Tables
+ if cmd.GetIncludeSchema() {
+ schema = schema_ref.TablesWithIncludedSchema
+ }
+ server.Alloc = memory.NewCheckedAllocator(memory.DefaultAllocator)
+ info := &flight.FlightInfo{
+ Endpoint: []*flight.FlightEndpoint{
+ {Ticket: &flight.Ticket{Ticket: desc.Cmd}},
+ },
+ FlightDescriptor: desc,
+ Schema: flight.SerializeSchema(schema, server.Alloc),
+ TotalRecords: -1,
+ TotalBytes: -1,
+ }
+
+ return info, nil
+}
+
+func (server *MultiTableTestServer) DoGetTables(ctx context.Context, cmd
flightsql.GetTables) (*arrow.Schema, <-chan flight.StreamChunk, error) {
+ bldr := array.NewRecordBuilder(server.Alloc, adbc.GetTableSchemaSchema)
+
+ bldr.Field(0).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(1).(*array.StringBuilder).AppendValues([]string{"", ""}, nil)
+ bldr.Field(2).(*array.StringBuilder).AppendValues([]string{"tbl1",
"tbl2"}, nil)
Review Comment:
Does this test fail without the fix above?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
