[GitHub] [arrow] rafael-telles opened a new pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles opened a new pull request #11989:
URL: https://github.com/apache/arrow/pull/11989


   This adds Flight SQL scenarios for integration tests for C++ and Java implementations.
   The integrations tests assert that:
   - RPC objects built on clients are parsed correctly on servers
   - Arrow vectors built on servers have the expected Arrow schemas when received on clients
   
   While working on this I found some inconsistencies between both implementation, this PR also fixes these problems.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r771725031



##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.h
##########
@@ -28,22 +28,27 @@
 
 namespace arrow {

Review comment:
       This isn't your mistake, but this file is missing a `#pragma once` up top - hence the CI failure.

##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java
##########
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import static com.google.protobuf.Any.pack;
+import static java.util.Collections.singletonList;
+
+import java.util.List;
+
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+
+/**
+ * Hardcoded Flight SQL producer used for cross-language integration tests.
+ */
+public class FlightSqlScenarioProducer implements FlightSqlProducer {
+  private final BufferAllocator allocator;
+
+  public FlightSqlScenarioProducer(BufferAllocator allocator) {
+    this.allocator = allocator;
+  }
+
+  static Schema getQuerySchema() {
+    return new Schema(
+        singletonList(
+            new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null)
+        )
+    );
+  }
+
+  @Override
+  public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementRequest request,
+                                      CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getQuery().equals("SELECT PREPARED STATEMENT") ||
+            request.getQuery().equals("UPDATE PREPARED STATEMENT"));
+
+    final FlightSql.ActionCreatePreparedStatementResult
+        result = FlightSql.ActionCreatePreparedStatementResult.newBuilder()
+        .setPreparedStatementHandle(ByteString.copyFromUtf8(request.getQuery() + " HANDLE"))
+        .build();
+    listener.onNext(new Result(pack(result).toByteArray()));
+    listener.onCompleted();
+  }
+
+  @Override
+  public void closePreparedStatement(FlightSql.ActionClosePreparedStatementRequest request,
+                                     CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getPreparedStatementHandle().toStringUtf8().equals("SELECT PREPARED STATEMENT HANDLE") ||
+            request.getPreparedStatementHandle().toStringUtf8().equals("UPDATE PREPARED STATEMENT HANDLE"));
+
+    listener.onCompleted();
+  }
+
+  @Override
+  public FlightInfo getFlightInfoStatement(FlightSql.CommandStatementQuery command,
+                                           CallContext context, FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT");
+
+    ByteString handle = ByteString.copyFromUtf8("SELECT STATEMENT HANDLE");
+
+    FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder()
+        .setStatementHandle(handle)
+        .build();
+    return getFlightInfoForSchema(ticket, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public FlightInfo getFlightInfoPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                                   CallContext context,
+                                                   FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return getFlightInfoForSchema(command, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public SchemaResult getSchemaStatement(FlightSql.CommandStatementQuery command,
+                                         CallContext context, FlightDescriptor descriptor) {
+    return new SchemaResult(getQuerySchema());
+  }
+
+  @Override
+  public void getStreamStatement(FlightSql.TicketStatementQuery ticket, CallContext context,
+                                 ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  @Override
+  public void getStreamPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                         CallContext context, ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  private Runnable acceptPutReturnConstant(StreamListener<PutResult> ackStream, int value) {
+    return () -> {
+      final FlightSql.DoPutUpdateResult build =
+          FlightSql.DoPutUpdateResult.newBuilder().setRecordCount(value).build();
+
+      try (final ArrowBuf buffer = allocator.buffer(build.getSerializedSize())) {
+        buffer.writeBytes(build.toByteArray());
+        ackStream.onNext(PutResult.metadata(buffer));
+        ackStream.onCompleted();
+      }
+    };
+  }
+
+  @Override
+  public Runnable acceptPutStatement(FlightSql.CommandStatementUpdate command, CallContext context,
+                                     FlightStream flightStream,
+                                     StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "UPDATE STATEMENT");
+
+    return acceptPutReturnConstant(ackStream, 10000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementUpdate(FlightSql.CommandPreparedStatementUpdate command,
+                                                   CallContext context, FlightStream flightStream,
+                                                   StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "UPDATE PREPARED STATEMENT HANDLE");
+
+    return acceptPutReturnConstant(ackStream, 20000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementQuery(FlightSql.CommandPreparedStatementQuery command,
+                                                  CallContext context, FlightStream flightStream,
+                                                  StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return null;
+  }
+
+  @Override
+  public FlightInfo getFlightInfoSqlInfo(FlightSql.CommandGetSqlInfo request, CallContext context,
+                                         FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(request.getInfoCount(), 2);
+    IntegrationAssertions.assertEquals(request.getInfo(0),
+        FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE);
+    IntegrationAssertions.assertEquals(request.getInfo(1),
+        FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE);
+
+    return getFlightInfoForSchema(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA);
+  }
+
+  @Override
+  public void getStreamSqlInfo(FlightSql.CommandGetSqlInfo command, CallContext context,
+                               ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, Schemas.GET_SQL_INFO_SCHEMA);
+  }
+
+  @Override
+  public FlightInfo getFlightInfoCatalogs(FlightSql.CommandGetCatalogs request, CallContext context,
+                                          FlightDescriptor descriptor) {
+    return getFlightInfoForSchema(request, descriptor, Schemas.GET_CATALOGS_SCHEMA);
+  }
+
+  private void serveJsonToStreamListener(ServerStreamListener stream, Schema schema) {

Review comment:
       What does this have to do with JSON?

##########
File path: cpp/src/arrow/CMakeLists.txt
##########
@@ -736,6 +736,10 @@ if(ARROW_FLIGHT_SQL)
   add_subdirectory(flight/sql)
 endif()
 
+if(ARROW_FLIGHT AND ARROW_BUILD_INTEGRATION)

Review comment:
       `AND ARROW_FLIGHT_SQL`?

##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java
##########
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import static com.google.protobuf.Any.pack;
+import static java.util.Collections.singletonList;
+
+import java.util.List;
+
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+
+/**
+ * Hardcoded Flight SQL producer used for cross-language integration tests.
+ */
+public class FlightSqlScenarioProducer implements FlightSqlProducer {
+  private final BufferAllocator allocator;
+
+  public FlightSqlScenarioProducer(BufferAllocator allocator) {
+    this.allocator = allocator;
+  }
+
+  static Schema getQuerySchema() {
+    return new Schema(
+        singletonList(
+            new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null)
+        )
+    );
+  }
+
+  @Override
+  public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementRequest request,
+                                      CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getQuery().equals("SELECT PREPARED STATEMENT") ||
+            request.getQuery().equals("UPDATE PREPARED STATEMENT"));
+
+    final FlightSql.ActionCreatePreparedStatementResult
+        result = FlightSql.ActionCreatePreparedStatementResult.newBuilder()
+        .setPreparedStatementHandle(ByteString.copyFromUtf8(request.getQuery() + " HANDLE"))
+        .build();
+    listener.onNext(new Result(pack(result).toByteArray()));
+    listener.onCompleted();
+  }
+
+  @Override
+  public void closePreparedStatement(FlightSql.ActionClosePreparedStatementRequest request,
+                                     CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getPreparedStatementHandle().toStringUtf8().equals("SELECT PREPARED STATEMENT HANDLE") ||
+            request.getPreparedStatementHandle().toStringUtf8().equals("UPDATE PREPARED STATEMENT HANDLE"));
+
+    listener.onCompleted();
+  }
+
+  @Override
+  public FlightInfo getFlightInfoStatement(FlightSql.CommandStatementQuery command,
+                                           CallContext context, FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT");
+
+    ByteString handle = ByteString.copyFromUtf8("SELECT STATEMENT HANDLE");
+
+    FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder()
+        .setStatementHandle(handle)
+        .build();
+    return getFlightInfoForSchema(ticket, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public FlightInfo getFlightInfoPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                                   CallContext context,
+                                                   FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return getFlightInfoForSchema(command, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public SchemaResult getSchemaStatement(FlightSql.CommandStatementQuery command,
+                                         CallContext context, FlightDescriptor descriptor) {
+    return new SchemaResult(getQuerySchema());
+  }
+
+  @Override
+  public void getStreamStatement(FlightSql.TicketStatementQuery ticket, CallContext context,
+                                 ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  @Override
+  public void getStreamPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                         CallContext context, ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  private Runnable acceptPutReturnConstant(StreamListener<PutResult> ackStream, int value) {
+    return () -> {
+      final FlightSql.DoPutUpdateResult build =
+          FlightSql.DoPutUpdateResult.newBuilder().setRecordCount(value).build();
+
+      try (final ArrowBuf buffer = allocator.buffer(build.getSerializedSize())) {
+        buffer.writeBytes(build.toByteArray());
+        ackStream.onNext(PutResult.metadata(buffer));
+        ackStream.onCompleted();
+      }
+    };
+  }
+
+  @Override
+  public Runnable acceptPutStatement(FlightSql.CommandStatementUpdate command, CallContext context,
+                                     FlightStream flightStream,
+                                     StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "UPDATE STATEMENT");
+
+    return acceptPutReturnConstant(ackStream, 10000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementUpdate(FlightSql.CommandPreparedStatementUpdate command,
+                                                   CallContext context, FlightStream flightStream,
+                                                   StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "UPDATE PREPARED STATEMENT HANDLE");
+
+    return acceptPutReturnConstant(ackStream, 20000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementQuery(FlightSql.CommandPreparedStatementQuery command,
+                                                  CallContext context, FlightStream flightStream,
+                                                  StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return null;
+  }
+
+  @Override
+  public FlightInfo getFlightInfoSqlInfo(FlightSql.CommandGetSqlInfo request, CallContext context,
+                                         FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(request.getInfoCount(), 2);
+    IntegrationAssertions.assertEquals(request.getInfo(0),
+        FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME_VALUE);
+    IntegrationAssertions.assertEquals(request.getInfo(1),
+        FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY_VALUE);
+
+    return getFlightInfoForSchema(request, descriptor, Schemas.GET_SQL_INFO_SCHEMA);
+  }
+
+  @Override
+  public void getStreamSqlInfo(FlightSql.CommandGetSqlInfo command, CallContext context,
+                               ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, Schemas.GET_SQL_INFO_SCHEMA);
+  }
+
+  @Override
+  public FlightInfo getFlightInfoCatalogs(FlightSql.CommandGetCatalogs request, CallContext context,
+                                          FlightDescriptor descriptor) {
+    return getFlightInfoForSchema(request, descriptor, Schemas.GET_CATALOGS_SCHEMA);
+  }
+
+  private void serveJsonToStreamListener(ServerStreamListener stream, Schema schema) {

Review comment:
       Maybe `putEmptyBatchToStreamListener`?

##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java
##########
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import static com.google.protobuf.Any.pack;
+import static java.util.Collections.singletonList;
+
+import java.util.List;
+
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+
+/**
+ * Hardcoded Flight SQL producer used for cross-language integration tests.
+ */
+public class FlightSqlScenarioProducer implements FlightSqlProducer {
+  private final BufferAllocator allocator;
+
+  public FlightSqlScenarioProducer(BufferAllocator allocator) {
+    this.allocator = allocator;
+  }
+
+  static Schema getQuerySchema() {
+    return new Schema(
+        singletonList(
+            new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null)
+        )
+    );
+  }
+
+  @Override
+  public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementRequest request,
+                                      CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getQuery().equals("SELECT PREPARED STATEMENT") ||
+            request.getQuery().equals("UPDATE PREPARED STATEMENT"));
+
+    final FlightSql.ActionCreatePreparedStatementResult
+        result = FlightSql.ActionCreatePreparedStatementResult.newBuilder()
+        .setPreparedStatementHandle(ByteString.copyFromUtf8(request.getQuery() + " HANDLE"))
+        .build();
+    listener.onNext(new Result(pack(result).toByteArray()));
+    listener.onCompleted();
+  }
+
+  @Override
+  public void closePreparedStatement(FlightSql.ActionClosePreparedStatementRequest request,
+                                     CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getPreparedStatementHandle().toStringUtf8().equals("SELECT PREPARED STATEMENT HANDLE") ||
+            request.getPreparedStatementHandle().toStringUtf8().equals("UPDATE PREPARED STATEMENT HANDLE"));
+
+    listener.onCompleted();
+  }
+
+  @Override
+  public FlightInfo getFlightInfoStatement(FlightSql.CommandStatementQuery command,
+                                           CallContext context, FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT");
+
+    ByteString handle = ByteString.copyFromUtf8("SELECT STATEMENT HANDLE");
+
+    FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder()
+        .setStatementHandle(handle)
+        .build();
+    return getFlightInfoForSchema(ticket, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public FlightInfo getFlightInfoPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                                   CallContext context,
+                                                   FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return getFlightInfoForSchema(command, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public SchemaResult getSchemaStatement(FlightSql.CommandStatementQuery command,
+                                         CallContext context, FlightDescriptor descriptor) {
+    return new SchemaResult(getQuerySchema());
+  }
+
+  @Override
+  public void getStreamStatement(FlightSql.TicketStatementQuery ticket, CallContext context,
+                                 ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  @Override
+  public void getStreamPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                         CallContext context, ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  private Runnable acceptPutReturnConstant(StreamListener<PutResult> ackStream, int value) {
+    return () -> {
+      final FlightSql.DoPutUpdateResult build =
+          FlightSql.DoPutUpdateResult.newBuilder().setRecordCount(value).build();
+
+      try (final ArrowBuf buffer = allocator.buffer(build.getSerializedSize())) {
+        buffer.writeBytes(build.toByteArray());
+        ackStream.onNext(PutResult.metadata(buffer));
+        ackStream.onCompleted();
+      }
+    };
+  }
+
+  @Override
+  public Runnable acceptPutStatement(FlightSql.CommandStatementUpdate command, CallContext context,
+                                     FlightStream flightStream,
+                                     StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "UPDATE STATEMENT");
+
+    return acceptPutReturnConstant(ackStream, 10000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementUpdate(FlightSql.CommandPreparedStatementUpdate command,
+                                                   CallContext context, FlightStream flightStream,
+                                                   StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "UPDATE PREPARED STATEMENT HANDLE");
+
+    return acceptPutReturnConstant(ackStream, 20000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementQuery(FlightSql.CommandPreparedStatementQuery command,
+                                                  CallContext context, FlightStream flightStream,
+                                                  StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return null;

Review comment:
       Hmm, this should throw an exception if it's not implemented.

##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java
##########
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import static com.google.protobuf.Any.pack;
+import static java.util.Collections.singletonList;
+
+import java.util.List;
+
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+
+/**
+ * Hardcoded Flight SQL producer used for cross-language integration tests.
+ */
+public class FlightSqlScenarioProducer implements FlightSqlProducer {
+  private final BufferAllocator allocator;
+
+  public FlightSqlScenarioProducer(BufferAllocator allocator) {
+    this.allocator = allocator;
+  }
+
+  static Schema getQuerySchema() {

Review comment:
       Can we explicitly document this as the schema returned for a mock select query (both here and in C++)?

##########
File path: cpp/src/arrow/flight/integration_tests/CMakeLists.txt
##########
@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+add_custom_target(arrow_flight_integration_tests)

Review comment:
       Is the target necessary?

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,665 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL server which only asserts
+/// that commands called during integration tests are being parsed correctly and returns the expected schemas to be
+/// validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return 10000;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    return Status::NotImplemented("Not implemented");
+  }
+
+  arrow::Result<int64_t> DoPutPreparedStatementUpdate(
+      const ServerCallContext& context, const sql::PreparedStatementUpdate& command,
+      FlightMessageReader* reader) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return 20000;
+  }
+
+ private:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoForCommand(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema) {
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{descriptor.cmd}, {}}};
+    ARROW_ASSIGN_OR_RAISE(auto result,
+                          FlightInfo::Make(*schema, descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetForTestCase(
+      std::shared_ptr<Schema> schema) {
+    ARROW_ASSIGN_OR_RAISE(auto reader2, RecordBatchReader::Make({}, schema));

Review comment:
       Also, why is it `reader2`?

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,665 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL server which only asserts
+/// that commands called during integration tests are being parsed correctly and returns the expected schemas to be
+/// validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return 10000;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    return Status::NotImplemented("Not implemented");
+  }
+
+  arrow::Result<int64_t> DoPutPreparedStatementUpdate(
+      const ServerCallContext& context, const sql::PreparedStatementUpdate& command,
+      FlightMessageReader* reader) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return 20000;
+  }
+
+ private:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoForCommand(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema) {
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{descriptor.cmd}, {}}};
+    ARROW_ASSIGN_OR_RAISE(auto result,
+                          FlightInfo::Make(*schema, descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetForTestCase(
+      std::shared_ptr<Schema> schema) {
+    ARROW_ASSIGN_OR_RAISE(auto reader2, RecordBatchReader::Make({}, schema));

Review comment:
       nit: `std::move(schema)`? (or declare the parameter as `const&`)

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,665 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL server which only asserts
+/// that commands called during integration tests are being parsed correctly and returns the expected schemas to be
+/// validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return 10000;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    return Status::NotImplemented("Not implemented");

Review comment:
       Shall we test this flow as well?

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,665 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL server which only asserts
+/// that commands called during integration tests are being parsed correctly and returns the expected schemas to be
+/// validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return 10000;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    return Status::NotImplemented("Not implemented");
+  }
+
+  arrow::Result<int64_t> DoPutPreparedStatementUpdate(
+      const ServerCallContext& context, const sql::PreparedStatementUpdate& command,
+      FlightMessageReader* reader) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return 20000;
+  }
+
+ private:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoForCommand(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema) {
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{descriptor.cmd}, {}}};
+    ARROW_ASSIGN_OR_RAISE(auto result,
+                          FlightInfo::Make(*schema, descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetForTestCase(
+      std::shared_ptr<Schema> schema) {
+    ARROW_ASSIGN_OR_RAISE(auto reader2, RecordBatchReader::Make({}, schema));
+    return std::unique_ptr<FlightDataStream>(new RecordBatchStream(reader2));
+  }
+};
+
+/// \brief Integration test scenario for validating Flight SQL specs across multiple
+/// implementations. This should ensure that RPC objects are being built and parsed
+/// correctly for multiple languages and that the Arrow schemas are returned as expected.
+class FlightSqlScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new FlightSqlScenarioServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status Validate(std::shared_ptr<Schema> expectedSchema,
+                  arrow::Result<std::unique_ptr<FlightInfo>> flightInfo,
+                  sql::FlightSqlClient* sql_client) {
+    FlightCallOptions call_options;
+
+    ARROW_ASSIGN_OR_RAISE(auto flight_info, flightInfo);
+    ARROW_ASSIGN_OR_RAISE(
+        auto reader, sql_client->DoGet(call_options, flight_info->endpoints()[0].ticket));
+
+    ARROW_ASSIGN_OR_RAISE(auto actual_schema, reader->GetSchema());
+
+    AssertSchemaEqual(expectedSchema, actual_schema);
+
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    sql::FlightSqlClient sql_client(std::move(client));
+
+    ARROW_RETURN_NOT_OK(ValidateMetadataRetrieval(&sql_client));
+
+    ARROW_RETURN_NOT_OK(ValidateStatementExecution(&sql_client));
+
+    ARROW_RETURN_NOT_OK(ValidatePreparedStatementExecution(&sql_client));
+
+    return Status::OK();
+  }
+
+  Status ValidateMetadataRetrieval(sql::FlightSqlClient* sql_client) {
+    FlightCallOptions options;
+
+    std::string catalog = "catalog";
+    std::string db_schema_filter_pattern = "db_schema_filter_pattern";
+    std::string table_filter_pattern = "table_filter_pattern";
+    std::string table = "table";
+    std::string db_schema = "db_schema";
+    std::vector<std::string> table_types = {"table", "view"};
+
+    sql::TableRef table_ref = {catalog, db_schema, table};
+    sql::TableRef pk_table_ref = {"pk_catalog", "pk_db_schema", "pk_table"};
+    sql::TableRef fk_table_ref = {"fk_catalog", "fk_db_schema", "fk_table"};
+
+    ARROW_RETURN_NOT_OK(Validate(sql::SqlSchema::GetCatalogsSchema(),
+                                 sql_client->GetCatalogs(options), sql_client));
+    ARROW_RETURN_NOT_OK(
+        Validate(sql::SqlSchema::GetDbSchemasSchema(),
+                 sql_client->GetDbSchemas(options, &catalog, &db_schema_filter_pattern),
+                 sql_client));
+    ARROW_RETURN_NOT_OK(
+        Validate(sql::SqlSchema::GetTablesSchemaWithIncludedSchema(),
+                 sql_client->GetTables(options, &catalog, &db_schema_filter_pattern,
+                                       &table_filter_pattern, true, &table_types),
+                 sql_client));
+    ARROW_RETURN_NOT_OK(Validate(sql::SqlSchema::GetTableTypesSchema(),
+                                 sql_client->GetTableTypes(options), sql_client));
+    ARROW_RETURN_NOT_OK(Validate(sql::SqlSchema::GetPrimaryKeysSchema(),
+                                 sql_client->GetPrimaryKeys(options, table_ref),
+                                 sql_client));
+    ARROW_RETURN_NOT_OK(Validate(sql::SqlSchema::GetExportedKeysSchema(),
+                                 sql_client->GetExportedKeys(options, table_ref),
+                                 sql_client));
+    ARROW_RETURN_NOT_OK(Validate(sql::SqlSchema::GetImportedKeysSchema(),
+                                 sql_client->GetImportedKeys(options, table_ref),
+                                 sql_client));
+    ARROW_RETURN_NOT_OK(Validate(
+        sql::SqlSchema::GetCrossReferenceSchema(),
+        sql_client->GetCrossReference(options, pk_table_ref, fk_table_ref), sql_client));
+    ARROW_RETURN_NOT_OK(Validate(
+        sql::SqlSchema::GetSqlInfoSchema(),
+        sql_client->GetSqlInfo(
+            options, {sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME,
+                      sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY}),
+        sql_client));
+
+    return Status::OK();
+  }
+
+  Status ValidateStatementExecution(sql::FlightSqlClient* sql_client) {
+    FlightCallOptions options;
+
+    ARROW_RETURN_NOT_OK(Validate(
+            GetQuerySchema(), sql_client->Execute(options, "SELECT STATEMENT"), sql_client));
+    ARROW_ASSIGN_OR_RAISE(auto update_statement_result,
+                          sql_client->ExecuteUpdate(options, "UPDATE STATEMENT"));
+    if (update_statement_result != 10000L) {
+      return Status::Invalid("Expected 'UPDATE STATEMENT' return 10000, got ",
+                             update_statement_result);
+    }
+
+    return Status::OK();
+  }
+
+  Status ValidatePreparedStatementExecution(sql::FlightSqlClient* sql_client) {
+    FlightCallOptions options;
+
+    ARROW_ASSIGN_OR_RAISE(auto select_prepared_statement,
+                          sql_client->Prepare(options, "SELECT PREPARED STATEMENT"));
+    ARROW_RETURN_NOT_OK(
+            Validate(GetQuerySchema(), select_prepared_statement->Execute(), sql_client));
+
+    ARROW_ASSIGN_OR_RAISE(auto update_prepared_statement,
+                          sql_client->Prepare(options, "UPDATE PREPARED STATEMENT"));
+    ARROW_ASSIGN_OR_RAISE(auto update_prepared_statement_result,
+                          update_prepared_statement->ExecuteUpdate());
+    if (update_prepared_statement_result != 20000L) {
+      return Status::Invalid("Expected 'UPDATE STATEMENT' return 20000, got ",
+                             update_prepared_statement_result);
+    }
+
+    return Status::OK();

Review comment:
       Can we explicitly close the prepared statement here?




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-997051330


   Hi @lidavidm ! We wrote the integration tests for Flight SQL as requested :)


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r773379954



##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,685 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+/// \brief Schema to be returned for mocking the statement/prepared statement results.
+/// Must be the same across all languages.
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+constexpr int64_t kUpdateStatementExpectedRows = 10000L;
+constexpr int64_t kUpdatePreparedStatementExpectedRows = 20000L;
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL
+/// server which only asserts that commands called during integration tests are being
+/// parsed correctly and returns the expected schemas to be validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return kUpdateStatementExpectedRows;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    ARROW_ASSIGN_OR_RAISE(auto actual_schema, reader->GetSchema());
+    ARROW_RETURN_NOT_OK(AssertEq<Schema>(*GetQuerySchema(), *actual_schema));
+
+    return Status::OK();
+  }
+
+  arrow::Result<int64_t> DoPutPreparedStatementUpdate(
+      const ServerCallContext& context, const sql::PreparedStatementUpdate& command,
+      FlightMessageReader* reader) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return kUpdatePreparedStatementExpectedRows;
+  }
+
+ private:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoForCommand(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema) {
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{descriptor.cmd}, {}}};
+    ARROW_ASSIGN_OR_RAISE(auto result,
+                          FlightInfo::Make(*schema, descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetForTestCase(
+      const std::shared_ptr<Schema>& schema) {
+    ARROW_ASSIGN_OR_RAISE(auto reader, RecordBatchReader::Make({}, schema));
+    return std::unique_ptr<FlightDataStream>(new RecordBatchStream(reader));
+  }
+};
+
+/// \brief Integration test scenario for validating Flight SQL specs across multiple
+/// implementations. This should ensure that RPC objects are being built and parsed
+/// correctly for multiple languages and that the Arrow schemas are returned as expected.
+class FlightSqlScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new FlightSqlScenarioServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status Validate(std::shared_ptr<Schema> expectedSchema,
+                  arrow::Result<std::unique_ptr<FlightInfo>> flightInfo,
+                  sql::FlightSqlClient* sql_client) {
+    FlightCallOptions call_options;
+
+    ARROW_ASSIGN_OR_RAISE(auto flight_info, flightInfo);
+    ARROW_ASSIGN_OR_RAISE(
+        auto reader, sql_client->DoGet(call_options, flight_info->endpoints()[0].ticket));
+
+    ARROW_ASSIGN_OR_RAISE(auto actual_schema, reader->GetSchema());
+
+    AssertSchemaEqual(expectedSchema, actual_schema);

Review comment:
       Fixed




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r772542505



##########
File path: cpp/src/arrow/flight/integration_tests/CMakeLists.txt
##########
@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+add_custom_target(arrow_flight_integration_tests)

Review comment:
       It was already there on arrow/flight/CMakeLists.txt before moving to the new directory




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998282667


   Hey @lidavidm , I don't know if these CI failures has something to do with these changes... Please let me know if there is anything else to do here


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r771763907



##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenarioProducer.java
##########
@@ -0,0 +1,343 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import static com.google.protobuf.Any.pack;
+import static java.util.Collections.singletonList;
+
+import java.util.List;
+
+import org.apache.arrow.flight.Criteria;
+import org.apache.arrow.flight.FlightDescriptor;
+import org.apache.arrow.flight.FlightEndpoint;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.PutResult;
+import org.apache.arrow.flight.Result;
+import org.apache.arrow.flight.SchemaResult;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.memory.ArrowBuf;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.ArrowType;
+import org.apache.arrow.vector.types.pojo.Field;
+import org.apache.arrow.vector.types.pojo.FieldType;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Message;
+
+/**
+ * Hardcoded Flight SQL producer used for cross-language integration tests.
+ */
+public class FlightSqlScenarioProducer implements FlightSqlProducer {
+  private final BufferAllocator allocator;
+
+  public FlightSqlScenarioProducer(BufferAllocator allocator) {
+    this.allocator = allocator;
+  }
+
+  static Schema getQuerySchema() {
+    return new Schema(
+        singletonList(
+            new Field("id", FieldType.nullable(new ArrowType.Int(64, true)), null)
+        )
+    );
+  }
+
+  @Override
+  public void createPreparedStatement(FlightSql.ActionCreatePreparedStatementRequest request,
+                                      CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getQuery().equals("SELECT PREPARED STATEMENT") ||
+            request.getQuery().equals("UPDATE PREPARED STATEMENT"));
+
+    final FlightSql.ActionCreatePreparedStatementResult
+        result = FlightSql.ActionCreatePreparedStatementResult.newBuilder()
+        .setPreparedStatementHandle(ByteString.copyFromUtf8(request.getQuery() + " HANDLE"))
+        .build();
+    listener.onNext(new Result(pack(result).toByteArray()));
+    listener.onCompleted();
+  }
+
+  @Override
+  public void closePreparedStatement(FlightSql.ActionClosePreparedStatementRequest request,
+                                     CallContext context, StreamListener<Result> listener) {
+    IntegrationAssertions.assertTrue("Expect to be one of the two queries used on tests",
+        request.getPreparedStatementHandle().toStringUtf8().equals("SELECT PREPARED STATEMENT HANDLE") ||
+            request.getPreparedStatementHandle().toStringUtf8().equals("UPDATE PREPARED STATEMENT HANDLE"));
+
+    listener.onCompleted();
+  }
+
+  @Override
+  public FlightInfo getFlightInfoStatement(FlightSql.CommandStatementQuery command,
+                                           CallContext context, FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "SELECT STATEMENT");
+
+    ByteString handle = ByteString.copyFromUtf8("SELECT STATEMENT HANDLE");
+
+    FlightSql.TicketStatementQuery ticket = FlightSql.TicketStatementQuery.newBuilder()
+        .setStatementHandle(handle)
+        .build();
+    return getFlightInfoForSchema(ticket, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public FlightInfo getFlightInfoPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                                   CallContext context,
+                                                   FlightDescriptor descriptor) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return getFlightInfoForSchema(command, descriptor, getQuerySchema());
+  }
+
+  @Override
+  public SchemaResult getSchemaStatement(FlightSql.CommandStatementQuery command,
+                                         CallContext context, FlightDescriptor descriptor) {
+    return new SchemaResult(getQuerySchema());
+  }
+
+  @Override
+  public void getStreamStatement(FlightSql.TicketStatementQuery ticket, CallContext context,
+                                 ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  @Override
+  public void getStreamPreparedStatement(FlightSql.CommandPreparedStatementQuery command,
+                                         CallContext context, ServerStreamListener listener) {
+    serveJsonToStreamListener(listener, getQuerySchema());
+  }
+
+  private Runnable acceptPutReturnConstant(StreamListener<PutResult> ackStream, int value) {
+    return () -> {
+      final FlightSql.DoPutUpdateResult build =
+          FlightSql.DoPutUpdateResult.newBuilder().setRecordCount(value).build();
+
+      try (final ArrowBuf buffer = allocator.buffer(build.getSerializedSize())) {
+        buffer.writeBytes(build.toByteArray());
+        ackStream.onNext(PutResult.metadata(buffer));
+        ackStream.onCompleted();
+      }
+    };
+  }
+
+  @Override
+  public Runnable acceptPutStatement(FlightSql.CommandStatementUpdate command, CallContext context,
+                                     FlightStream flightStream,
+                                     StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getQuery(), "UPDATE STATEMENT");
+
+    return acceptPutReturnConstant(ackStream, 10000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementUpdate(FlightSql.CommandPreparedStatementUpdate command,
+                                                   CallContext context, FlightStream flightStream,
+                                                   StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle().toStringUtf8(),
+        "UPDATE PREPARED STATEMENT HANDLE");
+
+    return acceptPutReturnConstant(ackStream, 20000);
+  }
+
+  @Override
+  public Runnable acceptPutPreparedStatementQuery(FlightSql.CommandPreparedStatementQuery command,
+                                                  CallContext context, FlightStream flightStream,
+                                                  StreamListener<PutResult> ackStream) {
+    IntegrationAssertions.assertEquals(command.getPreparedStatementHandle(),
+        "SELECT PREPARED STATEMENT HANDLE");
+
+    return null;

Review comment:
       Implemented this case as well




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998773336


   Hmm, I wonder if the test failure in Dev is just because it's getting confused since we're on a branch. I think we can resolve that when we make the final PR from flight-sql into master, then.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998295185


   There's some lint failures:
   
   ```
   
   [1/1] cd /tmp/arrow-lint-6fb2hzaq/cpp-build && /usr/local/bin/python /arrow/cpp/build-support/lint_cpp_cli.py /arrow/cpp/src
   INFO:archery:Running Python formatter (autopep8)
   INFO:archery:Running Python linter (flake8)
   /arrow/dev/archery/archery/integration/tester_java.py:52:80: E501 line too long (95 > 79 characters)
   /arrow/dev/archery/archery/integration/runner.py:387:80: E501 line too long (82 > 79 characters)
   INFO:archery:Running cmake-format linters
   ERROR __main__.py:618: Check failed: /arrow/cpp/src/arrow/CMakeLists.txt
   ```
   
   I think the other one is unrelated, but could you rebase on master to see if it goes away?


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998323182


   Ops @lidavidm ... Sorry, I don't know if I messed up when `git rebase master` or if also need to run it in the base branch `flight-sql`...


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-999103832


   Mega-PR is up at https://github.com/apache/arrow/pull/12013


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998975102


   Hey @lidavidm ! Are we good to merge this on `flight-sql` branch and then on `master`?
   
   Thank you


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998335552


   Hmm. I just rebased the branch…so do you mind trying to rebase again here? Sorry for the trouble.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998345145


   Well, looks like the Docker registry isn't doing so well right now. I'll kick all the CI pipelines in a bit to try again.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] github-actions[bot] commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

github-actions[bot] commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-997050056


   https://issues.apache.org/jira/browse/ARROW-15112


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998979560


   We should be. I just kicked off CI.


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r773374837



##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,685 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>

Review comment:
       Also, what are we using json_integration.h for?

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,685 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>
+#include <iostream>
+#include <memory>
+#include <string>
+#include <utility>
+#include <vector>
+
+namespace arrow {
+namespace flight {
+namespace integration_tests {
+
+/// \brief The server for the basic auth integration test.
+class AuthBasicProtoServer : public FlightServerBase {
+  Status DoAction(const ServerCallContext& context, const Action& action,
+                  std::unique_ptr<ResultStream>* result) override {
+    // Respond with the authenticated username.
+    auto buf = Buffer::FromString(context.peer_identity());
+    *result = std::unique_ptr<ResultStream>(new SimpleResultStream({Result{buf}}));
+    return Status::OK();
+  }
+};
+
+/// Validate the result of a DoAction.
+Status CheckActionResults(FlightClient* client, const Action& action,
+                          std::vector<std::string> results) {
+  std::unique_ptr<ResultStream> stream;
+  RETURN_NOT_OK(client->DoAction(action, &stream));
+  std::unique_ptr<Result> result;
+  for (const std::string& expected : results) {
+    RETURN_NOT_OK(stream->Next(&result));
+    if (!result) {
+      return Status::Invalid("Action result stream ended early");
+    }
+    const auto actual = result->body->ToString();
+    if (expected != actual) {
+      return Status::Invalid("Got wrong result; expected", expected, "but got", actual);
+    }
+  }
+  RETURN_NOT_OK(stream->Next(&result));
+  if (result) {
+    return Status::Invalid("Action result stream had too many entries");
+  }
+  return Status::OK();
+}
+
+// The expected username for the basic auth integration test.
+constexpr auto kAuthUsername = "arrow";
+// The expected password for the basic auth integration test.
+constexpr auto kAuthPassword = "flight";
+
+/// \brief A scenario testing the basic auth protobuf.
+class AuthBasicProtoScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new AuthBasicProtoServer());
+    options->auth_handler =
+        std::make_shared<TestServerBasicAuthHandler>(kAuthUsername, kAuthPassword);
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    Action action;
+    std::unique_ptr<ResultStream> stream;
+    std::shared_ptr<FlightStatusDetail> detail;
+    const auto& status = client->DoAction(action, &stream);
+    detail = FlightStatusDetail::UnwrapStatus(status);
+    // This client is unauthenticated and should fail.
+    if (detail == nullptr) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", status.ToString());
+    }
+    if (detail->code() != FlightStatusCode::Unauthenticated) {
+      return Status::Invalid("Expected UNAUTHENTICATED but got ", detail->ToString());
+    }
+
+    auto client_handler = std::unique_ptr<ClientAuthHandler>(
+        new TestClientBasicAuthHandler(kAuthUsername, kAuthPassword));
+    RETURN_NOT_OK(client->Authenticate({}, std::move(client_handler)));
+    return CheckActionResults(client.get(), action, {kAuthUsername});
+  }
+};
+
+/// \brief Test middleware that echoes back the value of a particular
+/// incoming header.
+///
+/// In Java, gRPC may consolidate this header with HTTP/2 trailers if
+/// the call fails, but C++ generally doesn't do this. The integration
+/// test confirms the presence of this header to ensure we can read it
+/// regardless of what gRPC does.
+class TestServerMiddleware : public ServerMiddleware {
+ public:
+  explicit TestServerMiddleware(std::string received) : received_(received) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) override {
+    outgoing_headers->AddHeader("x-middleware", received_);
+  }
+
+  void CallCompleted(const Status& status) override {}
+
+  std::string name() const override { return "GrpcTrailersMiddleware"; }
+
+ private:
+  std::string received_;
+};
+
+class TestServerMiddlewareFactory : public ServerMiddlewareFactory {
+ public:
+  Status StartCall(const CallInfo& info, const CallHeaders& incoming_headers,
+                   std::shared_ptr<ServerMiddleware>* middleware) override {
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    std::string received = "";
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      received = std::string(value);
+    }
+    *middleware = std::make_shared<TestServerMiddleware>(received);
+    return Status::OK();
+  }
+};
+
+/// \brief Test middleware that adds a header on every outgoing call,
+/// and gets the value of the expected header sent by the server.
+class TestClientMiddleware : public ClientMiddleware {
+ public:
+  explicit TestClientMiddleware(std::string* received_header)
+      : received_header_(received_header) {}
+
+  void SendingHeaders(AddCallHeaders* outgoing_headers) {
+    outgoing_headers->AddHeader("x-middleware", "expected value");
+  }
+
+  void ReceivedHeaders(const CallHeaders& incoming_headers) {
+    // We expect the server to always send this header. gRPC/Java may
+    // send it in trailers instead of headers, so we expect Flight to
+    // account for this.
+    const std::pair<CallHeaders::const_iterator, CallHeaders::const_iterator>& iter_pair =
+        incoming_headers.equal_range("x-middleware");
+    if (iter_pair.first != iter_pair.second) {
+      const util::string_view& value = (*iter_pair.first).second;
+      *received_header_ = std::string(value);
+    }
+  }
+
+  void CallCompleted(const Status& status) {}
+
+ private:
+  std::string* received_header_;
+};
+
+class TestClientMiddlewareFactory : public ClientMiddlewareFactory {
+ public:
+  void StartCall(const CallInfo& info, std::unique_ptr<ClientMiddleware>* middleware) {
+    *middleware =
+        std::unique_ptr<ClientMiddleware>(new TestClientMiddleware(&received_header_));
+  }
+
+  std::string received_header_;
+};
+
+/// \brief The server used for testing middleware. Implements only one
+/// endpoint, GetFlightInfo, in such a way that it either succeeds or
+/// returns an error based on the input, in order to test both paths.
+class MiddlewareServer : public FlightServerBase {
+  Status GetFlightInfo(const ServerCallContext& context,
+                       const FlightDescriptor& descriptor,
+                       std::unique_ptr<FlightInfo>* result) override {
+    if (descriptor.type == FlightDescriptor::DescriptorType::CMD &&
+        descriptor.cmd == "success") {
+      // Don't fail
+      std::shared_ptr<Schema> schema = arrow::schema({});
+      Location location;
+      // Return a fake location - the test doesn't read it
+      RETURN_NOT_OK(Location::ForGrpcTcp("localhost", 10010, &location));
+      std::vector<FlightEndpoint> endpoints{FlightEndpoint{{"foo"}, {location}}};
+      ARROW_ASSIGN_OR_RAISE(auto info,
+                            FlightInfo::Make(*schema, descriptor, endpoints, -1, -1));
+      *result = std::unique_ptr<FlightInfo>(new FlightInfo(info));
+      return Status::OK();
+    }
+    // Fail the call immediately. In some gRPC implementations, this
+    // means that gRPC sends only HTTP/2 trailers and not headers. We want
+    // Flight middleware to be agnostic to this difference.
+    return Status::UnknownError("Unknown");
+  }
+};
+
+/// \brief The middleware scenario.
+///
+/// This tests that the server and client get expected header values.
+class MiddlewareScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    options->middleware.push_back(
+        {"grpc_trailers", std::make_shared<TestServerMiddlewareFactory>()});
+    server->reset(new MiddlewareServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override {
+    client_middleware_ = std::make_shared<TestClientMiddlewareFactory>();
+    options->middleware.push_back(client_middleware_);
+    return Status::OK();
+  }
+
+  Status RunClient(std::unique_ptr<FlightClient> client) override {
+    std::unique_ptr<FlightInfo> info;
+    // This call is expected to fail. In gRPC/Java, this causes the
+    // server to combine headers and HTTP/2 trailers, so to read the
+    // expected header, Flight must check for both headers and
+    // trailers.
+    if (client->GetFlightInfo(FlightDescriptor::Command(""), &info).ok()) {
+      return Status::Invalid("Expected call to fail");
+    }
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got: '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on failing call." << std::endl;
+
+    // This call should succeed
+    client_middleware_->received_header_ = "";
+    RETURN_NOT_OK(client->GetFlightInfo(FlightDescriptor::Command("success"), &info));
+    if (client_middleware_->received_header_ != "expected value") {
+      return Status::Invalid(
+          "Expected to receive header 'x-middleware: expected value', but instead got '",
+          client_middleware_->received_header_, "'");
+    }
+    std::cerr << "Headers received successfully on passing call." << std::endl;
+    return Status::OK();
+  }
+
+  std::shared_ptr<TestClientMiddlewareFactory> client_middleware_;
+};
+
+/// \brief Schema to be returned for mocking the statement/prepared statement results.
+/// Must be the same across all languages.
+std::shared_ptr<Schema> GetQuerySchema() {
+  return arrow::schema({arrow::field("id", int64())});
+}
+
+constexpr int64_t kUpdateStatementExpectedRows = 10000L;
+constexpr int64_t kUpdatePreparedStatementExpectedRows = 20000L;
+
+template <typename T>
+arrow::Status AssertEq(const T& expected, const T& actual) {
+  if (expected != actual) {
+    return Status::Invalid("Expected \"", expected, "\", got \'", actual, "\"");
+  }
+  return Status::OK();
+}
+
+/// \brief The server used for testing Flight SQL, this implements a static Flight SQL
+/// server which only asserts that commands called during integration tests are being
+/// parsed correctly and returns the expected schemas to be validated on client.
+class FlightSqlScenarioServer : public sql::FlightSqlServerBase {
+ public:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoStatement(
+      const ServerCallContext& context, const sql::StatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT STATEMENT", command.query));
+
+    ARROW_ASSIGN_OR_RAISE(auto handle,
+                          sql::CreateStatementQueryTicket("SELECT STATEMENT HANDLE"));
+
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{handle}, {}}};
+    ARROW_ASSIGN_OR_RAISE(
+        auto result, FlightInfo::Make(*GetQuerySchema(), descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetStatement(
+      const ServerCallContext& context,
+      const sql::StatementQueryTicket& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPreparedStatement(
+      const ServerCallContext& context, const sql::PreparedStatementQuery& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return GetFlightInfoForCommand(descriptor, GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPreparedStatement(
+      const ServerCallContext& context,
+      const sql::PreparedStatementQuery& command) override {
+    return DoGetForTestCase(GetQuerySchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCatalogs(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCatalogs(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCatalogsSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.info.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_NAME, command.info[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(
+        sql::SqlInfoOptions::SqlInfo::FLIGHT_SQL_SERVER_READ_ONLY, command.info[1]));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetSqlInfo(
+      const ServerCallContext& context, const sql::GetSqlInfo& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetSqlInfoSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetDbSchemas(
+      const ServerCallContext& context, const sql::GetDbSchemas& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetDbSchemasSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTables(
+      const ServerCallContext& context, const sql::GetTables& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("catalog", command.catalog.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("db_schema_filter_pattern",
+                                              command.db_schema_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table_filter_pattern",
+                                              command.table_name_filter_pattern.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<int>(2, command.table_types.size()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_types[0]));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("view", command.table_types[1]));
+    ARROW_RETURN_NOT_OK(AssertEq<bool>(true, command.include_schema));
+
+    return GetFlightInfoForCommand(descriptor,
+                                   sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTables(
+      const ServerCallContext& context, const sql::GetTables& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTablesSchemaWithIncludedSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoTableTypes(
+      const ServerCallContext& context, const FlightDescriptor& descriptor) override {
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetTableTypes(
+      const ServerCallContext& context) override {
+    return DoGetForTestCase(sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetPrimaryKeys(
+      const ServerCallContext& context, const sql::GetPrimaryKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetPrimaryKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetExportedKeys(
+      const ServerCallContext& context, const sql::GetExportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetExportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("catalog", command.table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("db_schema", command.table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("table", command.table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetImportedKeys(
+      const ServerCallContext& context, const sql::GetImportedKeys& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetImportedKeysSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command,
+      const FlightDescriptor& descriptor) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_catalog", command.pk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("pk_db_schema", command.pk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("pk_table", command.pk_table_ref.table));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_catalog", command.fk_table_ref.catalog.value()));
+    ARROW_RETURN_NOT_OK(
+        AssertEq<std::string>("fk_db_schema", command.fk_table_ref.db_schema.value()));
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("fk_table", command.fk_table_ref.table));
+
+    return GetFlightInfoForCommand(descriptor, sql::SqlSchema::GetTableTypesSchema());
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetCrossReference(
+      const ServerCallContext& context, const sql::GetCrossReference& command) override {
+    return DoGetForTestCase(sql::SqlSchema::GetCrossReferenceSchema());
+  }
+
+  arrow::Result<int64_t> DoPutCommandStatementUpdate(
+      const ServerCallContext& context, const sql::StatementUpdate& command) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE STATEMENT", command.query));
+
+    return kUpdateStatementExpectedRows;
+  }
+
+  arrow::Result<sql::ActionCreatePreparedStatementResult> CreatePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionCreatePreparedStatementRequest& request) override {
+    ARROW_RETURN_NOT_OK(
+        AssertEq<bool>(true, request.query == "SELECT PREPARED STATEMENT" ||
+                                 request.query == "UPDATE PREPARED STATEMENT"));
+
+    sql::ActionCreatePreparedStatementResult result;
+    result.prepared_statement_handle = request.query + " HANDLE";
+
+    return result;
+  }
+
+  Status ClosePreparedStatement(
+      const ServerCallContext& context,
+      const sql::ActionClosePreparedStatementRequest& request) override {
+    return Status::OK();
+  }
+
+  Status DoPutPreparedStatementQuery(const ServerCallContext& context,
+                                     const sql::PreparedStatementQuery& command,
+                                     FlightMessageReader* reader,
+                                     FlightMetadataWriter* writer) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("SELECT PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    ARROW_ASSIGN_OR_RAISE(auto actual_schema, reader->GetSchema());
+    ARROW_RETURN_NOT_OK(AssertEq<Schema>(*GetQuerySchema(), *actual_schema));
+
+    return Status::OK();
+  }
+
+  arrow::Result<int64_t> DoPutPreparedStatementUpdate(
+      const ServerCallContext& context, const sql::PreparedStatementUpdate& command,
+      FlightMessageReader* reader) override {
+    ARROW_RETURN_NOT_OK(AssertEq<std::string>("UPDATE PREPARED STATEMENT HANDLE",
+                                              command.prepared_statement_handle));
+
+    return kUpdatePreparedStatementExpectedRows;
+  }
+
+ private:
+  arrow::Result<std::unique_ptr<FlightInfo>> GetFlightInfoForCommand(
+      const FlightDescriptor& descriptor, const std::shared_ptr<Schema>& schema) {
+    std::vector<FlightEndpoint> endpoints{FlightEndpoint{{descriptor.cmd}, {}}};
+    ARROW_ASSIGN_OR_RAISE(auto result,
+                          FlightInfo::Make(*schema, descriptor, endpoints, -1, -1))
+
+    return std::unique_ptr<FlightInfo>(new FlightInfo(result));
+  }
+
+  arrow::Result<std::unique_ptr<FlightDataStream>> DoGetForTestCase(
+      const std::shared_ptr<Schema>& schema) {
+    ARROW_ASSIGN_OR_RAISE(auto reader, RecordBatchReader::Make({}, schema));
+    return std::unique_ptr<FlightDataStream>(new RecordBatchStream(reader));
+  }
+};
+
+/// \brief Integration test scenario for validating Flight SQL specs across multiple
+/// implementations. This should ensure that RPC objects are being built and parsed
+/// correctly for multiple languages and that the Arrow schemas are returned as expected.
+class FlightSqlScenario : public Scenario {
+  Status MakeServer(std::unique_ptr<FlightServerBase>* server,
+                    FlightServerOptions* options) override {
+    server->reset(new FlightSqlScenarioServer());
+    return Status::OK();
+  }
+
+  Status MakeClient(FlightClientOptions* options) override { return Status::OK(); }
+
+  Status Validate(std::shared_ptr<Schema> expectedSchema,
+                  arrow::Result<std::unique_ptr<FlightInfo>> flightInfo,
+                  sql::FlightSqlClient* sql_client) {
+    FlightCallOptions call_options;
+
+    ARROW_ASSIGN_OR_RAISE(auto flight_info, flightInfo);
+    ARROW_ASSIGN_OR_RAISE(
+        auto reader, sql_client->DoGet(call_options, flight_info->endpoints()[0].ticket));
+
+    ARROW_ASSIGN_OR_RAISE(auto actual_schema, reader->GetSchema());
+
+    AssertSchemaEqual(expectedSchema, actual_schema);

Review comment:
       nit: `expected_schema`

##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,685 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>

Review comment:
       Sorry, I just noticed these: can we use `""` includes for Arrow headers?




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r773380088



##########
File path: cpp/src/arrow/flight/integration_tests/test_integration.cc
##########
@@ -0,0 +1,685 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "arrow/flight/integration_tests/test_integration.h"
+#include "arrow/flight/client_middleware.h"
+#include "arrow/flight/server_middleware.h"
+#include "arrow/flight/sql/client.h"
+#include "arrow/flight/sql/server.h"
+#include "arrow/flight/test_util.h"
+#include "arrow/flight/types.h"
+#include "arrow/ipc/dictionary.h"
+
+#include <arrow/testing/gtest_util.h>
+#include <arrow/testing/json_integration.h>

Review comment:
       Fixed!
   
   Actually didn't need json_integration.h, sorry




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998137221


   Done @lidavidm 


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998338029


   No problem! Just rebased :)


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-998334422


   Ah, right, we have a branch…How about I rebase the branch, and then we can rebase the PRs?


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-997949857


   We need to exclude Go and Rust from the new test scenario.
   
   Also, the PR needs to be formatted (see https://arrow.apache.org/docs/developers/cpp/development.html#code-style-linting-and-ci, note that we use ClangTools 12 now not version 8)


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r773195206



##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenario.java
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import java.util.Arrays;
+
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.flight.sql.util.TableRef;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+/**
+ * Integration test scenario for validating Flight SQL specs across multiple implementations.
+ * This should ensure that RPC objects are being built and parsed correctly for multiple languages
+ * and that the Arrow schemas are returned as expected.
+ */
+public class FlightSqlScenario implements Scenario {
+
+  @Override
+  public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception {
+    return new FlightSqlScenarioProducer(allocator);
+  }
+
+  @Override
+  public void buildServer(FlightServer.Builder builder) throws Exception {
+
+  }
+
+  @Override
+  public void client(BufferAllocator allocator, Location location, FlightClient client)
+      throws Exception {
+    final FlightSqlClient sqlClient = new FlightSqlClient(client);
+
+    validateMetadataRetrieval(sqlClient);
+
+    validateStatementExecution(sqlClient);
+
+    validatePreparedStatementExecution(sqlClient, allocator);
+  }
+
+  private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA, sqlClient.getCatalogs(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA,
+        sqlClient.getSchemas("catalog", "db_schema_filter_pattern", options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA,
+        sqlClient.getTables("catalog", "db_schema_filter_pattern", "table_filter_pattern",
+            Arrays.asList("table", "view"), true, options), sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA, sqlClient.getTableTypes(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_PRIMARY_KEYS_SCHEMA,
+        sqlClient.getPrimaryKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_EXPORTED_KEYS_SCHEMA,
+        sqlClient.getExportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_IMPORTED_KEYS_SCHEMA,
+        sqlClient.getImportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_CROSS_REFERENCE_SCHEMA,
+        sqlClient.getCrossReference(TableRef.of("pk_catalog", "pk_db_schema", "pk_table"),
+            TableRef.of("fk_catalog", "fk_db_schema", "fk_table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA,
+        sqlClient.getSqlInfo(new FlightSql.SqlInfo[] {FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME,
+            FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY}, options), sqlClient);
+  }
+
+  private void validateStatementExecution(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlScenarioProducer.getQuerySchema(),
+        sqlClient.execute("SELECT STATEMENT", options), sqlClient);
+
+    IntegrationAssertions.assertEquals(sqlClient.executeUpdate("UPDATE STATEMENT", options),
+        10000L);

Review comment:
       Thanks! Done




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] rafael-telles commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

rafael-telles commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-999094629


   Glad to see everything passing! Except to that "Dev/ Source Release and Merge Script", but it indeed seem to fail because we are in our own branch


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm closed pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm closed pull request #11989:
URL: https://github.com/apache/arrow/pull/11989


   


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on a change in pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on a change in pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#discussion_r773142629



##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenario.java
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import java.util.Arrays;
+
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.flight.sql.util.TableRef;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+/**
+ * Integration test scenario for validating Flight SQL specs across multiple implementations.
+ * This should ensure that RPC objects are being built and parsed correctly for multiple languages
+ * and that the Arrow schemas are returned as expected.
+ */
+public class FlightSqlScenario implements Scenario {
+
+  @Override
+  public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception {
+    return new FlightSqlScenarioProducer(allocator);
+  }
+
+  @Override
+  public void buildServer(FlightServer.Builder builder) throws Exception {
+
+  }
+
+  @Override
+  public void client(BufferAllocator allocator, Location location, FlightClient client)
+      throws Exception {
+    final FlightSqlClient sqlClient = new FlightSqlClient(client);
+
+    validateMetadataRetrieval(sqlClient);
+
+    validateStatementExecution(sqlClient);
+
+    validatePreparedStatementExecution(sqlClient, allocator);
+  }
+
+  private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA, sqlClient.getCatalogs(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA,
+        sqlClient.getSchemas("catalog", "db_schema_filter_pattern", options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA,
+        sqlClient.getTables("catalog", "db_schema_filter_pattern", "table_filter_pattern",
+            Arrays.asList("table", "view"), true, options), sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA, sqlClient.getTableTypes(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_PRIMARY_KEYS_SCHEMA,
+        sqlClient.getPrimaryKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_EXPORTED_KEYS_SCHEMA,
+        sqlClient.getExportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_IMPORTED_KEYS_SCHEMA,
+        sqlClient.getImportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_CROSS_REFERENCE_SCHEMA,
+        sqlClient.getCrossReference(TableRef.of("pk_catalog", "pk_db_schema", "pk_table"),
+            TableRef.of("fk_catalog", "fk_db_schema", "fk_table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA,
+        sqlClient.getSqlInfo(new FlightSql.SqlInfo[] {FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME,
+            FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY}, options), sqlClient);
+  }
+
+  private void validateStatementExecution(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlScenarioProducer.getQuerySchema(),
+        sqlClient.execute("SELECT STATEMENT", options), sqlClient);
+
+    IntegrationAssertions.assertEquals(sqlClient.executeUpdate("UPDATE STATEMENT", options),
+        10000L);

Review comment:
       Just a nit, but can we move these numbers to constants so it's clear what they represent? (`UPDATE_STATEMENT_EXPECTED_ROWS` or something)

##########
File path: java/flight/flight-integration-tests/src/main/java/org/apache/arrow/flight/integration/tests/FlightSqlScenario.java
##########
@@ -0,0 +1,136 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.arrow.flight.integration.tests;
+
+import java.util.Arrays;
+
+import org.apache.arrow.flight.CallOption;
+import org.apache.arrow.flight.FlightClient;
+import org.apache.arrow.flight.FlightInfo;
+import org.apache.arrow.flight.FlightProducer;
+import org.apache.arrow.flight.FlightServer;
+import org.apache.arrow.flight.FlightStream;
+import org.apache.arrow.flight.Location;
+import org.apache.arrow.flight.Ticket;
+import org.apache.arrow.flight.sql.FlightSqlClient;
+import org.apache.arrow.flight.sql.FlightSqlProducer;
+import org.apache.arrow.flight.sql.impl.FlightSql;
+import org.apache.arrow.flight.sql.util.TableRef;
+import org.apache.arrow.memory.BufferAllocator;
+import org.apache.arrow.vector.VectorSchemaRoot;
+import org.apache.arrow.vector.types.pojo.Schema;
+
+/**
+ * Integration test scenario for validating Flight SQL specs across multiple implementations.
+ * This should ensure that RPC objects are being built and parsed correctly for multiple languages
+ * and that the Arrow schemas are returned as expected.
+ */
+public class FlightSqlScenario implements Scenario {
+
+  @Override
+  public FlightProducer producer(BufferAllocator allocator, Location location) throws Exception {
+    return new FlightSqlScenarioProducer(allocator);
+  }
+
+  @Override
+  public void buildServer(FlightServer.Builder builder) throws Exception {
+
+  }
+
+  @Override
+  public void client(BufferAllocator allocator, Location location, FlightClient client)
+      throws Exception {
+    final FlightSqlClient sqlClient = new FlightSqlClient(client);
+
+    validateMetadataRetrieval(sqlClient);
+
+    validateStatementExecution(sqlClient);
+
+    validatePreparedStatementExecution(sqlClient, allocator);
+  }
+
+  private void validateMetadataRetrieval(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlProducer.Schemas.GET_CATALOGS_SCHEMA, sqlClient.getCatalogs(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SCHEMAS_SCHEMA,
+        sqlClient.getSchemas("catalog", "db_schema_filter_pattern", options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLES_SCHEMA,
+        sqlClient.getTables("catalog", "db_schema_filter_pattern", "table_filter_pattern",
+            Arrays.asList("table", "view"), true, options), sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_TABLE_TYPES_SCHEMA, sqlClient.getTableTypes(options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_PRIMARY_KEYS_SCHEMA,
+        sqlClient.getPrimaryKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_EXPORTED_KEYS_SCHEMA,
+        sqlClient.getExportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_IMPORTED_KEYS_SCHEMA,
+        sqlClient.getImportedKeys(TableRef.of("catalog", "db_schema", "table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_CROSS_REFERENCE_SCHEMA,
+        sqlClient.getCrossReference(TableRef.of("pk_catalog", "pk_db_schema", "pk_table"),
+            TableRef.of("fk_catalog", "fk_db_schema", "fk_table"), options),
+        sqlClient);
+    validate(FlightSqlProducer.Schemas.GET_SQL_INFO_SCHEMA,
+        sqlClient.getSqlInfo(new FlightSql.SqlInfo[] {FlightSql.SqlInfo.FLIGHT_SQL_SERVER_NAME,
+            FlightSql.SqlInfo.FLIGHT_SQL_SERVER_READ_ONLY}, options), sqlClient);
+  }
+
+  private void validateStatementExecution(FlightSqlClient sqlClient) throws Exception {
+    final CallOption[] options = new CallOption[0];
+
+    validate(FlightSqlScenarioProducer.getQuerySchema(),
+        sqlClient.execute("SELECT STATEMENT", options), sqlClient);
+
+    IntegrationAssertions.assertEquals(sqlClient.executeUpdate("UPDATE STATEMENT", options),
+        10000L);

Review comment:
       (The same goes for the producer, C++, etc.)




-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org

[GitHub] [arrow] lidavidm commented on pull request #11989: ARROW-15112: [Integration][C++][Java] Implement Flight SQL integration tests

[GitBox <gi...@apache.org>]

lidavidm commented on pull request #11989:
URL: https://github.com/apache/arrow/pull/11989#issuecomment-999009709


   I think the failures are spurious. (I filed ARROW-15181 for the crash in the Flight tests.)


-- 
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: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org