From 3795ba62e1cbc825329238669d5a79490bde1ced Mon Sep 17 00:00:00 2001 From: Maximilian Michels Date: Wed, 17 Feb 2016 10:34:01 +0100 Subject: [PATCH 1/2] [FLINK-3248] add constructor params and generic ConnectionFactory This adds more default constructor parameters to the RMQSource. In addition, users may override the setupConnectionFactory() method to return their onwn configured factory. --- .../connectors/rabbitmq/RMQSource.java | 79 +++++++++++++++++-- .../connectors/rabbitmq/RMQSourceTest.java | 47 ++++++++--- 2 files changed, 106 insertions(+), 20 deletions(-) diff --git a/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSource.java b/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSource.java index 09bb07c7082904..39f483b9eba3fb 100644 --- a/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSource.java +++ b/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSource.java @@ -60,6 +60,9 @@ * (correlation id is not set). * 3) No strong delivery guarantees (without checkpointing) with RabbitMQ auto-commit mode. * + * Users may overwrite the setupConnectionFactory() method to pass their setup their own + * ConnectionFactory in case the constructor parameters are not sufficient. + * * @param The type of the data read from RabbitMQ. */ public class RMQSource extends MultipleIdsMessageAcknowledgingSourceBase @@ -70,6 +73,9 @@ public class RMQSource extends MultipleIdsMessageAcknowledgingSourceBase schema; @@ -82,7 +88,6 @@ public class RMQSource extends MultipleIdsMessageAcknowledgingSourceBase extends MultipleIdsMessageAcknowledgingSourceBase deserializationSchema) { - this(hostName, queueName, false, deserializationSchema); + DeserializationSchema deserializationSchema) { + this(hostName, null, null, null, queueName, false, deserializationSchema); } /** @@ -104,7 +109,7 @@ public RMQSource(String hostName, String queueName, * at the producer. The correlation id must be unique. Otherwise the behavior of the source is * undefined. In doubt, set {@param usesCorrelationId} to false. When correlation ids are not * used, this source has at-least-once processing semantics when checkpointing is enabled. - * @param hostName The RabbiMQ broker's address to connect to. + * @param hostName The RabbitMQ broker's address to connect to. * @param queueName The queue to receive messages from. * @param usesCorrelationId Whether the messages received are supplied with a unique * id to deduplicate messages (in case of failed acknowledgments). @@ -113,20 +118,80 @@ public RMQSource(String hostName, String queueName, * into Java objects. */ public RMQSource(String hostName, String queueName, boolean usesCorrelationId, - DeserializationSchema deserializationSchema) { + DeserializationSchema deserializationSchema) { + this(hostName, null, null, null, queueName, usesCorrelationId, deserializationSchema); + } + + /** + * Creates a new RabbitMQ source. For exactly-once, you must set the correlation ids of messages + * at the producer. The correlation id must be unique. Otherwise the behavior of the source is + * undefined. In doubt, set {@param usesCorrelationId} to false. When correlation ids are not + * used, this source has at-least-once processing semantics when checkpointing is enabled. + * @param hostName The RabbitMQ broker's address to connect to. + * @param port The RabbitMQ broker's port. + * @param queueName The queue to receive messages from. + * @param usesCorrelationId Whether the messages received are supplied with a unique + * id to deduplicate messages (in case of failed acknowledgments). + * Only used when checkpointing is enabled. + * @param deserializationSchema A {@link DeserializationSchema} for turning the bytes received + * into Java objects. + */ + public RMQSource(String hostName, Integer port, + String queueName, boolean usesCorrelationId, + DeserializationSchema deserializationSchema) { + this(hostName, port, null, null, queueName, usesCorrelationId, deserializationSchema); + } + + /** + * Creates a new RabbitMQ source. For exactly-once, you must set the correlation ids of messages + * at the producer. The correlation id must be unique. Otherwise the behavior of the source is + * undefined. In doubt, set {@param usesCorrelationId} to false. When correlation ids are not + * used, this source has at-least-once processing semantics when checkpointing is enabled. + * @param hostName The RabbitMQ broker's address to connect to. + * @param port The RabbitMQ broker's port. + * @param queueName The queue to receive messages from. + * @param usesCorrelationId Whether the messages received are supplied with a unique + * id to deduplicate messages (in case of failed acknowledgments). + * Only used when checkpointing is enabled. + * @param deserializationSchema A {@link DeserializationSchema} for turning the bytes received + * into Java objects. + */ + public RMQSource(String hostName, Integer port, String username, String password, + String queueName, boolean usesCorrelationId, + DeserializationSchema deserializationSchema) { super(String.class); this.hostName = hostName; + this.port = port; + this.username = username; + this.password = password; this.queueName = queueName; this.usesCorrelationId = usesCorrelationId; this.schema = deserializationSchema; } + /** + * Initializes the connection to RMQ with a default connection factory. The user may override + * this method to setup and configure their own ConnectionFactory. + */ + protected ConnectionFactory setupConnectionFactory() { + return new ConnectionFactory(); + } + /** * Initializes the connection to RMQ. */ - protected void initializeConnection() { - ConnectionFactory factory = new ConnectionFactory(); + private void initializeConnection() { + ConnectionFactory factory = setupConnectionFactory(); factory.setHost(hostName); + if (port != null) { + factory.setPort(port); + } + if (username != null) { + factory.setUsername(username); + } + if (password != null) { + factory.setPassword(password); + } try { connection = factory.newConnection(); channel = connection.createChannel(); diff --git a/flink-streaming-connectors/flink-connector-rabbitmq/src/test/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSourceTest.java b/flink-streaming-connectors/flink-connector-rabbitmq/src/test/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSourceTest.java index e0eed70b9cb807..437c8081509d85 100644 --- a/flink-streaming-connectors/flink-connector-rabbitmq/src/test/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSourceTest.java +++ b/flink-streaming-connectors/flink-connector-rabbitmq/src/test/java/org/apache/flink/streaming/connectors/rabbitmq/RMQSourceTest.java @@ -19,8 +19,10 @@ import com.rabbitmq.client.AMQP; import com.rabbitmq.client.Channel; +import com.rabbitmq.client.ConnectionFactory; import com.rabbitmq.client.Envelope; import com.rabbitmq.client.QueueingConsumer; +import org.apache.flink.api.common.functions.RuntimeContext; import org.apache.flink.api.common.typeinfo.TypeInformation; import org.apache.flink.api.common.typeutils.base.StringSerializer; import org.apache.flink.api.java.tuple.Tuple2; @@ -28,6 +30,7 @@ import org.apache.flink.configuration.Configuration; import org.apache.flink.runtime.state.SerializedCheckpointData; import org.apache.flink.streaming.api.functions.source.SourceFunction; +import org.apache.flink.streaming.api.operators.StreamingRuntimeContext; import org.apache.flink.streaming.api.watermark.Watermark; import org.apache.flink.streaming.util.serialization.DeserializationSchema; import org.junit.After; @@ -72,18 +75,18 @@ public class RMQSourceTest { private volatile long messageId; - private boolean generateCorrelationIds = true; + private boolean generateCorrelationIds; private volatile Exception exception; @Before public void beforeTest() throws Exception { - source = new RMQTestSource<>("hostDummy", "queueDummy", true, new StringDeserializationScheme()); + source = new RMQTestSource<>("hostDummy", -1, "", "", "queueDummy", true, new StringDeserializationScheme()); source.open(config); - source.initializeConnection(); messageId = 0; + generateCorrelationIds = true; sourceThread = new Thread(new Runnable() { @Override @@ -240,31 +243,31 @@ public TypeInformation getProducedType() { private class RMQTestSource extends RMQSource { - public RMQTestSource(String hostName, String queueName, boolean usesCorrelationIds, - DeserializationSchema deserializationSchema) { - super(hostName, queueName, usesCorrelationIds, deserializationSchema); + public RMQTestSource(String hostName, Integer port, String username, String password, + String queueName, boolean usesCorrelationId, DeserializationSchema deserializationSchema) { + super(hostName, port, username, password, queueName, usesCorrelationId, deserializationSchema); } @Override - protected void initializeConnection() { - connection = Mockito.mock(Connection.class); - channel = Mockito.mock(Channel.class); + public void open(Configuration config) throws Exception { + super.open(config); + consumer = Mockito.mock(QueueingConsumer.class); // Mock for delivery final QueueingConsumer.Delivery deliveryMock = Mockito.mock(QueueingConsumer.Delivery.class); Mockito.when(deliveryMock.getBody()).thenReturn("test".getBytes()); - // Mock for envelope - Envelope envelope = Mockito.mock(Envelope.class); - Mockito.when(deliveryMock.getEnvelope()).thenReturn(envelope); - try { Mockito.when(consumer.nextDelivery()).thenReturn(deliveryMock); } catch (InterruptedException e) { fail("Couldn't setup up deliveryMock"); } + // Mock for envelope + Envelope envelope = Mockito.mock(Envelope.class); + Mockito.when(deliveryMock.getEnvelope()).thenReturn(envelope); + Mockito.when(envelope.getDeliveryTag()).thenAnswer(new Answer() { @Override public Long answer(InvocationOnMock invocation) throws Throwable { @@ -285,6 +288,24 @@ public String answer(InvocationOnMock invocation) throws Throwable { } + @Override + protected ConnectionFactory setupConnectionFactory() { + ConnectionFactory connectionFactory = Mockito.mock(ConnectionFactory.class); + Connection connection = Mockito.mock(Connection.class); + try { + Mockito.when(connectionFactory.newConnection()).thenReturn(connection); + Mockito.when(connection.createChannel()).thenReturn(Mockito.mock(Channel.class)); + } catch (IOException e) { + fail("Test environment couldn't be created."); + } + return connectionFactory; + } + + @Override + public RuntimeContext getRuntimeContext() { + return Mockito.mock(StreamingRuntimeContext.class); + } + @Override protected boolean addId(String uid) { assertEquals(false, autoAck); From 4112074f6340e57d66fe38fa203748a93d39ad88 Mon Sep 17 00:00:00 2001 From: Maximilian Michels Date: Thu, 18 Feb 2016 17:33:00 +0100 Subject: [PATCH 2/2] remove RMQTopology file --- .../connectors/rabbitmq/RMQTopology.java | 42 ------------------- 1 file changed, 42 deletions(-) delete mode 100644 flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQTopology.java diff --git a/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQTopology.java b/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQTopology.java deleted file mode 100644 index debcce4efc8cc7..00000000000000 --- a/flink-streaming-connectors/flink-connector-rabbitmq/src/main/java/org/apache/flink/streaming/connectors/rabbitmq/RMQTopology.java +++ /dev/null @@ -1,42 +0,0 @@ -/* - * 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.rabbitmq; - -import org.apache.flink.streaming.api.datastream.DataStreamSink; -import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment; -import org.apache.flink.streaming.util.serialization.SimpleStringSchema; - -public class RMQTopology { - - public static void main(String[] args) throws Exception { - - StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment(1); - - @SuppressWarnings("unused") - DataStreamSink dataStream1 = env.addSource( - new RMQSource<>("localhost", "hello", new SimpleStringSchema())).print(); - - @SuppressWarnings("unused") - DataStreamSink dataStream2 = env.fromElements("one", "two", "three", "four", "five", - "q").addSink( - new RMQSink<>("localhost", "hello", new SimpleStringSchema())); - - env.execute(); - } - -}