def test_set_delivery_guarantee(self):
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(self._build_serialization_schema()) \
.build()
guarantee = get_field_value(sink.get_java_function(),
'deliveryGuarantee')
self.assertEqual(guarantee.toString(), 'none')
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_delivery_guarantee(DeliveryGuarantee.AT_LEAST_ONCE) \
.set_record_serializer(self._build_serialization_schema()) \
.build()
guarantee = get_field_value(sink.get_java_function(),
'deliveryGuarantee')
self.assertEqual(guarantee.toString(), 'at-least-once')
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_delivery_guarantee(DeliveryGuarantee.EXACTLY_ONCE) \
.set_record_serializer(self._build_serialization_schema()) \
.build()
guarantee = get_field_value(sink.get_java_function(),
'deliveryGuarantee')
self.assertEqual(guarantee.toString(), 'exactly-once')
def _check_record(data, topic, serialized_data):
input_type = Types.ROW([Types.STRING()])
serialization_schema = KafkaRecordSerializationSchema.builder() \
.set_topic_selector(_select) \
.set_value_serialization_schema(
JsonRowSerializationSchema.builder().with_type_info(input_type).build()) \
.build()
jvm = get_gateway().jvm
serialization_schema._j_serialization_schema.open(
jvm.org.apache.flink.connector.testutils.formats.
DummyInitializationContext(),
jvm.org.apache.flink.connector.kafka.sink.
DefaultKafkaSinkContext(0, 1, jvm.java.util.Properties()))
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(serialization_schema) \
.build()
ds = MockDataStream(Types.ROW([Types.STRING()]))
ds.sink_to(sink)
row = Row(data)
topic_row = ds.feed(row) # type: Row
j_record = serialization_schema._j_serialization_schema.serialize(
to_java_data_structure(topic_row), None, None)
self.assertEqual(j_record.topic(), topic)
self.assertIsNone(j_record.key())
self.assertEqual(j_record.value(), serialized_data)
def test_set_property(self):
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(self._build_serialization_schema()) \
.set_property('test-key', 'test-value') \
.build()
config = get_field_value(sink.get_java_function(),
'kafkaProducerConfig')
self.assertEqual(config.get('test-key'), 'test-value')
def test_set_transactional_id_prefix(self):
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_transactional_id_prefix('test-prefix') \
.set_record_serializer(self._build_serialization_schema()) \
.build()
prefix = get_field_value(sink.get_java_function(),
'transactionalIdPrefix')
self.assertEqual(prefix, 'test-prefix')
def test_compile(self):
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(self._build_serialization_schema()) \
.build()
ds = self.env.from_collection([], type_info=Types.STRING())
ds.sink_to(sink)
plan = json.loads(self.env.get_execution_plan())
self.assertEqual(plan['nodes'][1]['type'], 'Sink: Writer')
self.assertEqual(plan['nodes'][2]['type'], 'Sink: Committer')
def test_set_record_serializer(self):
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(self._build_serialization_schema()) \
.build()
serializer = get_field_value(sink.get_java_function(),
'recordSerializer')
self.assertEqual(
serializer.getClass().getCanonicalName(),
'org.apache.flink.connector.kafka.sink.'
'KafkaRecordSerializationSchemaBuilder.'
'KafkaRecordSerializationSchemaWrapper')
topic_selector = get_field_value(serializer, 'topicSelector')
self.assertEqual(topic_selector.apply(None), 'test-topic')
value_serializer = get_field_value(serializer,
'valueSerializationSchema')
self.assertEqual(
value_serializer.getClass().getCanonicalName(),
'org.apache.flink.api.common.serialization.SimpleStringSchema')
def _check_record(data, topic, serialized_data):
input_type = Types.ROW([Types.STRING()])
serialization_schema = KafkaRecordSerializationSchema.builder() \
.set_topic_selector(_select) \
.set_value_serialization_schema(
JsonRowSerializationSchema.builder().with_type_info(input_type).build()) \
.build()
sink = KafkaSink.builder() \
.set_bootstrap_servers('localhost:9092') \
.set_record_serializer(serialization_schema) \
.build()
ds = MockDataStream(Types.ROW([Types.STRING()]))
ds.sink_to(sink)
row = Row(data)
topic_row = ds.feed(row) # type: Row
j_record = serialization_schema._j_serialization_schema.serialize(
to_java_data_structure(topic_row), None, None)
self.assertEqual(j_record.topic(), topic)
self.assertIsNone(j_record.key())
self.assertEqual(j_record.value(), serialized_data)