org.apache.flink.streaming.api.operators.StreamSink Java Examples

/**
 * This test is meant to assure that testAtLeastOnceProducer is valid by testing that if flushing is disabled,
 * the snapshot method does indeed finishes without waiting for pending records;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@SuppressWarnings("unchecked")
@Test(timeout = 5000)
public void testDoesNotWaitForPendingRecordsIfFlushingDisabled() throws Throwable {
	final DummyFlinkKafkaProducer<String> producer = new DummyFlinkKafkaProducer<>(
		FakeStandardProducerConfig.get(), new KeyedSerializationSchemaWrapper<>(new SimpleStringSchema()), null);
	producer.setFlushOnCheckpoint(false);

	final KafkaProducer<?, ?> mockProducer = producer.getMockKafkaProducer();

	final OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg"));

	// make sure that all callbacks have not been completed
	verify(mockProducer, times(1)).send(any(ProducerRecord.class), any(Callback.class));

	// should return even if there are pending records
	testHarness.snapshot(123L, 123L);

	testHarness.close();
}
 

@Override
protected OneInputStreamOperatorTestHarness<Integer, Object> createTestHarness() throws Exception {
	FlinkKafkaProducer011<Integer> kafkaProducer = new FlinkKafkaProducer011<>(
		TOPIC,
		integerKeyedSerializationSchema,
		createProperties(),
		FlinkKafkaProducer011.Semantic.EXACTLY_ONCE
	).ignoreFailuresAfterTransactionTimeout();

	return new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(kafkaProducer),
		1,
		1,
		0,
		IntSerializer.INSTANCE,
		new OperatorID(1, 1));
}
 

private OneInputStreamOperatorTestHarness<Integer, Object> createTestHarness(
		String topic,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		Semantic semantic) throws Exception {
	Properties properties = createProperties();

	FlinkKafkaProducer011<Integer> kafkaProducer = new FlinkKafkaProducer011<>(
		topic,
		integerKeyedSerializationSchema,
		properties,
		semantic);

	return new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(kafkaProducer),
		maxParallelism,
		parallelism,
		subtaskIndex,
		IntSerializer.INSTANCE,
		new OperatorID(42, 44));
}
 

private OneInputStreamOperatorTestHarness<Integer, Object> createTestHarness(
	String topic,
	int maxParallelism,
	int parallelism,
	int subtaskIndex,
	FlinkKafkaProducer.Semantic semantic) throws Exception {
	Properties properties = createProperties();

	FlinkKafkaProducer<Integer> kafkaProducer = new FlinkKafkaProducer<>(
		topic,
		integerKeyedSerializationSchema,
		properties,
		semantic);

	return new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(kafkaProducer),
		maxParallelism,
		parallelism,
		subtaskIndex,
		IntSerializer.INSTANCE,
		new OperatorID(42, 44));
}
 

/**
 * This test is meant to assure that testAtLeastOnceSink is valid by testing that if flushing is disabled,
 * the snapshot method does indeed finishes without waiting for pending requests;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@Test(timeout = 5000)
public void testDoesNotWaitForPendingRequestsIfFlushingDisabled() throws Exception {
	final DummyElasticsearchSink<String> sink = new DummyElasticsearchSink<>(
		new HashMap<String, String>(), new SimpleSinkFunction<String>(), new DummyRetryFailureHandler());
	sink.disableFlushOnCheckpoint(); // disable flushing

	final OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

	testHarness.open();

	// setup the next bulk request, and let bulk request succeed
	sink.setMockItemFailuresListForNextBulkItemResponses(Collections.singletonList(new Exception("artificial failure for record")));
	testHarness.processElement(new StreamRecord<>("msg-1"));
	verify(sink.getMockBulkProcessor(), times(1)).add(any(IndexRequest.class));

	// the snapshot should not block even though we haven't flushed the bulk request
	testHarness.snapshot(1L, 1000L);

	testHarness.close();
}
 

/**
 * This test is meant to assure that testAtLeastOnceProducer is valid by testing that if flushing is disabled,
 * the snapshot method does indeed finishes without waiting for pending records;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@SuppressWarnings("unchecked")
@Test//(timeout = 5000)
public void testDoesNotWaitForPendingRecordsIfFlushingDisabled() throws Throwable {
    Properties props = dummyProperties();
    props.setProperty("flushoncheckpoint", "false");

    final DummyFlinkPulsarSink<String> sink = new DummyFlinkPulsarSink<>(dummyClientConf(), props, mock(TopicKeyExtractor.class), null);

    final Producer mockProducer = sink.getProducer("tp");

    final OneInputStreamOperatorTestHarness<String, Object> testHarness =
            new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

    testHarness.open();

    testHarness.processElement(new StreamRecord<>("msg"));

    // make sure that all callbacks have not been completed
    verify(mockProducer, times(1)).newMessage();

    // should return even if there are pending records
    testHarness.snapshot(123L, 123L);

    testHarness.close();
}
 

static OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Object> createTestSinkWithBulkEncoder(
		final File outDir,
		final int totalParallelism,
		final int taskIdx,
		final long bucketCheckInterval,
		final BucketAssigner<Tuple2<String, Integer>, String> bucketer,
		final BulkWriter.Factory<Tuple2<String, Integer>> writer,
		final BucketFactory<Tuple2<String, Integer>, String> bucketFactory,
		final String partFilePrefix,
		final String partFileSuffix) throws Exception {

	StreamingFileSink<Tuple2<String, Integer>> sink = StreamingFileSink
		.forBulkFormat(new Path(outDir.toURI()), writer)
		.withBucketAssigner(bucketer)
		.withBucketCheckInterval(bucketCheckInterval)
		.withBucketFactory(bucketFactory)
		.withPartFilePrefix(partFilePrefix)
		.withPartFileSuffix(partFileSuffix)
		.build();

	return new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink), MAX_PARALLELISM, totalParallelism, taskIdx);
}
 

@Override
protected OneInputStreamOperatorTestHarness<Integer, Object> createTestHarness() throws Exception {
	FlinkKafkaProducer<Integer> kafkaProducer = new FlinkKafkaProducer<>(
		TOPIC,
		integerKeyedSerializationSchema,
		createProperties(),
		FlinkKafkaProducer.Semantic.EXACTLY_ONCE
	).ignoreFailuresAfterTransactionTimeout();

	return new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(kafkaProducer),
		1,
		1,
		0,
		IntSerializer.INSTANCE,
		new OperatorID(1, 1));
}
 

/**
 * This test is meant to assure that testAtLeastOnceSink is valid by testing that if flushing is disabled,
 * the snapshot method does indeed finishes without waiting for pending requests;
 * we set a timeout because the test will not finish if the logic is broken.
 */
