You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by tz...@apache.org on 2023/04/06 15:57:47 UTC
[flink-connector-kafka] branch main updated (40cf9994 -> ffa11654)
This is an automated email from the ASF dual-hosted git repository.
tzulitai pushed a change to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
from 40cf9994 [hotfix] Test PRs against Flink 1.17.0 and all nightly builds against 1.17-SNAPSHOT and 1.18-SNAPSHOT
new 85bdc841 [hotfix] Refactor test util method for collecting results of bounded query
new fc31198d [FLINK-31740] [upsert-kafka] Allow setting boundedness options for upsert-kafka
new ffa11654 [hotfix] Refactor MockPartitionOffsetsRetriever as a common test utility
The 3 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
Summary of changes:
.../kafka/table/KafkaConnectorOptionsUtil.java | 2 +-
.../table/UpsertKafkaDynamicTableFactory.java | 19 +-
.../kafka/table/KafkaDynamicTableFactoryTest.java | 72 +-------
.../connectors/kafka/table/KafkaTableITCase.java | 18 +-
.../kafka/table/KafkaTableTestUtils.java | 17 ++
.../table/UpsertKafkaDynamicTableFactoryTest.java | 157 +++++++++++++++-
.../kafka/table/UpsertKafkaTableITCase.java | 199 +++++++++++++++++++++
.../testutils/MockPartitionOffsetsRetriever.java | 102 +++++++++++
8 files changed, 494 insertions(+), 92 deletions(-)
create mode 100644 flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/testutils/MockPartitionOffsetsRetriever.java
[flink-connector-kafka] 01/03: [hotfix] Refactor test util method for collecting results of bounded query
Posted by tz...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tzulitai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
commit 85bdc841e72e8d4ac41c1824139b3e9730a1ed2d
Author: Tzu-Li (Gordon) Tai <tz...@apache.org>
AuthorDate: Wed Apr 5 14:58:32 2023 -0700
[hotfix] Refactor test util method for collecting results of bounded query
---
.../connectors/kafka/table/KafkaTableITCase.java | 18 +++---------------
.../connectors/kafka/table/KafkaTableTestUtils.java | 17 +++++++++++++++++
2 files changed, 20 insertions(+), 15 deletions(-)
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableITCase.java
index 515526f9..2674183f 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableITCase.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableITCase.java
@@ -29,7 +29,6 @@ import org.apache.flink.table.data.RowData;
import org.apache.flink.table.utils.EncodingUtils;
import org.apache.flink.test.util.SuccessException;
import org.apache.flink.types.Row;
-import org.apache.flink.util.CloseableIterator;
import org.apache.kafka.clients.consumer.NoOffsetForPartitionException;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -58,6 +57,7 @@ import java.util.stream.Collectors;
import java.util.stream.IntStream;
import static org.apache.flink.core.testutils.CommonTestUtils.waitUtil;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.collectAllRows;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.collectRows;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.readLines;
import static org.apache.flink.table.api.config.ExecutionConfigOptions.TABLE_EXEC_SOURCE_IDLE_TIMEOUT;
@@ -225,13 +225,7 @@ public class KafkaTableITCase extends KafkaTableTestBase {
// ---------- Consume stream from Kafka -------------------
- List<Row> results = new ArrayList<>();
- try (CloseableIterator<Row> resultsItr =
- tEnv.sqlQuery("SELECT * from kafka").execute().collect()) {
- while (resultsItr.hasNext()) {
- results.add(resultsItr.next());
- }
- }
+ List<Row> results = collectAllRows(tEnv.sqlQuery("SELECT * from kafka"));
assertThat(results)
.containsExactly(Row.of(1, 1102, "behavior 1"), Row.of(2, 1103, "behavior 2"));
@@ -286,13 +280,7 @@ public class KafkaTableITCase extends KafkaTableTestBase {
// ---------- Consume stream from Kafka -------------------
- List<Row> results = new ArrayList<>();
- try (CloseableIterator<Row> resultsItr =
- tEnv.sqlQuery("SELECT * from kafka").execute().collect()) {
- while (resultsItr.hasNext()) {
- results.add(resultsItr.next());
- }
- }
+ List<Row> results = collectAllRows(tEnv.sqlQuery("SELECT * from kafka"));
assertThat(results)
.containsExactly(
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableTestUtils.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableTestUtils.java
index c1aada46..793d8da7 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableTestUtils.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaTableTestUtils.java
@@ -68,6 +68,23 @@ public class KafkaTableTestUtils {
return collectedRows;
}
+ /**
+ * Variant of {@link #collectRows(Table, int)} for bounded queries. This should not run
+ * indefinitely if there is a bounded number of returned rows.
+ */
+ public static List<Row> collectAllRows(Table table) throws Exception {
+ final TableResult result = table.execute();
+
+ final List<Row> collectedRows = new ArrayList<>();
+ try (CloseableIterator<Row> iterator = result.collect()) {
+ while (iterator.hasNext()) {
+ collectedRows.add(iterator.next());
+ }
+ }
+
+ return collectedRows;
+ }
+
public static List<String> readLines(String resource) throws IOException {
final URL url = KafkaChangelogTableITCase.class.getClassLoader().getResource(resource);
assertThat(url).isNotNull();
[flink-connector-kafka] 02/03: [FLINK-31740] [upsert-kafka] Allow setting boundedness options for upsert-kafka
Posted by tz...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tzulitai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
commit fc31198d864ebafd36cf23ce302b6c1b4038b579
Author: Tzu-Li (Gordon) Tai <tz...@apache.org>
AuthorDate: Wed Apr 5 14:49:12 2023 -0700
[FLINK-31740] [upsert-kafka] Allow setting boundedness options for upsert-kafka
This closes #22.
---
.../kafka/table/KafkaConnectorOptionsUtil.java | 2 +-
.../table/UpsertKafkaDynamicTableFactory.java | 19 +-
.../table/UpsertKafkaDynamicTableFactoryTest.java | 227 ++++++++++++++++++++-
.../kafka/table/UpsertKafkaTableITCase.java | 199 ++++++++++++++++++
4 files changed, 441 insertions(+), 6 deletions(-)
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/KafkaConnectorOptionsUtil.java b/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/KafkaConnectorOptionsUtil.java
index ef70644e..d6390e27 100644
--- a/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/KafkaConnectorOptionsUtil.java
+++ b/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/KafkaConnectorOptionsUtil.java
@@ -189,7 +189,7 @@ class KafkaConnectorOptionsUtil {
});
}
- private static void validateScanBoundedMode(ReadableConfig tableOptions) {
+ static void validateScanBoundedMode(ReadableConfig tableOptions) {
tableOptions
.getOptional(SCAN_BOUNDED_MODE)
.ifPresent(
diff --git a/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactory.java b/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactory.java
index 254e1bf9..b9f2ea71 100644
--- a/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactory.java
+++ b/flink-connector-kafka/src/main/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactory.java
@@ -25,8 +25,8 @@ import org.apache.flink.configuration.ConfigOption;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.ReadableConfig;
import org.apache.flink.connector.base.DeliveryGuarantee;
-import org.apache.flink.streaming.connectors.kafka.config.BoundedMode;
import org.apache.flink.streaming.connectors.kafka.config.StartupMode;
+import org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.BoundedOptions;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.catalog.ResolvedCatalogTable;
import org.apache.flink.table.catalog.ResolvedSchema;
@@ -57,6 +57,9 @@ import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOp
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.KEY_FIELDS_PREFIX;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.KEY_FORMAT;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.PROPS_BOOTSTRAP_SERVERS;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SCAN_BOUNDED_MODE;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SCAN_BOUNDED_SPECIFIC_OFFSETS;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SCAN_BOUNDED_TIMESTAMP_MILLIS;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SINK_BUFFER_FLUSH_INTERVAL;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SINK_BUFFER_FLUSH_MAX_ROWS;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.SINK_PARALLELISM;
@@ -68,9 +71,11 @@ import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOp
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.autoCompleteSchemaRegistrySubject;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.createKeyFormatProjection;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.createValueFormatProjection;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.getBoundedOptions;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.getKafkaProperties;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.getSourceTopicPattern;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.getSourceTopics;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.validateScanBoundedMode;
/** Upsert-Kafka factory. */
public class UpsertKafkaDynamicTableFactory
@@ -101,6 +106,9 @@ public class UpsertKafkaDynamicTableFactory
options.add(SINK_PARALLELISM);
options.add(SINK_BUFFER_FLUSH_INTERVAL);
options.add(SINK_BUFFER_FLUSH_MAX_ROWS);
+ options.add(SCAN_BOUNDED_MODE);
+ options.add(SCAN_BOUNDED_SPECIFIC_OFFSETS);
+ options.add(SCAN_BOUNDED_TIMESTAMP_MILLIS);
return options;
}
@@ -129,6 +137,8 @@ public class UpsertKafkaDynamicTableFactory
// always use earliest to keep data integrity
StartupMode earliest = StartupMode.EARLIEST;
+ final BoundedOptions boundedOptions = getBoundedOptions(tableOptions);
+
return new KafkaDynamicSource(
context.getPhysicalRowDataType(),
keyDecodingFormat,
@@ -142,9 +152,9 @@ public class UpsertKafkaDynamicTableFactory
earliest,
Collections.emptyMap(),
0,
- BoundedMode.UNBOUNDED,
- Collections.emptyMap(),
- 0,
+ boundedOptions.boundedMode,
+ boundedOptions.specificOffsets,
+ boundedOptions.boundedTimestampMillis,
true,
context.getObjectIdentifier().asSummaryString());
}
@@ -228,6 +238,7 @@ public class UpsertKafkaDynamicTableFactory
Format valueFormat,
int[] primaryKeyIndexes) {
validateTopic(tableOptions);
+ validateScanBoundedMode(tableOptions);
validateFormat(keyFormat, valueFormat, tableOptions);
validatePKConstraints(primaryKeyIndexes);
}
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
index 5caaaa0a..959c44cf 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
@@ -21,11 +21,14 @@ package org.apache.flink.streaming.connectors.kafka.table;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.serialization.SerializationSchema;
import org.apache.flink.api.connector.sink2.Sink;
+import org.apache.flink.api.connector.source.Boundedness;
import org.apache.flink.api.dag.Transformation;
import org.apache.flink.connector.base.DeliveryGuarantee;
import org.apache.flink.connector.kafka.sink.KafkaSink;
import org.apache.flink.connector.kafka.source.KafkaSource;
+import org.apache.flink.connector.kafka.source.KafkaSourceTestUtils;
import org.apache.flink.connector.kafka.source.enumerator.KafkaSourceEnumState;
+import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
import org.apache.flink.connector.kafka.source.split.KafkaPartitionSplit;
import org.apache.flink.formats.avro.AvroRowDataSerializationSchema;
import org.apache.flink.formats.avro.RowDataToAvroConverters;
@@ -64,23 +67,29 @@ import org.apache.flink.table.types.logical.RowType;
import org.apache.flink.table.types.logical.VarCharType;
import org.apache.flink.util.TestLogger;
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
+import org.apache.kafka.common.TopicPartition;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.ExpectedException;
import java.util.Arrays;
+import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.function.Consumer;
+import java.util.function.Function;
import static org.apache.flink.core.testutils.FlinkMatchers.containsCause;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.ScanBoundedMode;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptionsUtil.AVRO_CONFLUENT;
import static org.apache.flink.table.factories.utils.FactoryMocks.createTableSink;
import static org.apache.flink.table.factories.utils.FactoryMocks.createTableSource;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
/** Test for {@link UpsertKafkaDynamicTableFactory}. */
public class UpsertKafkaDynamicTableFactoryTest extends TestLogger {
@@ -390,6 +399,134 @@ public class UpsertKafkaDynamicTableFactoryTest extends TestLogger {
RowDataToAvroConverters.createConverter(rowType));
}
+ // --------------------------------------------------------------------------------------------
+ // Bounded end-offset tests
+ // --------------------------------------------------------------------------------------------
+
+ @Test
+ public void testBoundedSpecificOffsetsValidate() {
+ final Map<String, String> options = getFullSourceOptions();
+ options.put(
+ KafkaConnectorOptions.SCAN_BOUNDED_MODE.key(),
+ ScanBoundedMode.SPECIFIC_OFFSETS.toString());
+
+ assertThatThrownBy(() -> createTableSource(SOURCE_SCHEMA, options))
+ .isInstanceOf(ValidationException.class)
+ .cause()
+ .hasMessageContaining(
+ "'scan.bounded.specific-offsets' is required in 'specific-offsets' bounded mode but missing.");
+ }
+
+ @Test
+ public void testBoundedSpecificOffsets() {
+ testBoundedOffsets(
+ ScanBoundedMode.SPECIFIC_OFFSETS,
+ options -> {
+ options.put("scan.bounded.specific-offsets", "partition:0,offset:2");
+ },
+ source -> {
+ assertThat(source.getBoundedness()).isEqualTo(Boundedness.BOUNDED);
+ OffsetsInitializer offsetsInitializer =
+ KafkaSourceTestUtils.getStoppingOffsetsInitializer(source);
+ TopicPartition partition = new TopicPartition(SOURCE_TOPIC, 0);
+ Map<TopicPartition, Long> partitionOffsets =
+ offsetsInitializer.getPartitionOffsets(
+ Collections.singletonList(partition),
+ MockPartitionOffsetsRetriever.noInteractions());
+ assertThat(partitionOffsets)
+ .containsOnlyKeys(partition)
+ .containsEntry(partition, 2L);
+ });
+ }
+
+ @Test
+ public void testBoundedLatestOffset() {
+ testBoundedOffsets(
+ ScanBoundedMode.LATEST_OFFSET,
+ options -> {},
+ source -> {
+ assertThat(source.getBoundedness()).isEqualTo(Boundedness.BOUNDED);
+ OffsetsInitializer offsetsInitializer =
+ KafkaSourceTestUtils.getStoppingOffsetsInitializer(source);
+ TopicPartition partition = new TopicPartition(SOURCE_TOPIC, 0);
+ Map<TopicPartition, Long> partitionOffsets =
+ offsetsInitializer.getPartitionOffsets(
+ Collections.singletonList(partition),
+ MockPartitionOffsetsRetriever.noInteractions());
+ assertThat(partitionOffsets)
+ .containsOnlyKeys(partition)
+ .containsEntry(partition, KafkaPartitionSplit.LATEST_OFFSET);
+ });
+ }
+
+ @Test
+ public void testBoundedGroupOffsets() {
+ testBoundedOffsets(
+ ScanBoundedMode.GROUP_OFFSETS,
+ options -> {
+ options.put("properties.group.id", "dummy");
+ },
+ source -> {
+ assertThat(source.getBoundedness()).isEqualTo(Boundedness.BOUNDED);
+ OffsetsInitializer offsetsInitializer =
+ KafkaSourceTestUtils.getStoppingOffsetsInitializer(source);
+ TopicPartition partition = new TopicPartition(SOURCE_TOPIC, 0);
+ Map<TopicPartition, Long> partitionOffsets =
+ offsetsInitializer.getPartitionOffsets(
+ Collections.singletonList(partition),
+ MockPartitionOffsetsRetriever.noInteractions());
+ assertThat(partitionOffsets)
+ .containsOnlyKeys(partition)
+ .containsEntry(partition, KafkaPartitionSplit.COMMITTED_OFFSET);
+ });
+ }
+
+ @Test
+ public void testBoundedTimestamp() {
+ testBoundedOffsets(
+ ScanBoundedMode.TIMESTAMP,
+ options -> {
+ options.put("scan.bounded.timestamp-millis", "1");
+ },
+ source -> {
+ assertThat(source.getBoundedness()).isEqualTo(Boundedness.BOUNDED);
+ OffsetsInitializer offsetsInitializer =
+ KafkaSourceTestUtils.getStoppingOffsetsInitializer(source);
+ TopicPartition partition = new TopicPartition(SOURCE_TOPIC, 0);
+ long offsetForTimestamp = 123L;
+ Map<TopicPartition, Long> partitionOffsets =
+ offsetsInitializer.getPartitionOffsets(
+ Collections.singletonList(partition),
+ MockPartitionOffsetsRetriever.timestampAndEnd(
+ partitions -> {
+ assertThat(partitions)
+ .containsOnlyKeys(partition)
+ .containsEntry(partition, 1L);
+ Map<TopicPartition, OffsetAndTimestamp> result =
+ new HashMap<>();
+ result.put(
+ partition,
+ new OffsetAndTimestamp(
+ offsetForTimestamp, 1L));
+ return result;
+ },
+ partitions -> {
+ Map<TopicPartition, Long> result = new HashMap<>();
+ result.put(
+ partition,
+ // the end offset is bigger than given by
+ // timestamp
+ // to make sure the one for timestamp is
+ // used
+ offsetForTimestamp + 1000L);
+ return result;
+ }));
+ assertThat(partitionOffsets)
+ .containsOnlyKeys(partition)
+ .containsEntry(partition, offsetForTimestamp);
+ });
+ }
+
// --------------------------------------------------------------------------------------------
// Negative tests
// --------------------------------------------------------------------------------------------
@@ -647,7 +784,7 @@ public class UpsertKafkaDynamicTableFactoryTest extends TestLogger {
null);
}
- private void assertKafkaSource(ScanTableSource.ScanRuntimeProvider provider) {
+ private KafkaSource<?> assertKafkaSource(ScanTableSource.ScanRuntimeProvider provider) {
assertThat(provider).isInstanceOf(DataStreamScanProvider.class);
final DataStreamScanProvider dataStreamScanProvider = (DataStreamScanProvider) provider;
final Transformation<RowData> transformation =
@@ -662,5 +799,93 @@ public class UpsertKafkaDynamicTableFactoryTest extends TestLogger {
(SourceTransformation<RowData, KafkaPartitionSplit, KafkaSourceEnumState>)
transformation;
assertThat(sourceTransformation.getSource()).isInstanceOf(KafkaSource.class);
+ return (KafkaSource<?>) sourceTransformation.getSource();
+ }
+
+ private void testBoundedOffsets(
+ ScanBoundedMode boundedMode,
+ Consumer<Map<String, String>> optionsConfig,
+ Consumer<KafkaSource<?>> validator) {
+ final Map<String, String> options = getFullSourceOptions();
+ options.put(KafkaConnectorOptions.SCAN_BOUNDED_MODE.key(), boundedMode.toString());
+ optionsConfig.accept(options);
+
+ final DynamicTableSource tableSource = createTableSource(SOURCE_SCHEMA, options);
+ assertThat(tableSource).isInstanceOf(KafkaDynamicSource.class);
+ ScanTableSource.ScanRuntimeProvider provider =
+ ((KafkaDynamicSource) tableSource)
+ .getScanRuntimeProvider(ScanRuntimeProviderContext.INSTANCE);
+ assertThat(provider).isInstanceOf(DataStreamScanProvider.class);
+ final KafkaSource<?> kafkaSource = assertKafkaSource(provider);
+ validator.accept(kafkaSource);
+ }
+
+ private interface OffsetsRetriever
+ extends Function<Collection<TopicPartition>, Map<TopicPartition, Long>> {}
+
+ private interface TimestampOffsetsRetriever
+ extends Function<Map<TopicPartition, Long>, Map<TopicPartition, OffsetAndTimestamp>> {}
+
+ private static final class MockPartitionOffsetsRetriever
+ implements OffsetsInitializer.PartitionOffsetsRetriever {
+
+ public static final OffsetsRetriever UNSUPPORTED_RETRIEVAL =
+ partitions -> {
+ throw new UnsupportedOperationException(
+ "The method was not supposed to be called");
+ };
+ private final OffsetsRetriever committedOffsets;
+ private final OffsetsRetriever endOffsets;
+ private final OffsetsRetriever beginningOffsets;
+ private final TimestampOffsetsRetriever offsetsForTimes;
+
+ static MockPartitionOffsetsRetriever noInteractions() {
+ return new MockPartitionOffsetsRetriever(
+ UNSUPPORTED_RETRIEVAL,
+ UNSUPPORTED_RETRIEVAL,
+ UNSUPPORTED_RETRIEVAL,
+ partitions -> {
+ throw new UnsupportedOperationException(
+ "The method was not supposed to be called");
+ });
+ }
+
+ static MockPartitionOffsetsRetriever timestampAndEnd(
+ TimestampOffsetsRetriever retriever, OffsetsRetriever endOffsets) {
+ return new MockPartitionOffsetsRetriever(
+ UNSUPPORTED_RETRIEVAL, endOffsets, UNSUPPORTED_RETRIEVAL, retriever);
+ }
+
+ private MockPartitionOffsetsRetriever(
+ OffsetsRetriever committedOffsets,
+ OffsetsRetriever endOffsets,
+ OffsetsRetriever beginningOffsets,
+ TimestampOffsetsRetriever offsetsForTimes) {
+ this.committedOffsets = committedOffsets;
+ this.endOffsets = endOffsets;
+ this.beginningOffsets = beginningOffsets;
+ this.offsetsForTimes = offsetsForTimes;
+ }
+
+ @Override
+ public Map<TopicPartition, Long> committedOffsets(Collection<TopicPartition> partitions) {
+ return committedOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions) {
+ return endOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions) {
+ return beginningOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
+ Map<TopicPartition, Long> timestampsToSearch) {
+ return offsetsForTimes.apply(timestampsToSearch);
+ }
}
}
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaTableITCase.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaTableITCase.java
index 109f4402..23f638f6 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaTableITCase.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaTableITCase.java
@@ -33,17 +33,20 @@ import org.junit.runners.Parameterized;
import java.time.Duration;
import java.time.LocalDateTime;
+import java.time.ZoneId;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.UUID;
import java.util.stream.Collectors;
import static org.apache.flink.api.common.typeinfo.Types.INT;
import static org.apache.flink.api.common.typeinfo.Types.LOCAL_DATE_TIME;
import static org.apache.flink.api.common.typeinfo.Types.ROW_NAMED;
import static org.apache.flink.api.common.typeinfo.Types.STRING;
+import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.collectAllRows;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.collectRows;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.comparedWithKeyAndOrder;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaTableTestUtils.waitingExpectedResults;
@@ -388,6 +391,202 @@ public class UpsertKafkaTableITCase extends KafkaTableTestBase {
deleteTestTopic(topic);
}
+ @Test
+ public void testUpsertKafkaSourceSinkWithBoundedSpecificOffsets() throws Exception {
+ final String topic = "bounded_upsert_" + format + "_" + UUID.randomUUID();
+ createTestTopic(topic, 1, 1);
+
+ // ---------- Produce an event time stream into Kafka -------------------
+ final String bootstraps = getBootstrapServers();
+
+ // table with upsert-kafka connector, bounded mode up to offset=2
+ final String createTableSql =
+ String.format(
+ "CREATE TABLE upsert_kafka (\n"
+ + " `user_id` BIGINT,\n"
+ + " `event_id` BIGINT,\n"
+ + " `payload` STRING,\n"
+ + " PRIMARY KEY (event_id, user_id) NOT ENFORCED"
+ + ") WITH (\n"
+ + " 'connector' = 'upsert-kafka',\n"
+ + " 'topic' = '%s',\n"
+ + " 'properties.bootstrap.servers' = '%s',\n"
+ + " 'key.format' = '%s',\n"
+ + " 'value.format' = '%s',\n"
+ + " 'value.fields-include' = 'ALL',\n"
+ + " 'scan.bounded.mode' = 'specific-offsets',\n"
+ + " 'scan.bounded.specific-offsets' = 'partition:0,offset:2'"
+ + ")",
+ topic, bootstraps, format, format);
+ tEnv.executeSql(createTableSql);
+
+ // insert multiple values to have more records past offset=2
+ final String insertValuesSql =
+ "INSERT INTO upsert_kafka\n"
+ + "VALUES\n"
+ + " (1, 100, 'payload 1'),\n"
+ + " (1, 100, 'payload 1-new'),\n"
+ + " (2, 101, 'payload 2'),\n"
+ + " (3, 102, 'payload 3')";
+ tEnv.executeSql(insertValuesSql).await();
+
+ // results should only have records up to offset=2
+ final List<Row> results = collectAllRows(tEnv.sqlQuery("SELECT * from upsert_kafka"));
+ final List<Row> expected =
+ Arrays.asList(
+ changelogRow("+I", 1L, 100L, "payload 1"),
+ changelogRow("-U", 1L, 100L, "payload 1"),
+ changelogRow("+U", 1L, 100L, "payload 1-new"));
+ assertThat(results).satisfies(matching(deepEqualTo(expected, true)));
+
+ // ------------- cleanup -------------------
+
+ deleteTestTopic(topic);
+ }
+
+ @Test
+ public void testUpsertKafkaSourceSinkWithBoundedTimestamp() throws Exception {
+ final String topic = "bounded_upsert_" + format + "_" + UUID.randomUUID();
+ createTestTopic(topic, 1, 1);
+
+ // ---------- Produce an event time stream into Kafka -------------------
+ final String bootstraps = getBootstrapServers();
+
+ // table with upsert-kafka connector, bounded mode up to timestamp 2023-03-10T14:00:00.000
+ final String createTableSql =
+ String.format(
+ "CREATE TABLE upsert_kafka (\n"
+ + " `user_id` BIGINT,\n"
+ + " `timestamp` TIMESTAMP(3) METADATA,\n"
+ + " `event_id` BIGINT,\n"
+ + " `payload` STRING,\n"
+ + " PRIMARY KEY (event_id, user_id) NOT ENFORCED"
+ + ") WITH (\n"
+ + " 'connector' = 'upsert-kafka',\n"
+ + " 'topic' = '%s',\n"
+ + " 'properties.bootstrap.servers' = '%s',\n"
+ + " 'key.format' = '%s',\n"
+ + " 'value.format' = '%s',\n"
+ + " 'value.fields-include' = 'ALL',\n"
+ + " 'scan.bounded.mode' = 'timestamp',\n"
+ + " 'scan.bounded.timestamp-millis' = '%d'"
+ + ")",
+ topic,
+ bootstraps,
+ format,
+ format,
+ LocalDateTime.parse("2023-03-10T14:00:00.000")
+ .atZone(ZoneId.systemDefault())
+ .toInstant()
+ .toEpochMilli());
+ tEnv.executeSql(createTableSql);
+
+ // insert multiple values with timestamp starting from 2023-03-08 up to 2023-03-11
+ final String insertValuesSql =
+ "INSERT INTO upsert_kafka\n"
+ + "VALUES\n"
+ + " (1, TIMESTAMP '2023-03-08 08:10:10.666', 100, 'payload 1'),\n"
+ + " (2, TIMESTAMP '2023-03-09 13:12:11.123', 101, 'payload 2'),\n"
+ + " (1, TIMESTAMP '2023-03-10 12:09:50.321', 100, 'payload 1-new'),\n"
+ + " (2, TIMESTAMP '2023-03-11 17:15:13.457', 101, 'payload 2-new')";
+ tEnv.executeSql(insertValuesSql).await();
+
+ // results should only have records up to timestamp 2023-03-10T14:00:00.000
+ final List<Row> results = collectAllRows(tEnv.sqlQuery("SELECT * from upsert_kafka"));
+ final List<Row> expected =
+ Arrays.asList(
+ changelogRow(
+ "+I",
+ 1L,
+ LocalDateTime.parse("2023-03-08T08:10:10.666"),
+ 100L,
+ "payload 1"),
+ changelogRow(
+ "+I",
+ 2L,
+ LocalDateTime.parse("2023-03-09T13:12:11.123"),
+ 101L,
+ "payload 2"),
+ changelogRow(
+ "-U",
+ 1L,
+ LocalDateTime.parse("2023-03-08T08:10:10.666"),
+ 100L,
+ "payload 1"),
+ changelogRow(
+ "+U",
+ 1L,
+ LocalDateTime.parse("2023-03-10T12:09:50.321"),
+ 100L,
+ "payload 1-new"));
+ assertThat(results).satisfies(matching(deepEqualTo(expected, true)));
+
+ // ------------- cleanup -------------------
+
+ deleteTestTopic(topic);
+ }
+
+ /**
+ * Tests that setting bounded end offset that is before the earliest offset results in 0
+ * results.
+ */
+ @Test
+ public void testUpsertKafkaSourceSinkWithZeroLengthBoundedness() throws Exception {
+ final String topic = "bounded_upsert_" + format + "_" + UUID.randomUUID();
+ createTestTopic(topic, 1, 1);
+
+ // ---------- Produce an event time stream into Kafka -------------------
+ final String bootstraps = getBootstrapServers();
+
+ // table with upsert-kafka connector, bounded mode up to timestamp 2023-03-10T14:00:00.000
+ final String createTableSql =
+ String.format(
+ "CREATE TABLE upsert_kafka (\n"
+ + " `user_id` BIGINT,\n"
+ + " `timestamp` TIMESTAMP(3) METADATA,\n"
+ + " `event_id` BIGINT,\n"
+ + " `payload` STRING,\n"
+ + " PRIMARY KEY (event_id, user_id) NOT ENFORCED"
+ + ") WITH (\n"
+ + " 'connector' = 'upsert-kafka',\n"
+ + " 'topic' = '%s',\n"
+ + " 'properties.bootstrap.servers' = '%s',\n"
+ + " 'key.format' = '%s',\n"
+ + " 'value.format' = '%s',\n"
+ + " 'value.fields-include' = 'ALL',\n"
+ + " 'scan.bounded.mode' = 'timestamp',\n"
+ + " 'scan.bounded.timestamp-millis' = '%d'"
+ + ")",
+ topic,
+ bootstraps,
+ format,
+ format,
+ LocalDateTime.parse("2023-03-10T14:00:00.000")
+ .atZone(ZoneId.systemDefault())
+ .toInstant()
+ .toEpochMilli());
+ tEnv.executeSql(createTableSql);
+
+ // insert multiple values with timestamp starting from 2023-03-11 (which is past the bounded
+ // end timestamp)
+ final String insertValuesSql =
+ "INSERT INTO upsert_kafka\n"
+ + "VALUES\n"
+ + " (1, TIMESTAMP '2023-03-11 08:10:10.666', 100, 'payload 1'),\n"
+ + " (2, TIMESTAMP '2023-03-12 13:12:11.123', 101, 'payload 2'),\n"
+ + " (1, TIMESTAMP '2023-03-13 12:09:50.321', 100, 'payload 1-new'),\n"
+ + " (2, TIMESTAMP '2023-03-14 17:15:13.457', 101, 'payload 2-new')";
+ tEnv.executeSql(insertValuesSql).await();
+
+ // results should be empty
+ final List<Row> results = collectAllRows(tEnv.sqlQuery("SELECT * from upsert_kafka"));
+ assertThat(results).satisfies(matching(deepEqualTo(Collections.emptyList(), true)));
+
+ // ------------- cleanup -------------------
+
+ deleteTestTopic(topic);
+ }
+
private void wordCountToUpsertKafka(String wordCountTable) throws Exception {
String bootstraps = getBootstrapServers();
[flink-connector-kafka] 03/03: [hotfix] Refactor MockPartitionOffsetsRetriever as a common test utility
Posted by tz...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
tzulitai pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-connector-kafka.git
commit ffa11654b67885cf7beb823182e640b863162d72
Author: Tzu-Li (Gordon) Tai <tz...@apache.org>
AuthorDate: Wed Apr 5 15:08:39 2023 -0700
[hotfix] Refactor MockPartitionOffsetsRetriever as a common test utility
---
.../kafka/table/KafkaDynamicTableFactoryTest.java | 72 +--------------
.../table/UpsertKafkaDynamicTableFactoryTest.java | 72 +--------------
.../testutils/MockPartitionOffsetsRetriever.java | 102 +++++++++++++++++++++
3 files changed, 104 insertions(+), 142 deletions(-)
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaDynamicTableFactoryTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaDynamicTableFactoryTest.java
index a0b74a5b..7ab05035 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaDynamicTableFactoryTest.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/KafkaDynamicTableFactoryTest.java
@@ -46,6 +46,7 @@ import org.apache.flink.streaming.connectors.kafka.internals.KafkaTopicPartition
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkFixedPartitioner;
import org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner;
import org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.ScanStartupMode;
+import org.apache.flink.streaming.connectors.kafka.testutils.MockPartitionOffsetsRetriever;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.ValidationException;
import org.apache.flink.table.catalog.Column;
@@ -87,7 +88,6 @@ import org.junit.jupiter.params.provider.ValueSource;
import javax.annotation.Nullable;
import java.util.Arrays;
-import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
@@ -97,7 +97,6 @@ import java.util.Optional;
import java.util.Properties;
import java.util.Set;
import java.util.function.Consumer;
-import java.util.function.Function;
import java.util.regex.Pattern;
import static org.apache.flink.core.testutils.FlinkAssertions.anyCauseMatches;
@@ -577,75 +576,6 @@ public class KafkaDynamicTableFactoryTest {
validator.accept(kafkaSource);
}
- private interface OffsetsRetriever
- extends Function<Collection<TopicPartition>, Map<TopicPartition, Long>> {}
-
- private interface TimestampOffsetsRetriever
- extends Function<Map<TopicPartition, Long>, Map<TopicPartition, OffsetAndTimestamp>> {}
-
- private static final class MockPartitionOffsetsRetriever
- implements OffsetsInitializer.PartitionOffsetsRetriever {
-
- public static final OffsetsRetriever UNSUPPORTED_RETRIEVAL =
- partitions -> {
- throw new UnsupportedOperationException(
- "The method was not supposed to be called");
- };
- private final OffsetsRetriever committedOffsets;
- private final OffsetsRetriever endOffsets;
- private final OffsetsRetriever beginningOffsets;
- private final TimestampOffsetsRetriever offsetsForTimes;
-
- static MockPartitionOffsetsRetriever noInteractions() {
- return new MockPartitionOffsetsRetriever(
- UNSUPPORTED_RETRIEVAL,
- UNSUPPORTED_RETRIEVAL,
- UNSUPPORTED_RETRIEVAL,
- partitions -> {
- throw new UnsupportedOperationException(
- "The method was not supposed to be called");
- });
- }
-
- static MockPartitionOffsetsRetriever timestampAndEnd(
- TimestampOffsetsRetriever retriever, OffsetsRetriever endOffsets) {
- return new MockPartitionOffsetsRetriever(
- UNSUPPORTED_RETRIEVAL, endOffsets, UNSUPPORTED_RETRIEVAL, retriever);
- }
-
- private MockPartitionOffsetsRetriever(
- OffsetsRetriever committedOffsets,
- OffsetsRetriever endOffsets,
- OffsetsRetriever beginningOffsets,
- TimestampOffsetsRetriever offsetsForTimes) {
- this.committedOffsets = committedOffsets;
- this.endOffsets = endOffsets;
- this.beginningOffsets = beginningOffsets;
- this.offsetsForTimes = offsetsForTimes;
- }
-
- @Override
- public Map<TopicPartition, Long> committedOffsets(Collection<TopicPartition> partitions) {
- return committedOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions) {
- return endOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions) {
- return beginningOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
- Map<TopicPartition, Long> timestampsToSearch) {
- return offsetsForTimes.apply(timestampsToSearch);
- }
- }
-
@Test
public void testTableSink() {
final Map<String, String> modifiedOptions =
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
index 959c44cf..ba424f79 100644
--- a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/table/UpsertKafkaDynamicTableFactoryTest.java
@@ -39,6 +39,7 @@ import org.apache.flink.streaming.api.operators.StreamOperatorFactory;
import org.apache.flink.streaming.api.transformations.SourceTransformation;
import org.apache.flink.streaming.connectors.kafka.config.BoundedMode;
import org.apache.flink.streaming.connectors.kafka.config.StartupMode;
+import org.apache.flink.streaming.connectors.kafka.testutils.MockPartitionOffsetsRetriever;
import org.apache.flink.streaming.runtime.operators.sink.SinkWriterOperatorFactory;
import org.apache.flink.table.api.DataTypes;
import org.apache.flink.table.api.ValidationException;
@@ -74,14 +75,12 @@ import org.junit.Test;
import org.junit.rules.ExpectedException;
import java.util.Arrays;
-import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;
import java.util.Properties;
import java.util.function.Consumer;
-import java.util.function.Function;
import static org.apache.flink.core.testutils.FlinkMatchers.containsCause;
import static org.apache.flink.streaming.connectors.kafka.table.KafkaConnectorOptions.ScanBoundedMode;
@@ -819,73 +818,4 @@ public class UpsertKafkaDynamicTableFactoryTest extends TestLogger {
final KafkaSource<?> kafkaSource = assertKafkaSource(provider);
validator.accept(kafkaSource);
}
-
- private interface OffsetsRetriever
- extends Function<Collection<TopicPartition>, Map<TopicPartition, Long>> {}
-
- private interface TimestampOffsetsRetriever
- extends Function<Map<TopicPartition, Long>, Map<TopicPartition, OffsetAndTimestamp>> {}
-
- private static final class MockPartitionOffsetsRetriever
- implements OffsetsInitializer.PartitionOffsetsRetriever {
-
- public static final OffsetsRetriever UNSUPPORTED_RETRIEVAL =
- partitions -> {
- throw new UnsupportedOperationException(
- "The method was not supposed to be called");
- };
- private final OffsetsRetriever committedOffsets;
- private final OffsetsRetriever endOffsets;
- private final OffsetsRetriever beginningOffsets;
- private final TimestampOffsetsRetriever offsetsForTimes;
-
- static MockPartitionOffsetsRetriever noInteractions() {
- return new MockPartitionOffsetsRetriever(
- UNSUPPORTED_RETRIEVAL,
- UNSUPPORTED_RETRIEVAL,
- UNSUPPORTED_RETRIEVAL,
- partitions -> {
- throw new UnsupportedOperationException(
- "The method was not supposed to be called");
- });
- }
-
- static MockPartitionOffsetsRetriever timestampAndEnd(
- TimestampOffsetsRetriever retriever, OffsetsRetriever endOffsets) {
- return new MockPartitionOffsetsRetriever(
- UNSUPPORTED_RETRIEVAL, endOffsets, UNSUPPORTED_RETRIEVAL, retriever);
- }
-
- private MockPartitionOffsetsRetriever(
- OffsetsRetriever committedOffsets,
- OffsetsRetriever endOffsets,
- OffsetsRetriever beginningOffsets,
- TimestampOffsetsRetriever offsetsForTimes) {
- this.committedOffsets = committedOffsets;
- this.endOffsets = endOffsets;
- this.beginningOffsets = beginningOffsets;
- this.offsetsForTimes = offsetsForTimes;
- }
-
- @Override
- public Map<TopicPartition, Long> committedOffsets(Collection<TopicPartition> partitions) {
- return committedOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions) {
- return endOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions) {
- return beginningOffsets.apply(partitions);
- }
-
- @Override
- public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
- Map<TopicPartition, Long> timestampsToSearch) {
- return offsetsForTimes.apply(timestampsToSearch);
- }
- }
}
diff --git a/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/testutils/MockPartitionOffsetsRetriever.java b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/testutils/MockPartitionOffsetsRetriever.java
new file mode 100644
index 00000000..175bddd6
--- /dev/null
+++ b/flink-connector-kafka/src/test/java/org/apache/flink/streaming/connectors/kafka/testutils/MockPartitionOffsetsRetriever.java
@@ -0,0 +1,102 @@
+/*
+ * 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.flink.streaming.connectors.kafka.testutils;
+
+import org.apache.flink.connector.kafka.source.enumerator.initializer.OffsetsInitializer;
+
+import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
+import org.apache.kafka.common.TopicPartition;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.function.Function;
+
+/** Fake {@link OffsetsInitializer.PartitionOffsetsRetriever} for unit tests. */
+public final class MockPartitionOffsetsRetriever
+ implements OffsetsInitializer.PartitionOffsetsRetriever {
+
+ /** Fake offsets retriever for a given set of topic partitions. */
+ public interface OffsetsRetriever
+ extends Function<Collection<TopicPartition>, Map<TopicPartition, Long>> {}
+
+ /**
+ * Fake offsets retrieve for a given set of topic partitions and their target timestamp
+ * position.
+ */
+ public interface TimestampOffsetsRetriever
+ extends Function<Map<TopicPartition, Long>, Map<TopicPartition, OffsetAndTimestamp>> {}
+
+ public static final OffsetsRetriever UNSUPPORTED_RETRIEVAL =
+ partitions -> {
+ throw new UnsupportedOperationException("The method was not supposed to be called");
+ };
+ private final OffsetsRetriever committedOffsets;
+ private final OffsetsRetriever endOffsets;
+ private final OffsetsRetriever beginningOffsets;
+ private final TimestampOffsetsRetriever offsetsForTimes;
+
+ public static MockPartitionOffsetsRetriever noInteractions() {
+ return new MockPartitionOffsetsRetriever(
+ UNSUPPORTED_RETRIEVAL,
+ UNSUPPORTED_RETRIEVAL,
+ UNSUPPORTED_RETRIEVAL,
+ partitions -> {
+ throw new UnsupportedOperationException(
+ "The method was not supposed to be called");
+ });
+ }
+
+ public static MockPartitionOffsetsRetriever timestampAndEnd(
+ TimestampOffsetsRetriever retriever, OffsetsRetriever endOffsets) {
+ return new MockPartitionOffsetsRetriever(
+ UNSUPPORTED_RETRIEVAL, endOffsets, UNSUPPORTED_RETRIEVAL, retriever);
+ }
+
+ private MockPartitionOffsetsRetriever(
+ OffsetsRetriever committedOffsets,
+ OffsetsRetriever endOffsets,
+ OffsetsRetriever beginningOffsets,
+ TimestampOffsetsRetriever offsetsForTimes) {
+ this.committedOffsets = committedOffsets;
+ this.endOffsets = endOffsets;
+ this.beginningOffsets = beginningOffsets;
+ this.offsetsForTimes = offsetsForTimes;
+ }
+
+ @Override
+ public Map<TopicPartition, Long> committedOffsets(Collection<TopicPartition> partitions) {
+ return committedOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, Long> endOffsets(Collection<TopicPartition> partitions) {
+ return endOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, Long> beginningOffsets(Collection<TopicPartition> partitions) {
+ return beginningOffsets.apply(partitions);
+ }
+
+ @Override
+ public Map<TopicPartition, OffsetAndTimestamp> offsetsForTimes(
+ Map<TopicPartition, Long> timestampsToSearch) {
+ return offsetsForTimes.apply(timestampsToSearch);
+ }
+}