org.apache.flink.streaming.connectors.kafka.partitioner.FlinkKafkaPartitioner Java Examples

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka08TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema
	);
}
 

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka011TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema
	);
}
 

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka08TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema
	);
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaProducer(Properties producerConfig, KeyedSerializationSchema<T> schema, FlinkKafkaPartitioner partitioner) {

	super(DUMMY_TOPIC, schema, producerConfig, partitioner);

	this.mockProducer = mock(KafkaProducer.class);
	when(mockProducer.send(any(ProducerRecord.class), any(Callback.class))).thenAnswer(new Answer<Object>() {
		@Override
		public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
			pendingCallbacks.add(invocationOnMock.getArgument(1));
			return null;
		}
	});

	this.pendingCallbacks = new ArrayList<>();
	this.flushLatch = new MultiShotLatch();
}
 

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka010TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema
	);
}
 

@SuppressWarnings("unchecked")
private Optional<FlinkKafkaPartitioner<Row>> getFlinkKafkaPartitioner(DescriptorProperties descriptorProperties) {
	return descriptorProperties
		.getOptionalString(CONNECTOR_SINK_PARTITIONER)
		.flatMap((String partitionerString) -> {
			switch (partitionerString) {
				case CONNECTOR_SINK_PARTITIONER_VALUE_FIXED:
					return Optional.of(new FlinkFixedPartitioner<>());
				case CONNECTOR_SINK_PARTITIONER_VALUE_ROUND_ROBIN:
					return Optional.empty();
				case CONNECTOR_SINK_PARTITIONER_VALUE_CUSTOM:
					final Class<? extends FlinkKafkaPartitioner> partitionerClass =
						descriptorProperties.getClass(CONNECTOR_SINK_PARTITIONER_CLASS, FlinkKafkaPartitioner.class);
					return Optional.of((FlinkKafkaPartitioner<Row>) InstantiationUtil.instantiate(partitionerClass));
				default:
					throw new TableException("Unsupported sink partitioner. Validator should have checked that.");
			}
		});
}
 

@SuppressWarnings("unchecked")
DummyFlinkKafkaProducer(Properties producerConfig, KeyedSerializationSchema<T> schema, FlinkKafkaPartitioner partitioner) {

	super(DUMMY_TOPIC, schema, producerConfig, partitioner);

	this.mockProducer = mock(KafkaProducer.class);
	when(mockProducer.send(any(ProducerRecord.class), any(Callback.class))).thenAnswer(new Answer<Object>() {
		@Override
		public Object answer(InvocationOnMock invocationOnMock) throws Throwable {
			pendingCallbacks.add(invocationOnMock.getArgument(1));
			return null;
		}
	});

	this.pendingCallbacks = new ArrayList<>();
	this.flushLatch = new MultiShotLatch();
}
 

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka010TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema
	);
}
 

@Override
protected FlinkKafkaProducerBase<Row> createKafkaProducer(
		String topic,
		Properties properties,
		SerializationSchema<Row> serializationSchema,
		Optional<FlinkKafkaPartitioner<Row>> partitioner) {
	return new FlinkKafkaProducer08<>(
		topic,
		serializationSchema,
		properties,
		partitioner.orElse(null));
}
 

public Kafka09TableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {
	super(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

public KafkaTableSink(
	TableSchema schema,
	String topic,
	Properties properties,
	Optional<FlinkKafkaPartitioner<Row>> partitioner,
	SerializationSchema<Row> serializationSchema) {

	super(schema, topic, properties, partitioner, serializationSchema);
}
 

@Override KafkaTableSinkBase newTableSink(TableSchema schema, String topic, Properties properties,
    Optional<FlinkKafkaPartitioner<Row>> partitioner, SerializationSchema<Row> serializationSchema) {
    return new Kafka010TableSink(
        schema,
        topic,
        properties,
        partitioner,
        serializationSchema
    );
}
 

@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);
}
 

@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
protected KafkaTableSinkBase createKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka09TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

/**
 * Tests that partitions list is determinate and correctly provided to custom partitioner.
 */
@SuppressWarnings("unchecked")
@Test
public void testPartitionerInvokedWithDeterminatePartitionList() throws Exception {
	FlinkKafkaPartitioner<String> mockPartitioner = mock(FlinkKafkaPartitioner.class);

	RuntimeContext mockRuntimeContext = mock(StreamingRuntimeContext.class);
	when(mockRuntimeContext.getIndexOfThisSubtask()).thenReturn(0);
	when(mockRuntimeContext.getNumberOfParallelSubtasks()).thenReturn(1);

	// out-of-order list of 4 partitions
	List<PartitionInfo> mockPartitionsList = new ArrayList<>(4);
	mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 3, null, null, null));
	mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 1, null, null, null));
	mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 0, null, null, null));
	mockPartitionsList.add(new PartitionInfo(DummyFlinkKafkaProducer.DUMMY_TOPIC, 2, null, null, null));

	final DummyFlinkKafkaProducer<String> producer = new DummyFlinkKafkaProducer<>(
		FakeStandardProducerConfig.get(), new KeyedSerializationSchemaWrapper<>(new SimpleStringSchema()), mockPartitioner);
	producer.setRuntimeContext(mockRuntimeContext);

	final KafkaProducer mockProducer = producer.getMockKafkaProducer();
	when(mockProducer.partitionsFor(anyString())).thenReturn(mockPartitionsList);
	when(mockProducer.metrics()).thenReturn(null);

	producer.open(new Configuration());
	verify(mockPartitioner, times(1)).open(0, 1);

	producer.invoke("foobar", SinkContextUtil.forTimestamp(0));
	verify(mockPartitioner, times(1)).partition(
		"foobar", null, "foobar".getBytes(), DummyFlinkKafkaProducer.DUMMY_TOPIC, new int[] {0, 1, 2, 3});
}
 