@Test(timeout = 5000)
public void testDoesNotWaitForPendingRequestsIfFlushingDisabled() throws Exception {
	final DummyElasticsearchSink<String> sink = new DummyElasticsearchSink<>(
		new HashMap<String, String>(), new SimpleSinkFunction<String>(), new DummyRetryFailureHandler());
	sink.disableFlushOnCheckpoint(); // disable flushing

	final OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

	testHarness.open();

	// setup the next bulk request, and let bulk request succeed
	sink.setMockItemFailuresListForNextBulkItemResponses(Collections.singletonList(new Exception("artificial failure for record")));
	testHarness.processElement(new StreamRecord<>("msg-1"));
	verify(sink.getMockBulkProcessor(), times(1)).add(any(IndexRequest.class));

	// the snapshot should not block even though we haven't flushed the bulk request
	testHarness.snapshot(1L, 1000L);

	testHarness.close();
}
 

static OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Object> createTestSinkWithBulkEncoder(
		final File outDir,
		final int totalParallelism,
		final int taskIdx,
		final long bucketCheckInterval,
		final BucketAssigner<Tuple2<String, Integer>, String> bucketer,
		final BulkWriter.Factory<Tuple2<String, Integer>> writer,
		final BucketFactory<Tuple2<String, Integer>, String> bucketFactory) throws Exception {

	StreamingFileSink<Tuple2<String, Integer>> sink = StreamingFileSink
			.forBulkFormat(new Path(outDir.toURI()), writer)
			.withBucketAssigner(bucketer)
			.withBucketCheckInterval(bucketCheckInterval)
			.withBucketFactory(bucketFactory)
			.build();

	return new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink), MAX_PARALLELISM, totalParallelism, taskIdx);
}
 

static OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Object> createCustomRescalingTestSink(
		final File outDir,
		final int totalParallelism,
		final int taskIdx,
		final long bucketCheckInterval,
		final BucketAssigner<Tuple2<String, Integer>, String> bucketer,
		final Encoder<Tuple2<String, Integer>> writer,
		final RollingPolicy<Tuple2<String, Integer>, String> rollingPolicy,
		final BucketFactory<Tuple2<String, Integer>, String> bucketFactory) throws Exception {

	StreamingFileSink<Tuple2<String, Integer>> sink = StreamingFileSink
			.forRowFormat(new Path(outDir.toURI()), writer)
			.withBucketAssigner(bucketer)
			.withRollingPolicy(rollingPolicy)
			.withBucketCheckInterval(bucketCheckInterval)
			.withBucketFactory(bucketFactory)
			.build();

	return new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink), MAX_PARALLELISM, totalParallelism, taskIdx);
}
 

private OneInputStreamOperatorTestHarness<Integer, Object> createTestHarness(
		String topic,
		int maxParallelism,
		int parallelism,
		int subtaskIndex,
		Semantic semantic) throws Exception {
	Properties properties = createProperties();

	FlinkKafkaProducer011<Integer> kafkaProducer = new FlinkKafkaProducer011<>(
		topic,
		integerKeyedSerializationSchema,
		properties,
		semantic);

	return new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(kafkaProducer),
		maxParallelism,
		parallelism,
		subtaskIndex,
		IntSerializer.INSTANCE,
		new OperatorID(42, 44));
}
 

static OneInputStreamOperatorTestHarness<Tuple2<String, Integer>, Object> createCustomRescalingTestSink(
		final File outDir,
		final int totalParallelism,
		final int taskIdx,
		final long bucketCheckInterval,
		final BucketAssigner<Tuple2<String, Integer>, String> bucketer,
		final Encoder<Tuple2<String, Integer>> writer,
		final RollingPolicy<Tuple2<String, Integer>, String> rollingPolicy,
		final BucketFactory<Tuple2<String, Integer>, String> bucketFactory) throws Exception {

	StreamingFileSink<Tuple2<String, Integer>> sink = StreamingFileSink
			.forRowFormat(new Path(outDir.toURI()), writer)
			.withBucketAssigner(bucketer)
			.withRollingPolicy(rollingPolicy)
			.withBucketCheckInterval(bucketCheckInterval)
			.withBucketFactory(bucketFactory)
			.build();

	return new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink), MAX_PARALLELISM, totalParallelism, taskIdx);
}
 

/**
 * Adds the given sink to this DataStream. Only streams with sinks added
 * will be executed once the {@link StreamExecutionEnvironment#execute()}
 * method is called.
 *
 * @param sinkFunction
 *            The object containing the sink's invoke function.
 * @return The closed DataStream.
 */
public DataStreamSink<T> addSink(SinkFunction<T> sinkFunction) {

	// read the output type of the input Transform to coax out errors about MissingTypeInfo
	transformation.getOutputType();

	// configure the type if needed
	if (sinkFunction instanceof InputTypeConfigurable) {
		((InputTypeConfigurable) sinkFunction).setInputType(getType(), getExecutionConfig());
	}

	StreamSink<T> sinkOperator = new StreamSink<>(clean(sinkFunction));

	DataStreamSink<T> sink = new DataStreamSink<>(this, sinkOperator);

	getExecutionEnvironment().addOperator(sink.getTransformation());
	return sink;
}
 

@Test
public void testAppendTableSink() throws IOException {
	JDBCAppendTableSink sink = JDBCAppendTableSink.builder()
		.setDrivername("foo")
		.setDBUrl("bar")
		.setQuery("insert into %s (id) values (?)")
		.setParameterTypes(FIELD_TYPES)
		.build();

	StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();

	DataStream<Row> ds = env.fromCollection(Collections.singleton(Row.of("foo")), ROW_TYPE);
	sink.emitDataStream(ds);

	Collection<Integer> sinkIds = env.getStreamGraph().getSinkIDs();
	assertEquals(1, sinkIds.size());
	int sinkId = sinkIds.iterator().next();

	StreamSink planSink = (StreamSink) env.getStreamGraph().getStreamNode(sinkId).getOperator();
	assertTrue(planSink.getUserFunction() instanceof JDBCSinkFunction);

	JDBCSinkFunction sinkFunction = (JDBCSinkFunction) planSink.getUserFunction();
	assertSame(sink.getOutputFormat(), sinkFunction.outputFormat);
}
 

