|
99 | 99 | import org.slf4j.Logger; |
100 | 100 | import org.slf4j.LoggerFactory; |
101 | 101 |
|
| 102 | +import static io.confluent.connect.s3.S3SinkConnectorConfig.PARTITIONER_MAX_OPEN_FILES_CONFIG; |
102 | 103 | import static io.confluent.connect.s3.S3SinkConnectorConfig.SCHEMA_PARTITION_AFFIX_TYPE_CONFIG; |
103 | 104 | import static io.confluent.connect.s3.util.Utils.sinkRecordToLoggableString; |
104 | 105 | import static io.confluent.connect.storage.StorageSinkConnectorConfig.FLUSH_SIZE_CONFIG; |
@@ -658,6 +659,125 @@ public void testWriteRecordTimeBasedPartitionRecordTimestampHoursOutOfOrder() th |
658 | 659 | verify(expectedFiles, 2, schema, records); |
659 | 660 | } |
660 | 661 |
|
| 662 | + @Test |
| 663 | + public void testWriteRecordFieldBasedPartitionAndRotateOnMaxSize() throws Exception { |
| 664 | + localProps.put(PARTITIONER_MAX_OPEN_FILES_CONFIG, "5"); |
| 665 | + localProps.put(FLUSH_SIZE_CONFIG, "10"); |
| 666 | + setUp(); |
| 667 | + |
| 668 | + // Define the partitioner |
| 669 | + Partitioner<?> partitioner = new FieldPartitioner<>(); |
| 670 | + partitioner.configure(parsedConfig); |
| 671 | + |
| 672 | + TopicPartitionWriter topicPartitionWriter = new TopicPartitionWriter( |
| 673 | + TOPIC_PARTITION, storage, writerProvider, partitioner, connectorConfig, context, null); |
| 674 | + |
| 675 | + String key = "key"; |
| 676 | + Schema schema = createSchema(); |
| 677 | + List<Struct> records = createRecordBatches(schema, 6, 1); |
| 678 | + |
| 679 | + Collection<SinkRecord> sinkRecords = createSinkRecords(records, key, schema); |
| 680 | + |
| 681 | + for (SinkRecord record : sinkRecords) { |
| 682 | + topicPartitionWriter.buffer(record); |
| 683 | + } |
| 684 | + |
| 685 | + // Test actual write |
| 686 | + topicPartitionWriter.write(); |
| 687 | + topicPartitionWriter.close(); |
| 688 | + |
| 689 | + @SuppressWarnings("unchecked") |
| 690 | + List<String> partitionFields = (List<String>) parsedConfig.get(PARTITION_FIELD_NAME_CONFIG); |
| 691 | + String partitionField = partitionFields.get(0); |
| 692 | + |
| 693 | + List<String> expectedFiles = new ArrayList<>(); |
| 694 | + for (int i = 0; i < 5; i++) { |
| 695 | + String dirPrefix = partitioner.generatePartitionedPath(TOPIC, partitionField + "=" + String.valueOf(16 + i)); |
| 696 | + expectedFiles.add(FileUtils.fileKeyToCommit(topicsDir, dirPrefix, TOPIC_PARTITION, i, extension, ZERO_PAD_FMT)); |
| 697 | + } |
| 698 | + verify(expectedFiles, 1, schema, records); |
| 699 | + } |
| 700 | + |
| 701 | + @Test |
| 702 | + public void testWriteRecordFieldBasedPartitionAndRotateOnMaxSizeWithMultipleFields() throws Exception { |
| 703 | + localProps.put(PARTITIONER_MAX_OPEN_FILES_CONFIG, "5"); |
| 704 | + localProps.put(FLUSH_SIZE_CONFIG, "10"); |
| 705 | + setUp(); |
| 706 | + |
| 707 | + // Define the partitioner |
| 708 | + Partitioner<?> partitioner = new FieldPartitioner<>(); |
| 709 | + parsedConfig.put(PARTITION_FIELD_NAME_CONFIG, Arrays.asList("int", "boolean")); |
| 710 | + partitioner.configure(parsedConfig); |
| 711 | + |
| 712 | + TopicPartitionWriter topicPartitionWriter = new TopicPartitionWriter( |
| 713 | + TOPIC_PARTITION, storage, writerProvider, partitioner, connectorConfig, context, null); |
| 714 | + |
| 715 | + String key = "key"; |
| 716 | + Schema schema = createSchema(); |
| 717 | + List<Struct> records = createRecordBatches(schema, 6, 1); |
| 718 | + |
| 719 | + Collection<SinkRecord> sinkRecords = createSinkRecords(records, key, schema); |
| 720 | + |
| 721 | + for (SinkRecord record : sinkRecords) { |
| 722 | + topicPartitionWriter.buffer(record); |
| 723 | + } |
| 724 | + |
| 725 | + // Test actual write |
| 726 | + topicPartitionWriter.write(); |
| 727 | + topicPartitionWriter.close(); |
| 728 | + |
| 729 | + @SuppressWarnings("unchecked") |
| 730 | + List<String> partitionFields = (List<String>) parsedConfig.get(PARTITION_FIELD_NAME_CONFIG); |
| 731 | + String partitionField = partitionFields.get(0); |
| 732 | + |
| 733 | + List<String> expectedFiles = new ArrayList<>(); |
| 734 | + for (int i = 0; i < 5; i++) { |
| 735 | + String dirPrefix = partitioner.generatePartitionedPath(TOPIC, partitionField + "=" + String.valueOf(16 + i)); |
| 736 | + dirPrefix = partitioner.generatePartitionedPath(dirPrefix, "boolean=true"); |
| 737 | + expectedFiles.add(FileUtils.fileKeyToCommit(topicsDir, dirPrefix, TOPIC_PARTITION, i, extension, ZERO_PAD_FMT)); |
| 738 | + } |
| 739 | + verify(expectedFiles, 1, schema, records); |
| 740 | + } |
| 741 | + |
| 742 | + @Test |
| 743 | + public void testWriteRecordFieldBasedPartitionAndRotateOnFlushSizeWhenMaxFilesIsUnset() throws Exception { |
| 744 | + localProps.put(PARTITIONER_MAX_OPEN_FILES_CONFIG, "-1"); |
| 745 | + localProps.put(FLUSH_SIZE_CONFIG, "10"); |
| 746 | + setUp(); |
| 747 | + |
| 748 | + // Define the partitioner |
| 749 | + Partitioner<?> partitioner = new FieldPartitioner<>(); |
| 750 | + partitioner.configure(parsedConfig); |
| 751 | + |
| 752 | + TopicPartitionWriter topicPartitionWriter = new TopicPartitionWriter( |
| 753 | + TOPIC_PARTITION, storage, writerProvider, partitioner, connectorConfig, context, null); |
| 754 | + |
| 755 | + String key = "key"; |
| 756 | + Schema schema = createSchema(); |
| 757 | + List<Struct> records = createRecordBatches(schema, 1, 20); |
| 758 | + |
| 759 | + Collection<SinkRecord> sinkRecords = createSinkRecords(records, key, schema); |
| 760 | + |
| 761 | + for (SinkRecord record : sinkRecords) { |
| 762 | + topicPartitionWriter.buffer(record); |
| 763 | + } |
| 764 | + |
| 765 | + // Test actual write |
| 766 | + topicPartitionWriter.write(); |
| 767 | + topicPartitionWriter.close(); |
| 768 | + |
| 769 | + @SuppressWarnings("unchecked") |
| 770 | + List<String> partitionFields = (List<String>) parsedConfig.get(PARTITION_FIELD_NAME_CONFIG); |
| 771 | + String partitionField = partitionFields.get(0); |
| 772 | + |
| 773 | + List<String> expectedFiles = new ArrayList<>(); |
| 774 | + String dirPrefix = partitioner.generatePartitionedPath(TOPIC, partitionField + "=" + String.valueOf(16)); |
| 775 | + expectedFiles.add(FileUtils.fileKeyToCommit(topicsDir, dirPrefix, TOPIC_PARTITION, 0, extension, ZERO_PAD_FMT)); |
| 776 | + expectedFiles.add(FileUtils.fileKeyToCommit(topicsDir, dirPrefix, TOPIC_PARTITION, 10, extension, ZERO_PAD_FMT)); |
| 777 | + |
| 778 | + verify(expectedFiles, 10, schema, records); |
| 779 | + } |
| 780 | + |
661 | 781 | @Test |
662 | 782 | public void testWriteRecordTimeBasedPartitionRecordTimestampHoursOutOfOrderAndRotateOnPartitionChange() throws Exception { |
663 | 783 | localProps.put(FLUSH_SIZE_CONFIG, "1000"); |
|
0 commit comments