@Override
public <T> DataStreamSink<T> produceIntoKafka(DataStream<T> stream, String topic, KeyedSerializationSchema<T> serSchema, Properties props, FlinkKafkaPartitioner<T> partitioner) {
	return stream.addSink(new FlinkKafkaProducer<T>(
		topic,
		serSchema,
		props,
		Optional.ofNullable(partitioner),
		producerSemantic,
		FlinkKafkaProducer.DEFAULT_KAFKA_PRODUCERS_POOL_SIZE));
}
 

public Kafka010TableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {
	super(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

@Override
protected KafkaTableSinkBase createKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka011TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

@Override
protected KafkaTableSinkBase getExpectedKafkaTableSink(
	TableSchema schema,
	String topic,
	Properties properties,
	Optional<FlinkKafkaPartitioner<Row>> partitioner,
	SerializationSchema<Row> serializationSchema) {

	return new KafkaTableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

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

/**
 * The main constructor for creating a FlinkKafkaProducer.
 *
 * @param defaultTopicId The default topic to write data to
 * @param serializationSchema A serializable serialization schema for turning user objects into a kafka-consumable byte[] supporting key/value messages
 * @param producerConfig Configuration properties for the KafkaProducer. 'bootstrap.servers.' is the only required argument.
 * @param customPartitioner A serializable partitioner for assigning messages to Kafka partitions. Passing null will use Kafka's partitioner.
 */
public FlinkKafkaProducerBase(String defaultTopicId, KeyedSerializationSchema<IN> serializationSchema, Properties producerConfig, FlinkKafkaPartitioner<IN> customPartitioner) {
	requireNonNull(defaultTopicId, "TopicID not set");
	requireNonNull(serializationSchema, "serializationSchema not set");
	requireNonNull(producerConfig, "producerConfig not set");
	ClosureCleaner.clean(customPartitioner, ExecutionConfig.ClosureCleanerLevel.RECURSIVE, true);
	ClosureCleaner.ensureSerializable(serializationSchema);

	this.defaultTopicId = defaultTopicId;
	this.schema = serializationSchema;
	this.producerConfig = producerConfig;
	this.flinkKafkaPartitioner = customPartitioner;

	// set the producer configuration properties for kafka record key value serializers.
	if (!producerConfig.containsKey(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG)) {
		this.producerConfig.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
	} else {
		LOG.warn("Overwriting the '{}' is not recommended", ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG);
	}

	if (!producerConfig.containsKey(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG)) {
		this.producerConfig.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class.getName());
	} else {
		LOG.warn("Overwriting the '{}' is not recommended", ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG);
	}

	// eagerly ensure that bootstrap servers are set.
	if (!this.producerConfig.containsKey(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG)) {
		throw new IllegalArgumentException(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG + " must be supplied in the producer config properties.");
	}

	this.topicPartitionsMap = new HashMap<>();
}
 

abstract KafkaTableSinkBase newTableSink(
    TableSchema schema,
    String topic,
    Properties properties,
    Optional<FlinkKafkaPartitioner<Row>> partitioner,
    SerializationSchema<Row> serializationSchema
);
 

@Override
protected KafkaTableSinkBase createKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka010TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

@Override
protected FlinkKafkaProducerBase<Row> createKafkaProducer(
		String topic,
		Properties properties,
		SerializationSchema<Row> serializationSchema,
		Optional<FlinkKafkaPartitioner<Row>> partitioner) {
	return new FlinkKafkaProducer010<>(
		topic,
		serializationSchema,
		properties,
		partitioner.orElse(null));
}
 

@Override
protected FlinkKafkaProducerBase<Row> createKafkaProducer(
		String topic,
		Properties properties,
		SerializationSchema<Row> serializationSchema,
		Optional<FlinkKafkaPartitioner<Row>> partitioner) {
	return new FlinkKafkaProducer010<>(
		topic,
		serializationSchema,
		properties,
		partitioner.orElse(null));
}
 

@Override
protected KafkaTableSinkBase createKafkaTableSink(
		TableSchema schema,
		String topic,
		Properties properties,
		Optional<FlinkKafkaPartitioner<Row>> partitioner,
		SerializationSchema<Row> serializationSchema) {

	return new Kafka09TableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}
 

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

@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));
}
 

@Override
protected KafkaTableSinkBase createKafkaTableSink(
	TableSchema schema,
	String topic,
	Properties properties,
	Optional<FlinkKafkaPartitioner<Row>> partitioner,
	SerializationSchema<Row> serializationSchema) {

	return new KafkaTableSink(
		schema,
		topic,
		properties,
		partitioner,
		serializationSchema);
}