/**
 * Adds the given sink to this DataStream. Only streams with sinks added
 * will be executed once the {@link StreamExecutionEnvironment#execute()}
 * method is called.
 *
 * @param sinkFunction
 *            The object containing the sink's invoke function.
 * @return The closed DataStream.
 */
public DataStreamSink<T> addSink(SinkFunction<T> sinkFunction) {

	// read the output type of the input Transform to coax out errors about MissingTypeInfo
	transformation.getOutputType();

	// configure the type if needed
	if (sinkFunction instanceof InputTypeConfigurable) {
		((InputTypeConfigurable) sinkFunction).setInputType(getType(), getExecutionConfig());
	}

	StreamSink<T> sinkOperator = new StreamSink<>(clean(sinkFunction));

	DataStreamSink<T> sink = new DataStreamSink<>(this, sinkOperator);

	getExecutionEnvironment().addOperator(sink.getTransformation());
	return sink;
}
 

/**
 * Manually run this to write binary snapshot data. Remove @Ignore to run.
 */
@Ignore
@Test
public void writeSnapshot() throws Exception {

	final File outDir = tempFolder.newFolder();

	BucketingSink<String> sink = new BucketingSink<String>(outDir.getAbsolutePath())
		.setWriter(new StringWriter<String>())
		.setBatchSize(5)
		.setPartPrefix(PART_PREFIX)
		.setInProgressPrefix("")
		.setPendingPrefix("")
		.setValidLengthPrefix("")
		.setInProgressSuffix(IN_PROGRESS_SUFFIX)
		.setPendingSuffix(PENDING_SUFFIX)
		.setValidLengthSuffix(VALID_LENGTH_SUFFIX);

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

	testHarness.setup();
	testHarness.open();

	testHarness.processElement(new StreamRecord<>("test1", 0L));
	testHarness.processElement(new StreamRecord<>("test2", 0L));

	checkLocalFs(outDir, 1, 1, 0, 0);

	testHarness.processElement(new StreamRecord<>("test3", 0L));
	testHarness.processElement(new StreamRecord<>("test4", 0L));
	testHarness.processElement(new StreamRecord<>("test5", 0L));

	checkLocalFs(outDir, 1, 4, 0, 0);

	OperatorSubtaskState snapshot = testHarness.snapshot(0L, 0L);

	OperatorSnapshotUtil.writeStateHandle(snapshot, "src/test/resources/bucketing-sink-migration-test-flink" + flinkGenerateSavepointVersion + "-snapshot");
	testHarness.close();
}
 

private OneInputStreamOperatorTestHarness<String, Object> createOpenedTestHarness(
	TestCassandraSink testCassandraSink) throws Exception {
	final StreamSink<String> testStreamSink = new StreamSink<>(testCassandraSink);
	final OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(testStreamSink);
	testHarness.open();
	return testHarness;
}
 

/**
 * Test ensuring that if a snapshot call happens right after an async exception is caught, it should be rethrown.
 */
@SuppressWarnings("ResultOfMethodCallIgnored")
@Test
public void testAsyncErrorRethrownOnCheckpoint() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));

	producer.getPendingRecordFutures().get(0).setException(new Exception("artificial async exception"));

	try {
		testHarness.snapshot(123L, 123L);
	} catch (Exception e) {
		// the next checkpoint should rethrow the async exception
		Assert.assertTrue(ExceptionUtils.findThrowableWithMessage(e, "artificial async exception").isPresent());

		// test succeeded
		return;
	}

	Assert.fail();
}
 

/**
 * Test ensuring that if an invoke call happens right after an async exception is caught, it should be rethrown.
 */
@SuppressWarnings("ResultOfMethodCallIgnored")
@Test
public void testAsyncErrorRethrownOnInvoke() throws Throwable {
	final DummyFlinkKinesisProducer<String> producer = new DummyFlinkKinesisProducer<>(new SimpleStringSchema());

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));

	producer.getPendingRecordFutures().get(0).setException(new Exception("artificial async exception"));

	try {
		testHarness.processElement(new StreamRecord<>("msg-2"));
	} catch (Exception e) {
		// the next invoke should rethrow the async exception
		Assert.assertTrue(ExceptionUtils.findThrowableWithMessage(e, "artificial async exception").isPresent());

		// test succeeded
		return;
	}

	Assert.fail();
}
 

