|
25 | 25 | import static org.hamcrest.CoreMatchers.instanceOf; |
26 | 26 | import static org.hamcrest.CoreMatchers.is; |
27 | 27 | import static org.hamcrest.MatcherAssert.assertThat; |
28 | | -import static org.junit.jupiter.api.Assertions.assertArrayEquals; |
29 | | -import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; |
30 | | -import static org.junit.jupiter.api.Assertions.assertEquals; |
31 | | -import static org.junit.jupiter.api.Assertions.assertThrows; |
32 | | -import static org.junit.jupiter.api.Assertions.assertTrue; |
33 | | -import static org.junit.jupiter.api.Assertions.fail; |
| 28 | +import static org.junit.jupiter.api.Assertions.*; |
34 | 29 |
|
35 | 30 | import com.google.common.collect.ImmutableSet; |
36 | 31 | import java.nio.charset.StandardCharsets; |
@@ -645,6 +640,139 @@ public void testFallbackSecondFlightServer() throws Exception { |
645 | 640 | } |
646 | 641 | } |
647 | 642 |
|
| 643 | + @Test |
| 644 | + public void testFallbackUnresolvableFlightServer() throws Exception { |
| 645 | + final Schema schema = |
| 646 | + new Schema( |
| 647 | + Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType()))); |
| 648 | + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); |
| 649 | + VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) { |
| 650 | + resultData.setRowCount(1); |
| 651 | + ((IntVector) resultData.getVector(0)).set(0, 1); |
| 652 | + |
| 653 | + try (final FallbackFlightSqlProducer rootProducer = |
| 654 | + new FallbackFlightSqlProducer(resultData); |
| 655 | + FlightServer rootServer = |
| 656 | + FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer) |
| 657 | + .build() |
| 658 | + .start(); |
| 659 | + Connection newConnection = |
| 660 | + DriverManager.getConnection( |
| 661 | + String.format( |
| 662 | + "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false", |
| 663 | + rootServer.getLocation().getUri().getHost(), rootServer.getPort()))) { |
| 664 | + // This first attempt should take a measurable amount of time. |
| 665 | + long start = System.nanoTime(); |
| 666 | + try (Statement newStatement = newConnection.createStatement()) { |
| 667 | + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { |
| 668 | + List<Integer> actualData = new ArrayList<>(); |
| 669 | + while (result.next()) { |
| 670 | + actualData.add(result.getInt(1)); |
| 671 | + } |
| 672 | + |
| 673 | + // Assert |
| 674 | + assertEquals(resultData.getRowCount(), actualData.size()); |
| 675 | + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); |
| 676 | + } |
| 677 | + } |
| 678 | + long attempt1 = System.nanoTime(); |
| 679 | + double elapsedMs = (attempt1 - start) / 1_000_000.; |
| 680 | + assertTrue( |
| 681 | + elapsedMs >= 5000., |
| 682 | + String.format( |
| 683 | + "Expected first attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); |
| 684 | + |
| 685 | + // Once the client cache is implemented (GH-661), this second attempt should take less time, |
| 686 | + // since the failure from before should be cached. |
| 687 | + start = System.nanoTime(); |
| 688 | + try (Statement newStatement = newConnection.createStatement()) { |
| 689 | + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { |
| 690 | + List<Integer> actualData = new ArrayList<>(); |
| 691 | + while (result.next()) { |
| 692 | + actualData.add(result.getInt(1)); |
| 693 | + } |
| 694 | + |
| 695 | + // Assert |
| 696 | + assertEquals(resultData.getRowCount(), actualData.size()); |
| 697 | + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); |
| 698 | + } |
| 699 | + } |
| 700 | + attempt1 = System.nanoTime(); |
| 701 | + elapsedMs = (attempt1 - start) / 1_000_000.; |
| 702 | + // TODO(GH-661): this assertion should be flipped to assertTrue. |
| 703 | + assertFalse( |
| 704 | + elapsedMs < 5000., |
| 705 | + String.format("Expected second attempt to be the same, but %f ms elapsed", elapsedMs)); |
| 706 | + } |
| 707 | + } |
| 708 | + } |
| 709 | + |
| 710 | + @Test |
| 711 | + public void testFallbackUnresolvableFlightServerDisableCache() throws Exception { |
| 712 | + final Schema schema = |
| 713 | + new Schema( |
| 714 | + Collections.singletonList(Field.nullable("int_column", Types.MinorType.INT.getType()))); |
| 715 | + try (BufferAllocator allocator = new RootAllocator(Long.MAX_VALUE); |
| 716 | + VectorSchemaRoot resultData = VectorSchemaRoot.create(schema, allocator)) { |
| 717 | + resultData.setRowCount(1); |
| 718 | + ((IntVector) resultData.getVector(0)).set(0, 1); |
| 719 | + |
| 720 | + try (final FallbackFlightSqlProducer rootProducer = |
| 721 | + new FallbackFlightSqlProducer(resultData); |
| 722 | + FlightServer rootServer = |
| 723 | + FlightServer.builder(allocator, forGrpcInsecure("localhost", 0), rootProducer) |
| 724 | + .build() |
| 725 | + .start(); |
| 726 | + Connection newConnection = |
| 727 | + DriverManager.getConnection( |
| 728 | + String.format( |
| 729 | + "jdbc:arrow-flight-sql://%s:%d/?useEncryption=false&useClientCache=false", |
| 730 | + rootServer.getLocation().getUri().getHost(), rootServer.getPort()))) { |
| 731 | + // This first attempt should take a measurable amount of time. |
| 732 | + long start = System.nanoTime(); |
| 733 | + try (Statement newStatement = newConnection.createStatement()) { |
| 734 | + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { |
| 735 | + List<Integer> actualData = new ArrayList<>(); |
| 736 | + while (result.next()) { |
| 737 | + actualData.add(result.getInt(1)); |
| 738 | + } |
| 739 | + |
| 740 | + // Assert |
| 741 | + assertEquals(resultData.getRowCount(), actualData.size()); |
| 742 | + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); |
| 743 | + } |
| 744 | + } |
| 745 | + long attempt1 = System.nanoTime(); |
| 746 | + double elapsedMs = (attempt1 - start) / 1_000_000.; |
| 747 | + assertTrue( |
| 748 | + elapsedMs >= 5000., |
| 749 | + String.format( |
| 750 | + "Expected first attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); |
| 751 | + |
| 752 | + // This second attempt should take a long time still, since we disabled the cache. |
| 753 | + start = System.nanoTime(); |
| 754 | + try (Statement newStatement = newConnection.createStatement()) { |
| 755 | + try (ResultSet result = newStatement.executeQuery("fallback with unresolvable")) { |
| 756 | + List<Integer> actualData = new ArrayList<>(); |
| 757 | + while (result.next()) { |
| 758 | + actualData.add(result.getInt(1)); |
| 759 | + } |
| 760 | + |
| 761 | + // Assert |
| 762 | + assertEquals(resultData.getRowCount(), actualData.size()); |
| 763 | + assertTrue(actualData.contains(((IntVector) resultData.getVector(0)).get(0))); |
| 764 | + } |
| 765 | + } |
| 766 | + attempt1 = System.nanoTime(); |
| 767 | + elapsedMs = (attempt1 - start) / 1_000_000.; |
| 768 | + assertTrue( |
| 769 | + elapsedMs >= 5000., |
| 770 | + String.format( |
| 771 | + "Expected second attempt to hit the timeout, but only %f ms elapsed", elapsedMs)); |
| 772 | + } |
| 773 | + } |
| 774 | + } |
| 775 | + |
648 | 776 | @Test |
649 | 777 | public void testShouldRunSelectQueryWithEmptyVectorsEmbedded() throws Exception { |
650 | 778 | try (Statement statement = connection.createStatement(); |
|
0 commit comments