/*
 * 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.driver.jdbc;

import static java.lang.String.format;
import static java.util.Collections.synchronizedSet;
import static org.apache.arrow.flight.Location.forGrpcInsecure;
import static org.hamcrest.CoreMatchers.allOf;
import static org.hamcrest.CoreMatchers.anyOf;
import static org.hamcrest.CoreMatchers.containsString;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.is;
import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

import com.google.common.collect.ImmutableSet;
import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLTimeoutException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import java.util.Random;
import java.util.Set;
import java.util.concurrent.CountDownLatch;
import org.apache.arrow.driver.jdbc.utils.CoreMockedSqlProducers;
import org.apache.arrow.driver.jdbc.utils.FallbackFlightSqlProducer;
import org.apache.arrow.driver.jdbc.utils.PartitionedFlightSqlProducer;
import org.apache.arrow.flight.FlightEndpoint;
import org.apache.arrow.flight.FlightProducer;
import org.apache.arrow.flight.FlightRuntimeException;
import org.apache.arrow.flight.FlightServer;
import org.apache.arrow.flight.FlightStatusCode;
import org.apache.arrow.flight.Location;
import org.apache.arrow.flight.Ticket;
import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.memory.RootAllocator;
import org.apache.arrow.vector.IntVector;
import org.apache.arrow.vector.VectorSchemaRoot;
import org.apache.arrow.vector.types.Types;
import org.apache.arrow.vector.types.pojo.ArrowType;
import org.apache.arrow.vector.types.pojo.Field;
import org.apache.arrow.vector.types.pojo.Schema;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;

public class ResultSetTest {
  private static final Random RANDOM = new Random(10);

  @RegisterExtension
  public static final FlightServerTestExtension FLIGHT_SERVER_TEST_EXTENSION =
      FlightServerTestExtension.createStandardTestExtension(
          CoreMockedSqlProducers.getLegacyProducer());

  private static Connection connection;

  @BeforeAll
  public static void setup() throws SQLException {
    connection = FLIGHT_SERVER_TEST_EXTENSION.getConnection(false);
  }

  @AfterAll
  public static void tearDown() throws SQLException {
    connection.close();
  }

  private static void resultSetNextUntilDone(ResultSet resultSet) throws SQLException {
    while (resultSet.next()) {
      // TODO: implement resultSet.last()
      // Pass to the next until resultSet is done
    }
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} can run a query successfully.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldRunSelectQuery() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
      CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet);
    }
  }

  @Test
  public void testShouldExecuteQueryNotBlockIfClosedBeforeEnd() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {

      for (int i = 0; i < 7500; i++) {
        assertTrue(resultSet.next());
      }
    }
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} query only returns only the amount of value set
   * by {@link org.apache.calcite.avatica.AvaticaStatement#setMaxRows(int)}.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldRunSelectQuerySettingMaxRowLimit() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {

      final int maxRowsLimit = 3;
      statement.setMaxRows(maxRowsLimit);

      assertThat(statement.getMaxRows(), is(maxRowsLimit));

      int count = 0;
      int columns = 6;
      for (; resultSet.next(); count++) {
        for (int column = 1; column <= columns; column++) {
          resultSet.getObject(column);
        }
        assertThat("Test Name #" + count, is(resultSet.getString(2)));
      }

      assertThat(maxRowsLimit, is(count));
    }
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} fails upon attempting to run an invalid query.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldThrowExceptionUponAttemptingToExecuteAnInvalidSelectQuery() {
    assertThrows(
        SQLException.class,
        () -> {
          try (Statement statement = connection.createStatement();
              ResultSet resultSet = statement.executeQuery("SELECT * FROM SHOULD-FAIL")) {
            fail();
          }
        });
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} query only returns only the amount of value set
   * by {@link org.apache.calcite.avatica.AvaticaStatement#setLargeMaxRows(long)} (int)}.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldRunSelectQuerySettingLargeMaxRowLimit() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
      final long maxRowsLimit = 3;
      statement.setLargeMaxRows(maxRowsLimit);

      assertThat(statement.getLargeMaxRows(), is(maxRowsLimit));

      int count = 0;
      int columns = resultSet.getMetaData().getColumnCount();
      for (; resultSet.next(); count++) {
        for (int column = 1; column <= columns; column++) {
          resultSet.getObject(column);
        }
        assertEquals("Test Name #" + count, resultSet.getString(2));
      }

      assertEquals(maxRowsLimit, count);
    }
  }

  @Test
  public void testColumnCountShouldRemainConsistentForResultSetThroughoutEntireDuration()
      throws SQLException {
    final Set<Integer> counts = new HashSet<>();
    try (final Statement statement = connection.createStatement();
        final ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
      while (resultSet.next()) {
        counts.add(resultSet.getMetaData().getColumnCount());
      }
    }
    assertThat(counts, is(ImmutableSet.of(6)));
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet
   * when call {@link org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldCloseStatementWhenIsCloseOnCompletion() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {

      statement.closeOnCompletion();

      resultSetNextUntilDone(resultSet);

      assertThat(statement.isClosed(), is(true));
    }
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} close the statement after complete ResultSet
   * with max rows limit when call {@link
   * org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldCloseStatementWhenIsCloseOnCompletionWithMaxRowsLimit() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {

      final long maxRowsLimit = 3;
      statement.setLargeMaxRows(maxRowsLimit);
      statement.closeOnCompletion();

      resultSetNextUntilDone(resultSet);

      assertThat(statement.isClosed(), is(true));
    }
  }

  /**
   * Tests whether the {@link ArrowFlightJdbcDriver} not close the statement after complete
   * ResultSet with max rows limit when call {@link
   * org.apache.calcite.avatica.AvaticaStatement#closeOnCompletion()}.
   *
   * @throws Exception If the connection fails to be established.
   */
  @Test
  public void testShouldNotCloseStatementWhenIsNotCloseOnCompletionWithMaxRowsLimit()
      throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {

      final long maxRowsLimit = 3;
      statement.setLargeMaxRows(maxRowsLimit);

      assertThat(statement.isClosed(), is(false));
      resultSetNextUntilDone(resultSet);
      assertThat(resultSet.isClosed(), is(false));
      assertThat(resultSet, is(instanceOf(ArrowFlightJdbcFlightStreamResultSet.class)));
    }
  }

  @Test
  public void testShouldCancelQueryUponCancelAfterQueryingResultSet() throws SQLException {
    try (final Statement statement = connection.createStatement();
        final ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
      final int column = RANDOM.nextInt(resultSet.getMetaData().getColumnCount()) + 1;
      assertThat(resultSet.isClosed(), is(false));
      assertThat(resultSet.next(), is(true));
      assertDoesNotThrow(() -> resultSet.getObject(column));
      statement.cancel();
      // Should reset `ResultSet`; keep both `ResultSet` and `Connection` open.
      assertThat(statement.isClosed(), is(false));
      assertThat(resultSet.isClosed(), is(false));
      assertThat(resultSet.getMetaData().getColumnCount(), is(0));
    }
  }

  @Test
  public void testShouldInterruptFlightStreamsIfQueryIsCancelledMidQuerying()
      throws SQLException, InterruptedException {
    try (final Statement statement = connection.createStatement()) {
      final CountDownLatch latch = new CountDownLatch(1);
      final Set<Exception> exceptions = synchronizedSet(new HashSet<>(1));
      final Thread thread =
          new Thread(
              () -> {
                try (final ResultSet resultSet =
                    statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
                  final int cachedColumnCount = resultSet.getMetaData().getColumnCount();
                  Thread.sleep(300);
                  while (resultSet.next()) {
                    resultSet.getObject(RANDOM.nextInt(cachedColumnCount) + 1);
                  }
                } catch (final SQLException | InterruptedException e) {
                  exceptions.add(e);
                } finally {
                  latch.countDown();
                }
              });
      thread.setName("Test Case: interrupt query execution before first retrieval");
      thread.start();
      statement.cancel();
      thread.join();
      assertThat(
          exceptions.stream()
              .map(Exception::getMessage)
              .map(StringBuilder::new)
              .reduce(StringBuilder::append)
              .orElseThrow(IllegalArgumentException::new)
              .toString(),
          is("Statement canceled"));
    }
  }

  @Test
  public void
      testShouldInterruptFlightStreamsIfQueryIsCancelledMidProcessingForTimeConsumingQueries()
          throws SQLException, InterruptedException {
    final String query = CoreMockedSqlProducers.LEGACY_CANCELLATION_SQL_CMD;
    try (final Statement statement = connection.createStatement()) {
      final Set<Exception> exceptions = synchronizedSet(new HashSet<>(1));
      final Thread thread =
          new Thread(
              () -> {
                try (final ResultSet ignored = statement.executeQuery(query)) {
                  fail();
                } catch (final SQLException e) {
                  exceptions.add(e);
                }
              });
      thread.setName("Test Case: interrupt query execution mid-process");
      thread.setPriority(Thread.MAX_PRIORITY);
      thread.start();
      Thread.sleep(5000); // Let the other thread attempt to retrieve results.
      statement.cancel();
      thread.join();
      assertThat(
          exceptions.stream()
              .map(Exception::getMessage)
              .map(StringBuilder::new)
              .reduce(StringBuilder::append)
              .orElseThrow(IllegalStateException::new)
              .toString(),
          anyOf(
              is(format("Error while executing SQL \"%s\": Query canceled", query)),
              allOf(
                  containsString(format("Error while executing SQL \"%s\"", query)),
                  anyOf(containsString("CANCELLED"), containsString("Cancelling")))));
    }
  }

  @Test
  public void testShouldInterruptFlightStreamsIfQueryTimeoutIsOver() throws SQLException {
    final String query = CoreMockedSqlProducers.LEGACY_CANCELLATION_SQL_CMD;
    final int timeoutValue = 2;
    final String timeoutUnit = "SECONDS";
    try (final Statement statement = connection.createStatement()) {
      statement.setQueryTimeout(timeoutValue);
      final Set<Exception> exceptions = new HashSet<>(1);
      try {
        statement.executeQuery(query);
      } catch (final Exception e) {
        exceptions.add(e);
      }
      final Throwable comparisonCause =
          exceptions.stream().findFirst().orElseThrow(RuntimeException::new).getCause().getCause();
      assertThat(comparisonCause, is(instanceOf(SQLTimeoutException.class)));
      assertThat(
          comparisonCause.getMessage(),
          is(format("Query timed out after %d %s", timeoutValue, timeoutUnit)));
    }
  }

  @Test
  public void testFlightStreamsQueryShouldNotTimeout() throws SQLException {
    final String query = CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD;
    final int timeoutValue = 5;
    try (Statement statement = connection.createStatement()) {
      statement.setQueryTimeout(timeoutValue);
      try (ResultSet resultSet = statement.executeQuery(query)) {
        CoreMockedSqlProducers.assertLegacyRegularSqlResultSet(resultSet);
      }
    }
  }

  @Test
  public void testPartitionedFlightServer() throws Exception {
    // Arrange
    final Schema schema =
        new Schema(
            Arrays.asList(Field.nullablePrimitive("int_column", new ArrowType.Int(32, true))));
    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        VectorSchemaRoot firstPartition = VectorSchemaRoot.create(schema, allocator);
        VectorSchemaRoot secondPartition = VectorSchemaRoot.create(schema, allocator)) {
      firstPartition.setRowCount(1);
      ((IntVector) firstPartition.getVector(0)).set(0, 1);
      secondPartition.setRowCount(1);
      ((IntVector) secondPartition.getVector(0)).set(0, 2);

      // Construct the data-only nodes first.
      FlightProducer firstProducer =
          new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer(
              new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstPartition);
      FlightProducer secondProducer =
          new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer(
              new Ticket("second".getBytes(StandardCharsets.UTF_8)), secondPartition);

      final FlightServer.Builder firstBuilder =
          FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), firstProducer);

      final FlightServer.Builder secondBuilder =
          FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), secondProducer);

      // Run the data-only nodes so that we can get the Locations they are running at.
      try (FlightServer firstServer = firstBuilder.build();
          FlightServer secondServer = secondBuilder.build()) {
        firstServer.start();
        secondServer.start();
        final FlightEndpoint firstEndpoint =
            new FlightEndpoint(
                new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstServer.getLocation());

        final FlightEndpoint secondEndpoint =
            new FlightEndpoint(
                new Ticket("second".getBytes(StandardCharsets.UTF_8)), secondServer.getLocation());

        // Finally start the root node.
        try (final PartitionedFlightSqlProducer rootProducer =
                new PartitionedFlightSqlProducer(schema, firstEndpoint, secondEndpoint);
            FlightServer rootServer =
                FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer)
                    .build()
                    .start();
            Connection newConnection =
                DriverManager.getConnection(
                    String.format(
                        "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false",
                        rootServer.getLocation().getUri().getHost(), rootServer.getPort()));
            Statement newStatement = newConnection.createStatement();
            // Act
            ResultSet result = newStatement.executeQuery("Select partitioned_data")) {
          List<Integer> resultData = new ArrayList<>();
          while (result.next()) {
            resultData.add(result.getInt(1));
          }

          // Assert
          assertEquals(
              firstPartition.getRowCount() + secondPartition.getRowCount(), resultData.size());
          assertTrue(resultData.contains(((IntVector) firstPartition.getVector(0)).get(0)));
          assertTrue(resultData.contains(((IntVector) secondPartition.getVector(0)).get(0)));
        }
      }
    }
  }

  @Test
  public void testPartitionedFlightServerIgnoreFailure() throws Exception {
    final Schema schema =
        new Schema(
            Collections.singletonList(
                Field.nullablePrimitive("int_column", new ArrowType.Int(32, true))));
    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
      final FlightEndpoint firstEndpoint =
          new FlightEndpoint(
              new Ticket("first".getBytes(StandardCharsets.UTF_8)),
              Location.forGrpcInsecure("127.0.0.2", 1234),
              Location.forGrpcInsecure("127.0.0.3", 1234));

      try (final PartitionedFlightSqlProducer rootProducer =
              new PartitionedFlightSqlProducer(schema, firstEndpoint);
          FlightServer rootServer =
              FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer)
                  .build()
                  .start();
          Connection newConnection =
              DriverManager.getConnection(
                  String.format(
                      "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false",
                      rootServer.getLocation().getUri().getHost(), rootServer.getPort()));
          Statement newStatement = newConnection.createStatement()) {
        final SQLException e =
            assertThrows(
                SQLException.class,
                () -> {
                  ResultSet result = newStatement.executeQuery("Select partitioned_data");
                  while (result.next()) {}
                });
        final Throwable cause = e.getCause();
        assertTrue(cause instanceof FlightRuntimeException);
        final FlightRuntimeException fre = (FlightRuntimeException) cause;
        assertEquals(FlightStatusCode.UNAVAILABLE, fre.status().code());
      }
    }
  }

  @Test
  public void testPartitionedFlightServerAllFailure() throws Exception {
    // Arrange
    final Schema schema =
        new Schema(
            Collections.singletonList(
                Field.nullablePrimitive("int_column", new ArrowType.Int(32, true))));
    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        VectorSchemaRoot firstPartition = VectorSchemaRoot.create(schema, allocator)) {
      firstPartition.setRowCount(1);
      ((IntVector) firstPartition.getVector(0)).set(0, 1);

      // Construct the data-only nodes first.
      FlightProducer firstProducer =
          new PartitionedFlightSqlProducer.DataOnlyFlightSqlProducer(
              new Ticket("first".getBytes(StandardCharsets.UTF_8)), firstPartition);

      final FlightServer.Builder firstBuilder =
          FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), firstProducer);

      // Run the data-only nodes so that we can get the Locations they are running at.
      try (FlightServer firstServer = firstBuilder.build()) {
        firstServer.start();
        final Location badLocation = Location.forGrpcInsecure("127.0.0.2", 1234);
        final FlightEndpoint firstEndpoint =
            new FlightEndpoint(
                new Ticket("first".getBytes(StandardCharsets.UTF_8)),
                badLocation,
                firstServer.getLocation());

        // Finally start the root node.
        try (final PartitionedFlightSqlProducer rootProducer =
                new PartitionedFlightSqlProducer(schema, firstEndpoint);
            FlightServer rootServer =
                FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer)
                    .build()
                    .start();
            Connection newConnection =
                DriverManager.getConnection(
                    String.format(
                        "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false",
                        rootServer.getLocation().getUri().getHost(), rootServer.getPort()));
            Statement newStatement = newConnection.createStatement();
            // Act
            ResultSet result = newStatement.executeQuery("Select partitioned_data")) {
          List<Integer> resultData = new ArrayList<>();
          while (result.next()) {
            resultData.add(result.getInt(1));
          }

          // Assert
          assertEquals(firstPartition.getRowCount(), resultData.size());
          assertTrue(resultData.contains(((IntVector) firstPartition.getVector(0)).get(0)));
        }
      }
    }
  }

  @Test
  public void testFallbackFlightServer() throws Exception {
    final Schema schema =
        new Schema(
            Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType())));
    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) {
      resultData.setRowCount(1);
      ((IntVector) resultData.getVector(0)).set(0, 1);

      try (final FallbackFlightSqlProducer rootProducer =
              new FallbackFlightSqlProducer(resultData);
          FlightServer rootServer =
              FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer)
                  .build()
                  .start();
          Connection newConnection =
              DriverManager.getConnection(
                  String.format(
                      "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false",
                      rootServer.getLocation().getUri().getHost(), rootServer.getPort()));
          Statement newStatement = newConnection.createStatement();
          ResultSet result = newStatement.executeQuery("fallback")) {
        List<Integer> actualData = new ArrayList<>();
        while (result.next()) {
          actualData.add(result.getInt(1));
        }

        // Assert
        assertEquals(resultData.getRowCount(), actualData.size());
        assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0)));
      }
    }
  }

  @Test
  public void testFallbackSecondFlightServer() throws Exception {
    final Schema schema =
        new Schema(
            Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType())));
    try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE);
        VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) {
      resultData.setRowCount(1);
      ((IntVector) resultData.getVector(0)).set(0, 1);

      try (final FallbackFlightSqlProducer rootProducer =
              new FallbackFlightSqlProducer(resultData);
          FlightServer rootServer =
              FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer)
                  .build()
                  .start();
          Connection newConnection =
              DriverManager.getConnection(
                  String.format(
                      "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false",
                      rootServer.getLocation().getUri().getHost(), rootServer.getPort()));
          Statement newStatement = newConnection.createStatement();
          ResultSet result = newStatement.executeQuery("fallback with error")) {
        List<Integer> actualData = new ArrayList<>();
        while (result.next()) {
          actualData.add(result.getInt(1));
        }

        // Assert
        assertEquals(resultData.getRowCount(), actualData.size());
        assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0)));
      }
    }
  }

  @Test
  public void testShouldRunSelectQueryWithEmptyVectorsEmbedded() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_WITH_EMPTY_SQL_CMD)) {
      long rowCount = 0;
      while (resultSet.next()) {
        ++rowCount;
      }
      assertEquals(2, rowCount);
    }
  }

  @Test
  public void testResultSetAppMetadata() throws Exception {
    try (Statement statement = connection.createStatement();
        ResultSet resultSet =
            statement.executeQuery(CoreMockedSqlProducers.LEGACY_REGULAR_SQL_CMD)) {
      assertArrayEquals(
          ((ArrowFlightJdbcFlightStreamResultSet) resultSet).getAppMetadata(),
          "foo".getBytes(StandardCharsets.UTF_8));
    }
  }
}