/**
 * Test ensuring that if a snapshot call happens right after an async exception is caught, it should be rethrown.
 */
@Test
public void testAsyncErrorRethrownOnCheckpoint() throws Throwable {
	final DummyFlinkKafkaProducer<String> producer = new DummyFlinkKafkaProducer<>(
		FakeStandardProducerConfig.get(), new KeyedSerializationSchemaWrapper<>(new SimpleStringSchema()), null);

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));

	// let the message request return an async exception
	producer.getPendingCallbacks().get(0).onCompletion(null, new Exception("artificial async exception"));

	try {
		testHarness.snapshot(123L, 123L);
	} catch (Exception e) {
		// the next invoke should rethrow the async exception
		Assert.assertTrue(e.getCause().getMessage().contains("artificial async exception"));

		// test succeeded
		return;
	}

	Assert.fail();
}
 

/** Tests that any item failure in the listener callbacks is rethrown on an immediately following checkpoint. */
@Test
public void testItemFailureRethrownOnCheckpoint() throws Throwable {
	final DummyElasticsearchSink<String> sink = new DummyElasticsearchSink<>(
		new HashMap<String, String>(), new SimpleSinkFunction<String>(), new NoOpFailureHandler());

	final OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

	testHarness.open();

	// setup the next bulk request, and its mock item failures
	sink.setMockItemFailuresListForNextBulkItemResponses(Collections.singletonList(new Exception("artificial failure for record")));
	testHarness.processElement(new StreamRecord<>("msg"));
	verify(sink.getMockBulkProcessor(), times(1)).add(any(IndexRequest.class));

	// manually execute the next bulk request
	sink.manualBulkRequestWithAllPendingRequests();

	try {
		testHarness.snapshot(1L, 1000L);
	} catch (Exception e) {
		// the snapshot should have failed with the failure
		Assert.assertTrue(e.getCause().getCause().getMessage().contains("artificial failure for record"));

		// test succeeded
		return;
	}

	Assert.fail();
}
 

@Override
public <T> StreamSink<T> getProducerSink(String topic, KeyedSerializationSchema<T> serSchema, Properties props, FlinkKafkaPartitioner<T> partitioner) {
	return new StreamSink<>(new FlinkKafkaProducer011<>(
		topic,
		serSchema,
		props,
		Optional.ofNullable(partitioner),
		producerSemantic,
		FlinkKafkaProducer011.DEFAULT_KAFKA_PRODUCERS_POOL_SIZE));
}
 

@Test
public void testRestore() throws Exception {
	final File outDir = tempFolder.newFolder();

	ValidatingBucketingSink<String> sink = (ValidatingBucketingSink<String>)
			new ValidatingBucketingSink<String>(outDir.getAbsolutePath(), expectedBucketFilesPrefix)
		.setWriter(new StringWriter<String>())
		.setBatchSize(5)
		.setPartPrefix(PART_PREFIX)
		.setInProgressPrefix("")
		.setPendingPrefix("")
		.setValidLengthPrefix("")
		.setInProgressSuffix(IN_PROGRESS_SUFFIX)
		.setPendingSuffix(PENDING_SUFFIX)
		.setValidLengthSuffix(VALID_LENGTH_SUFFIX)
		.setUseTruncate(false); // don't use truncate because files do not exist

	OneInputStreamOperatorTestHarness<String, Object> testHarness = new OneInputStreamOperatorTestHarness<>(
		new StreamSink<>(sink), 10, 1, 0);
	testHarness.setup();

	testHarness.initializeState(
		OperatorSnapshotUtil.getResourceFilename(
			"bucketing-sink-migration-test-flink" + testMigrateVersion + "-snapshot"));

	testHarness.open();

	assertTrue(sink.initializeCalled);

	testHarness.processElement(new StreamRecord<>("test1", 0L));
	testHarness.processElement(new StreamRecord<>("test2", 0L));

	checkLocalFs(outDir, 1, 1, 0, 0);

	testHarness.close();
}
 

/**
 * Test ensuring that if an invoke call happens right after an async exception is caught, it should be rethrown.
 */
@Test
public void testAsyncErrorRethrownOnInvoke() throws Throwable {
	final DummyFlinkKafkaProducer<String> producer = new DummyFlinkKafkaProducer<>(
		FakeStandardProducerConfig.get(), new KeyedSerializationSchemaWrapper<>(new SimpleStringSchema()), null);

	OneInputStreamOperatorTestHarness<String, Object> testHarness =
		new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

	testHarness.open();

	testHarness.processElement(new StreamRecord<>("msg-1"));

	// let the message request return an async exception
	producer.getPendingCallbacks().get(0).onCompletion(null, new Exception("artificial async exception"));

	try {
		testHarness.processElement(new StreamRecord<>("msg-2"));
	} catch (Exception e) {
		// the next invoke should rethrow the async exception
		Assert.assertTrue(e.getCause().getMessage().contains("artificial async exception"));

		// test succeeded
		return;
	}

	Assert.fail();
}
 

@Override
public <T> StreamSink<T> getProducerSink(
		String topic,
		KeyedSerializationSchema<T> serSchema,
		Properties props,
		FlinkKafkaPartitioner<T> partitioner) {
	FlinkKafkaProducer08<T> prod = new FlinkKafkaProducer08<>(
			topic,
			serSchema,
			props,
			partitioner);
	prod.setFlushOnCheckpoint(true);
	return new StreamSink<>(prod);
}
 

@Override
public <T> StreamSink<T> getProducerSink(
		String topic,
		KeyedSerializationSchema<T> serSchema,
		Properties props,
		FlinkKafkaPartitioner<T> partitioner) {
	FlinkKafkaProducer09<T> prod = new FlinkKafkaProducer09<>(topic, serSchema, props, partitioner);
	prod.setFlushOnCheckpoint(true);
	return new StreamSink<>(prod);
}
 

/**
 * Test ensuring that if a snapshot call happens right after an async exception is caught, it should be rethrown.
 */
@Test
public void testAsyncErrorRethrownOnCheckpoint() throws Throwable {
    final DummyFlinkPulsarSink<String> producer = new DummyFlinkPulsarSink<>(
            dummyClientConf(), dummyProperties(), mock(TopicKeyExtractor.class), null);

    OneInputStreamOperatorTestHarness<String, Object> testHarness =
            new OneInputStreamOperatorTestHarness<>(new StreamSink<>(producer));

    testHarness.open();

    testHarness.processElement(new StreamRecord<>("msg-1"));

    // let the message request return an async exception
    producer.getPendingCallbacks().get(0).accept(null, new Exception("artificial async exception"));

    try {
        testHarness.snapshot(123L, 123L);
    } catch (Exception e) {
        // the next invoke should rethrow the async exception
        Assert.assertTrue(e.getCause().getMessage().contains("artificial async exception"));

        // test succeeded
        return;
    }

    Assert.fail();
}
 

/**
 * Test ensuring that if an invoke call happens right after an async exception is caught, it should be rethrown.
 */
@Test
public void testAsyncErrorRethrownOnInvoke() throws Throwable {
    DummyFlinkPulsarSink<String> sink = new DummyFlinkPulsarSink<>(dummyClientConf(), dummyProperties(), mock(TopicKeyExtractor.class), null);

    OneInputStreamOperatorTestHarness<String, Object> testHarness =
            new OneInputStreamOperatorTestHarness<>(new StreamSink<>(sink));

    testHarness.open();

    testHarness.processElement(new StreamRecord<>("msg-1"));

    // let the message request return an async exception
    sink.getPendingCallbacks().get(0).accept(null, new Exception("artificial async exception"));

    try {
        testHarness.processElement(new StreamRecord<>("msg-2"));
    } catch (Exception e) {
        // the next invoke should rethrow the async exception
        Assert.assertTrue(e.getCause().getMessage().contains("artificial async exception"));

        // test succeeded
        return;
    }

    Assert.fail();
}
 

/**
 * Creates a new {@code SinkTransformation} from the given input {@code StreamTransformation}.
 *
 * @param input The input {@code StreamTransformation}
 * @param name The name of the {@code StreamTransformation}, this will be shown in Visualizations and the Log
 * @param operator The sink operator
 * @param parallelism The parallelism of this {@code SinkTransformation}
 */
public SinkTransformation(
		StreamTransformation<T> input,
		String name,
		StreamSink<T> operator,
		int parallelism) {
	super(name, TypeExtractor.getForClass(Object.class), parallelism);
	this.input = input;
	this.operator = operator;
}