You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@beam.apache.org by yh...@apache.org on 2024/06/03 14:33:35 UTC
(beam) branch master updated: add support for ConnectionFactory ProviderFn in JmsIO (#31264)
This is an automated email from the ASF dual-hosted git repository.
yhu pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new 349d7b9f9c6 add support for ConnectionFactory ProviderFn in JmsIO (#31264)
349d7b9f9c6 is described below
commit 349d7b9f9c6f641015576c9da42d7ff2db27ddec
Author: Maciej Szwaja <ms...@google.com>
AuthorDate: Mon Jun 3 16:33:28 2024 +0200
add support for ConnectionFactory ProviderFn in JmsIO (#31264)
* add connection factory supplier support to JmsIO.Read
* add connection factory supplier support to JmsIO.Write
* use extends in type param
* update the JmsIOIT test
* add DRY-ness
* use SerializableFunction instead of Supplier
* fix spotless and checkstyle
* add a hint to use the withConnectionFactoryProviderFn
* swap the parameters in CommonJms
* limit the number of tests for ConnectionFactory provider function
* add option in JmsIOIT to run the tests with a connection factory provider
* remove empty test
---
.../java/org/apache/beam/sdk/io/jms/JmsIO.java | 134 ++++++++++--
.../java/org/apache/beam/sdk/io/jms/CommonJms.java | 58 ++++--
.../java/org/apache/beam/sdk/io/jms/JmsIOIT.java | 35 +++-
.../java/org/apache/beam/sdk/io/jms/JmsIOTest.java | 226 ++++++++++++++-------
.../jms/MockNonSerializableConnectionFactory.java | 34 ++++
5 files changed, 382 insertions(+), 105 deletions(-)
diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
index 81994da93ad..1fa459a733d 100644
--- a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
+++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java
@@ -30,6 +30,7 @@ import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.NoSuchElementException;
+import java.util.Optional;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.stream.Stream;
@@ -186,12 +187,23 @@ public class JmsIO {
return new AutoValue_JmsIO_Write.Builder<EventT>().build();
}
+ public interface ConnectionFactoryContainer<T extends ConnectionFactoryContainer<T>> {
+
+ T withConnectionFactory(ConnectionFactory connectionFactory);
+
+ T withConnectionFactoryProviderFn(
+ SerializableFunction<Void, ? extends ConnectionFactory> connectionFactoryProviderFn);
+ }
+
/**
* A {@link PTransform} to read from a JMS destination. See {@link JmsIO} for more information on
* usage and configuration.
*/
@AutoValue
- public abstract static class Read<T> extends PTransform<PBegin, PCollection<T>> {
+ public abstract static class Read<T> extends PTransform<PBegin, PCollection<T>>
+ implements ConnectionFactoryContainer<Read<T>> {
+
+ private @Nullable transient ConnectionFactory connectionFactory;
/**
* NB: According to http://docs.oracle.com/javaee/1.4/api/javax/jms/ConnectionFactory.html "It
@@ -201,9 +213,27 @@ public class JmsIO {
* they can be stored in all JNDI naming contexts. In addition, it is recommended that these
* implementations follow the JavaBeansTM design patterns."
*
- * <p>So, a {@link ConnectionFactory} implementation is serializable.
+ * <p>So, a {@link ConnectionFactory} implementation should be serializable.
+ */
+ @Nullable
+ ConnectionFactory getConnectionFactory() {
+ if (connectionFactory == null) {
+ connectionFactory =
+ Optional.ofNullable(getConnectionFactoryProviderFn())
+ .map(provider -> provider.apply(null))
+ .orElse(null);
+ }
+ return connectionFactory;
+ }
+
+ /**
+ * In case the {@link ConnectionFactory} is not serializable it can be constructed separately on
+ * each worker from scratch by providing a {@link SerializableFunction} instead of a ready-made
+ * object.
*/
- abstract @Nullable ConnectionFactory getConnectionFactory();
+ @Nullable
+ abstract SerializableFunction<Void, ? extends ConnectionFactory>
+ getConnectionFactoryProviderFn();
abstract @Nullable String getQueue();
@@ -233,7 +263,8 @@ public class JmsIO {
@AutoValue.Builder
abstract static class Builder<T> {
- abstract Builder<T> setConnectionFactory(ConnectionFactory connectionFactory);
+ abstract Builder<T> setConnectionFactoryProviderFn(
+ SerializableFunction<Void, ? extends ConnectionFactory> connectionFactoryProviderFn);
abstract Builder<T> setQueue(String queue);
@@ -265,6 +296,9 @@ public class JmsIO {
/**
* Specify the JMS connection factory to connect to the JMS broker.
*
+ * <p>The {@link ConnectionFactory} object has to be serializable, if it is not consider using
+ * the {@link #withConnectionFactoryProviderFn}
+ *
* <p>For instance:
*
* <pre>{@code
@@ -275,9 +309,33 @@ public class JmsIO {
* @param connectionFactory The JMS {@link ConnectionFactory}.
* @return The corresponding {@link JmsIO.Read}.
*/
+ @Override
public Read<T> withConnectionFactory(ConnectionFactory connectionFactory) {
checkArgument(connectionFactory != null, "connectionFactory can not be null");
- return builder().setConnectionFactory(connectionFactory).build();
+ return builder().setConnectionFactoryProviderFn(__ -> connectionFactory).build();
+ }
+
+ /**
+ * Specify a JMS connection factory provider function to connect to the JMS broker. Use this
+ * method in case your {@link ConnectionFactory} objects are themselves not serializable, but
+ * you can recreate them as needed with a {@link SerializableFunction} provider
+ *
+ * <p>For instance:
+ *
+ * <pre>
+ * {@code pipeline.apply(JmsIO.read().withConnectionFactoryProviderFn(() -> new MyJmsConnectionFactory());}
+ * </pre>
+ *
+ * @param connectionFactoryProviderFn a {@link SerializableFunction} that creates a {@link
+ * ConnectionFactory}
+ * @return The corresponding {@link JmsIO.Read}
+ */
+ @Override
+ public Read<T> withConnectionFactoryProviderFn(
+ SerializableFunction<Void, ? extends ConnectionFactory> connectionFactoryProviderFn) {
+ checkArgument(
+ connectionFactoryProviderFn != null, "connectionFactoryProviderFn cannot be null");
+ return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build();
}
/**
@@ -426,7 +484,9 @@ public class JmsIO {
@Override
public PCollection<T> expand(PBegin input) {
- checkArgument(getConnectionFactory() != null, "withConnectionFactory() is required");
+ checkArgument(
+ getConnectionFactoryProviderFn() != null,
+ "Either withConnectionFactory() or withConnectionFactoryProviderFn() is required");
checkArgument(
getQueue() != null || getTopic() != null,
"Either withQueue() or withTopic() is required");
@@ -786,9 +846,25 @@ public class JmsIO {
*/
@AutoValue
public abstract static class Write<EventT>
- extends PTransform<PCollection<EventT>, WriteJmsResult<EventT>> {
+ extends PTransform<PCollection<EventT>, WriteJmsResult<EventT>>
+ implements ConnectionFactoryContainer<Write<EventT>> {
+
+ private @Nullable transient ConnectionFactory connectionFactory;
+
+ @Nullable
+ ConnectionFactory getConnectionFactory() {
+ if (connectionFactory == null) {
+ connectionFactory =
+ Optional.ofNullable(getConnectionFactoryProviderFn())
+ .map(provider -> provider.apply(null))
+ .orElse(null);
+ }
+
+ return connectionFactory;
+ }
- abstract @Nullable ConnectionFactory getConnectionFactory();
+ abstract @Nullable SerializableFunction<Void, ? extends ConnectionFactory>
+ getConnectionFactoryProviderFn();
abstract @Nullable String getQueue();
@@ -808,7 +884,8 @@ public class JmsIO {
@AutoValue.Builder
abstract static class Builder<EventT> {
- abstract Builder<EventT> setConnectionFactory(ConnectionFactory connectionFactory);
+ abstract Builder<EventT> setConnectionFactoryProviderFn(
+ SerializableFunction<Void, ? extends ConnectionFactory> connectionFactoryProviderFn);
abstract Builder<EventT> setQueue(String queue);
@@ -832,6 +909,9 @@ public class JmsIO {
/**
* Specify the JMS connection factory to connect to the JMS broker.
*
+ * <p>The {@link ConnectionFactory} object has to be serializable, if it is not consider using
+ * the {@link #withConnectionFactoryProviderFn}
+ *
* <p>For instance:
*
* <pre>{@code
@@ -842,9 +922,33 @@ public class JmsIO {
* @param connectionFactory The JMS {@link ConnectionFactory}.
* @return The corresponding {@link JmsIO.Read}.
*/
+ @Override
public Write<EventT> withConnectionFactory(ConnectionFactory connectionFactory) {
checkArgument(connectionFactory != null, "connectionFactory can not be null");
- return builder().setConnectionFactory(connectionFactory).build();
+ return builder().setConnectionFactoryProviderFn(__ -> connectionFactory).build();
+ }
+
+ /**
+ * Specify a JMS connection factory provider function to connect to the JMS broker. Use this
+ * method in case your {@link ConnectionFactory} objects are themselves not serializable, but
+ * you can recreate them as needed with a {@link SerializableFunction} provider
+ *
+ * <p>For instance:
+ *
+ * <pre>
+ * {@code pipeline.apply(JmsIO.write().withConnectionFactoryProviderFn(() -> new MyJmsConnectionFactory());}
+ * </pre>
+ *
+ * @param connectionFactoryProviderFn a {@link SerializableFunction} that creates a {@link
+ * ConnectionFactory}
+ * @return The corresponding {@link JmsIO.Write}
+ */
+ @Override
+ public Write<EventT> withConnectionFactoryProviderFn(
+ SerializableFunction<Void, ? extends ConnectionFactory> connectionFactoryProviderFn) {
+ checkArgument(
+ connectionFactoryProviderFn != null, "connectionFactoryProviderFn can not be null");
+ return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build();
}
/**
@@ -1010,14 +1114,18 @@ public class JmsIO {
@Override
public WriteJmsResult<EventT> expand(PCollection<EventT> input) {
- checkArgument(getConnectionFactory() != null, "withConnectionFactory() is required");
+ checkArgument(
+ getConnectionFactoryProviderFn() != null,
+ "Either withConnectionFactory() or withConnectionFactoryProviderFn() is required");
checkArgument(
getTopicNameMapper() != null || getQueue() != null || getTopic() != null,
- "Either withTopicNameMapper(topicNameMapper), withQueue(queue), or withTopic(topic) is required");
+ "Either withTopicNameMapper(topicNameMapper), withQueue(queue), or withTopic(topic) is"
+ + " required");
boolean exclusiveTopicQueue = isExclusiveTopicQueue();
checkArgument(
exclusiveTopicQueue,
- "Only one of withQueue(queue), withTopic(topic), or withTopicNameMapper(function) must be set.");
+ "Only one of withQueue(queue), withTopic(topic), or withTopicNameMapper(function) must be"
+ + " set.");
checkArgument(getValueMapper() != null, "withValueMapper() is required");
return input.apply(new Writer<>(this));
diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java
index 9670e60e345..c0f8cf258d2 100644
--- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java
+++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/CommonJms.java
@@ -18,6 +18,7 @@
package org.apache.beam.sdk.io.jms;
import java.io.Serializable;
+import java.lang.reflect.InvocationTargetException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
@@ -32,6 +33,9 @@ import org.apache.activemq.security.SimpleAuthenticationPlugin;
import org.apache.activemq.store.memory.MemoryPersistenceAdapter;
import org.apache.activemq.transport.TransportFactory;
import org.apache.activemq.transport.amqp.AmqpTransportFactory;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+import org.apache.beam.sdk.util.ThrowingSupplier;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Supplier;
/**
* A common test fixture to create a broker and connection factories for {@link JmsIOIT} & {@link
@@ -40,6 +44,31 @@ import org.apache.activemq.transport.amqp.AmqpTransportFactory;
public class CommonJms implements Serializable {
private static final String BROKER_WITHOUT_PREFETCH_PARAM = "?jms.prefetchPolicy.all=0&";
+ // convenient typedefs and a helper conversion functions
+ interface ThrowingSerializableSupplier<T> extends ThrowingSupplier<T>, Serializable {}
+
+ private interface SerializableSupplier<T> extends Serializable, Supplier<T> {}
+
+ private static <T> SerializableSupplier<T> toSerializableSupplier(
+ ThrowingSerializableSupplier<T> throwingSerializableSupplier) {
+ return () -> {
+ try {
+ return throwingSerializableSupplier.get();
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ };
+ }
+
+ private static <T> SerializableFunction<Void, T> toSerializableFunction(Supplier<T> supplier) {
+ return __ -> supplier.get();
+ }
+
+ static <T> SerializableFunction<Void, T> toSerializableFunction(
+ ThrowingSerializableSupplier<T> supplier) {
+ return toSerializableFunction(toSerializableSupplier(supplier));
+ }
+
static final String USERNAME = "test_user";
static final String PASSWORD = "test_password";
static final String QUEUE = "test_queue";
@@ -50,9 +79,7 @@ public class CommonJms implements Serializable {
private final String forceAsyncAcksParam;
private transient BrokerService broker;
- protected ConnectionFactory connectionFactory;
protected final Class<? extends ConnectionFactory> connectionFactoryClass;
- protected ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch;
public CommonJms(
String brokerUrl,
@@ -91,13 +118,20 @@ public class CommonJms implements Serializable {
broker.start();
broker.waitUntilStarted();
+ }
+
+ ConnectionFactory createConnectionFactory()
+ throws NoSuchMethodException, InvocationTargetException, InstantiationException,
+ IllegalAccessException {
+ return connectionFactoryClass.getConstructor(String.class).newInstance(brokerUrl);
+ }
- // create JMS connection factory
- connectionFactory = connectionFactoryClass.getConstructor(String.class).newInstance(brokerUrl);
- connectionFactoryWithSyncAcksAndWithoutPrefetch =
- connectionFactoryClass
- .getConstructor(String.class)
- .newInstance(brokerUrl + BROKER_WITHOUT_PREFETCH_PARAM + forceAsyncAcksParam);
+ ConnectionFactory createConnectionFactoryWithSyncAcksAndWithoutPrefetch()
+ throws NoSuchMethodException, InvocationTargetException, InstantiationException,
+ IllegalAccessException {
+ return connectionFactoryClass
+ .getConstructor(String.class)
+ .newInstance(brokerUrl + BROKER_WITHOUT_PREFETCH_PARAM + forceAsyncAcksParam);
}
void stopBroker() throws Exception {
@@ -110,14 +144,6 @@ public class CommonJms implements Serializable {
return this.connectionFactoryClass;
}
- ConnectionFactory getConnectionFactory() {
- return this.connectionFactory;
- }
-
- ConnectionFactory getConnectionFactoryWithSyncAcksAndWithoutPrefetch() {
- return this.connectionFactoryWithSyncAcksAndWithoutPrefetch;
- }
-
/** A test class that maps a {@link javax.jms.BytesMessage} into a {@link String}. */
public static class BytesMessageToStringMessageMapper implements JmsIO.MessageMapper<String> {
diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java
index 6f5a00760e3..eddcb0de556 100644
--- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java
+++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOIT.java
@@ -120,6 +120,13 @@ public class JmsIOIT implements Serializable {
Integer getReadTimeout();
void setReadTimeout(Integer timeout);
+
+ @Description(
+ "Use ConnectionFactory provider function instead of pre-instantiated ConnectionFactory object")
+ @Default.Boolean(false)
+ boolean getUseConnectionFactoryProviderFn();
+
+ void setUseConnectionFactoryProviderFn(boolean value);
}
private static final JmsIOITOptions OPTIONS =
@@ -169,7 +176,7 @@ public class JmsIOIT implements Serializable {
public void setup() throws Exception {
if (OPTIONS.isLocalJmsBrokerEnabled()) {
this.commonJms.startBroker();
- connectionFactory = this.commonJms.getConnectionFactory();
+ connectionFactory = this.commonJms.createConnectionFactory();
connectionFactoryClass = this.commonJms.getConnectionFactoryClass();
// use a small number of record for local integration test
OPTIONS.setNumberOfRecords(10000);
@@ -229,15 +236,22 @@ public class JmsIOIT implements Serializable {
private PipelineResult readMessages() {
pipelineRead.getOptions().as(JmsIOITOptions.class).setStreaming(true);
pipelineRead.getOptions().as(JmsIOITOptions.class).setBlockOnRun(false);
+ JmsIO.Read<String> jmsIORead = JmsIO.readMessage();
+ if (pipelineRead.getOptions().as(JmsIOITOptions.class).getUseConnectionFactoryProviderFn()) {
+ jmsIORead =
+ jmsIORead.withConnectionFactoryProviderFn(
+ CommonJms.toSerializableFunction(commonJms::createConnectionFactory));
+ } else {
+ jmsIORead = jmsIORead.withConnectionFactory(connectionFactory);
+ }
pipelineRead
.apply(
"Read Messages",
- JmsIO.<String>readMessage()
+ jmsIORead
.withQueue(QUEUE)
.withUsername(USERNAME)
.withPassword(PASSWORD)
.withCoder(SerializableCoder.of(String.class))
- .withConnectionFactory(connectionFactory)
.withMessageMapper(getJmsMessageMapper()))
.apply(ParDo.of(new TimeMonitor<>(NAMESPACE, READ_TIME_METRIC)))
.apply("Counting element", ParDo.of(new CountingFn(NAMESPACE, READ_ELEMENT_METRIC_NAME)));
@@ -245,18 +259,27 @@ public class JmsIOIT implements Serializable {
}
private PipelineResult publishingMessages() {
+
+ JmsIO.Write<String> jmsIOWrite = JmsIO.write();
+ if (pipelineWrite.getOptions().as(JmsIOITOptions.class).getUseConnectionFactoryProviderFn()) {
+ jmsIOWrite =
+ jmsIOWrite.withConnectionFactoryProviderFn(
+ CommonJms.toSerializableFunction(commonJms::createConnectionFactory));
+ } else {
+ jmsIOWrite = jmsIOWrite.withConnectionFactory(connectionFactory);
+ }
+
pipelineWrite
.apply("Generate Sequence Data", GenerateSequence.from(0).to(OPTIONS.getNumberOfRecords()))
.apply("Convert to String", ParDo.of(new ToString()))
.apply("Collect write time", ParDo.of(new TimeMonitor<>(NAMESPACE, WRITE_TIME_METRIC)))
.apply(
"Publish to Jms Broker",
- JmsIO.<String>write()
+ jmsIOWrite
.withQueue(QUEUE)
.withUsername(USERNAME)
.withPassword(PASSWORD)
- .withValueMapper(new TextMessageMapper())
- .withConnectionFactory(connectionFactory));
+ .withValueMapper(new TextMessageMapper()));
return pipelineWrite.run();
}
diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
index 4e967276746..6fe65520873 100644
--- a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
+++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsIOTest.java
@@ -22,6 +22,7 @@ import static org.apache.beam.sdk.io.jms.CommonJms.PASSWORD;
import static org.apache.beam.sdk.io.jms.CommonJms.QUEUE;
import static org.apache.beam.sdk.io.jms.CommonJms.TOPIC;
import static org.apache.beam.sdk.io.jms.CommonJms.USERNAME;
+import static org.apache.beam.sdk.io.jms.CommonJms.toSerializableFunction;
import static org.apache.beam.sdk.io.jms.JmsIO.Writer.JMS_IO_PRODUCER_METRIC_NAME;
import static org.apache.beam.sdk.io.jms.JmsIO.Writer.PUBLICATION_RETRIES_METRIC_NAME;
import static org.hamcrest.CoreMatchers.allOf;
@@ -30,6 +31,7 @@ import static org.hamcrest.Matchers.contains;
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
import static org.hamcrest.Matchers.hasProperty;
import static org.hamcrest.Matchers.is;
+import static org.hamcrest.Matchers.isA;
import static org.hamcrest.Matchers.lessThan;
import static org.hamcrest.core.StringContains.containsString;
import static org.hamcrest.object.HasToString.hasToString;
@@ -49,7 +51,9 @@ import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import java.io.IOException;
+import java.io.NotSerializableException;
import java.io.Serializable;
+import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.nio.ByteBuffer;
import java.nio.charset.StandardCharsets;
@@ -81,6 +85,7 @@ import org.apache.activemq.util.Callback;
import org.apache.beam.sdk.PipelineResult;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.SerializableCoder;
+import org.apache.beam.sdk.coders.StringUtf8Coder;
import org.apache.beam.sdk.io.UnboundedSource.CheckpointMark;
import org.apache.beam.sdk.io.jms.JmsIO.UnboundedJmsReader;
import org.apache.beam.sdk.metrics.MetricNameFilter;
@@ -95,6 +100,7 @@ import org.apache.beam.sdk.testing.TestPipeline;
import org.apache.beam.sdk.transforms.Count;
import org.apache.beam.sdk.transforms.Create;
import org.apache.beam.sdk.transforms.SerializableBiFunction;
+import org.apache.beam.sdk.util.SerializableUtils;
import org.apache.beam.sdk.values.PCollection;
import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Throwables;
import org.apache.qpid.jms.JmsAcknowledgeCallback;
@@ -120,6 +126,24 @@ import org.slf4j.LoggerFactory;
public class JmsIOTest {
private static final Logger LOG = LoggerFactory.getLogger(JmsIOTest.class);
+
+ private interface TestMessageCreator<T extends Message> {
+ T createMessage(Session session, long messageIndex) throws JMSException;
+ }
+
+ private static BytesMessage createBytesMessage(Session session, long messageIndex)
+ throws JMSException {
+ BytesMessage message = session.createBytesMessage();
+ message.writeBytes(
+ String.format("This Is A Test %d", messageIndex).getBytes(StandardCharsets.UTF_8));
+ return message;
+ }
+
+ private static TextMessage createTextMessage(Session session, long messageIndex)
+ throws JMSException {
+ return session.createTextMessage(String.format("This Is A Test %d", messageIndex));
+ }
+
private final RetryConfiguration retryConfiguration =
RetryConfiguration.create(1, Duration.standardSeconds(1), null);
@Rule public final transient TestPipeline pipeline = TestPipeline.create();
@@ -136,34 +160,33 @@ public class JmsIOTest {
}
private final CommonJms commonJms;
- private ConnectionFactory connectionFactory;
- private Class<? extends ConnectionFactory> connectionFactoryClass;
- private ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch;
+ private final ConnectionFactory connectionFactory;
+ private final Class<? extends ConnectionFactory> connectionFactoryClass;
+ private final ConnectionFactory connectionFactoryWithSyncAcksAndWithoutPrefetch;
public JmsIOTest(
String brokerUrl,
Integer brokerPort,
String forceAsyncAcksParam,
- Class<? extends ConnectionFactory> connectionFactoryClass) {
+ Class<? extends ConnectionFactory> connectionFactoryClass)
+ throws InvocationTargetException, NoSuchMethodException, InstantiationException,
+ IllegalAccessException {
this.commonJms =
new CommonJms(brokerUrl, brokerPort, forceAsyncAcksParam, connectionFactoryClass);
+ this.connectionFactoryClass = connectionFactoryClass;
+ this.connectionFactory = commonJms.createConnectionFactory();
+ this.connectionFactoryWithSyncAcksAndWithoutPrefetch =
+ commonJms.createConnectionFactoryWithSyncAcksAndWithoutPrefetch();
}
@Before
- public void beforeEeach() throws Exception {
+ public void beforeEach() throws Exception {
this.commonJms.startBroker();
- connectionFactory = this.commonJms.getConnectionFactory();
- connectionFactoryClass = this.commonJms.getConnectionFactoryClass();
- connectionFactoryWithSyncAcksAndWithoutPrefetch =
- this.commonJms.getConnectionFactoryWithSyncAcksAndWithoutPrefetch();
}
@After
public void tearDown() throws Exception {
this.commonJms.stopBroker();
- connectionFactory = null;
- connectionFactoryClass = null;
- connectionFactoryWithSyncAcksAndWithoutPrefetch = null;
}
private void runPipelineExpectingJmsConnectException(String innerMessage) {
@@ -203,20 +226,8 @@ public class JmsIOTest {
@Test
public void testReadMessages() throws Exception {
-
- // produce message
- Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
- Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
- MessageProducer producer = session.createProducer(session.createQueue(QUEUE));
- TextMessage message = session.createTextMessage("This Is A Test");
- producer.send(message);
- producer.send(message);
- producer.send(message);
- producer.send(message);
- producer.send(message);
- producer.close();
- session.close();
- connection.close();
+ long count = 5;
+ produceTestMessages(count, JmsIOTest::createTextMessage);
// read from the queue
PCollection<JmsRecord> output =
@@ -226,15 +237,73 @@ public class JmsIOTest {
.withQueue(QUEUE)
.withUsername(USERNAME)
.withPassword(PASSWORD)
- .withMaxNumRecords(5));
+ .withMaxNumRecords(count));
- PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(5L);
+ PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(count);
pipeline.run();
assertQueueIsEmpty();
}
- private void assertQueueIsEmpty() throws JMSException {
+ @Test
+ public void testPipelineWithNonSerializableCF() {
+ SerializableUtils.ensureSerializable(
+ JmsIO.read()
+ .withConnectionFactoryProviderFn(__ -> new MockNonSerializableConnectionFactory()));
+ try {
+ SerializableUtils.ensureSerializable(
+ JmsIO.read().withConnectionFactory(new MockNonSerializableConnectionFactory()));
+ fail();
+ } catch (Exception e) {
+ assertThat(Throwables.getRootCause(e), isA(NotSerializableException.class));
+ }
+ }
+
+ @Test
+ public void testReadMessagesWithCFProviderFn() throws Exception {
+ long count = 5;
+ produceTestMessages(count, JmsIOTest::createTextMessage);
+
+ PCollection<JmsRecord> output =
+ pipeline.apply(
+ JmsIO.read()
+ .withConnectionFactoryProviderFn(
+ toSerializableFunction(commonJms::createConnectionFactory))
+ .withQueue(QUEUE)
+ .withUsername(USERNAME)
+ .withPassword(PASSWORD)
+ .withMaxNumRecords(count));
+
+ PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(count);
+ pipeline.run();
+
+ assertQueueIsEmpty();
+ }
+
+ @Test
+ public void testReadBytesMessagesWithCFProviderFn() throws Exception {
+ long count = 5;
+ produceTestMessages(count, JmsIOTest::createBytesMessage);
+
+ PCollection<String> output =
+ pipeline.apply(
+ JmsIO.<String>readMessage()
+ .withConnectionFactoryProviderFn(
+ toSerializableFunction(commonJms::createConnectionFactory))
+ .withQueue(QUEUE)
+ .withUsername(USERNAME)
+ .withPassword(PASSWORD)
+ .withMaxNumRecords(count)
+ .withMessageMapper(new CommonJms.BytesMessageToStringMessageMapper())
+ .withCoder(StringUtf8Coder.of()));
+
+ PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(count);
+ pipeline.run();
+
+ assertQueueIsEmpty();
+ }
+
+ private void assertQueueContainsMessages(int expectedCount) throws JMSException {
// we need to disable the prefetch, otherwise the receiveNoWait could still return null even
// when there are messages waiting on the queue
Connection connection =
@@ -243,26 +312,43 @@ public class JmsIOTest {
Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
MessageConsumer consumer = session.createConsumer(session.createQueue(QUEUE));
connection.start();
- assertNull(consumer.receiveNoWait());
+ int actualCount = 0;
+ while (actualCount < expectedCount && consumer.receive(1000L) != null) {
+ actualCount += 1;
+ }
+ assertEquals("queue didn't have expected number of messages", expectedCount, actualCount);
+ assertNull(
+ String.format("verifying queue is empty after consuming %d messages - ", expectedCount),
+ consumer.receiveNoWait());
} finally {
connection.close();
}
}
- @Test
- public void testReadBytesMessages() throws Exception {
+ private void assertQueueIsEmpty() throws JMSException {
+ assertQueueContainsMessages(0);
+ }
- // produce message
+ private void produceTestMessages(long count, TestMessageCreator testMessageCreator)
+ throws JMSException {
Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
- Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
- MessageProducer producer = session.createProducer(session.createQueue(QUEUE));
- BytesMessage message = session.createBytesMessage();
- message.writeBytes("This Is A Test".getBytes(StandardCharsets.UTF_8));
- producer.send(message);
- producer.close();
- session.close();
- connection.close();
+ try {
+ Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
+ MessageProducer producer = session.createProducer(session.createQueue(QUEUE));
+ for (long i = 0; i < count; i += 1) {
+ Message message = testMessageCreator.createMessage(session, i);
+ producer.send(message);
+ }
+ } finally {
+ connection.close();
+ }
+ }
+ @Test
+ public void testReadBytesMessages() throws Exception {
+ long count = 1L;
+ // produce message
+ produceTestMessages(count, JmsIOTest::createBytesMessage);
// read from the queue
PCollection<String> output =
pipeline.apply(
@@ -275,7 +361,7 @@ public class JmsIOTest {
.withCoder(SerializableCoder.of(String.class))
.withMessageMapper(new CommonJms.BytesMessageToStringMessageMapper()));
- PAssert.thatSingleton(output.apply("Count", Count.<String>globally())).isEqualTo(1L);
+ PAssert.thatSingleton(output.apply("Count", Count.<String>globally())).isEqualTo(count);
pipeline.run();
assertQueueIsEmpty();
@@ -283,16 +369,8 @@ public class JmsIOTest {
@Test
public void testRequiresDedup() throws Exception {
- // produce message
- Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
- Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
- MessageProducer producer = session.createProducer(session.createQueue(QUEUE));
- for (int i = 0; i < 10; i++) {
- producer.send(session.createTextMessage("test " + i));
- }
- producer.close();
- session.close();
- connection.close();
+ int count = 10;
+ produceTestMessages(count, JmsIOTest::createTextMessage);
JmsIO.Read spec =
JmsIO.read()
@@ -308,7 +386,7 @@ public class JmsIOTest {
Set<ByteBuffer> uniqueIds = new HashSet<>();
uniqueIds.add(ByteBuffer.wrap(reader.getCurrentRecordId()));
- for (int i = 0; i < 10; i++) {
+ for (int i = 0; i < count; i++) {
if (reader.advance()) {
assertTrue(uniqueIds.add(ByteBuffer.wrap(reader.getCurrentRecordId())));
}
@@ -335,15 +413,31 @@ public class JmsIOTest {
pipeline.run();
- Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
- connection.start();
- Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
- MessageConsumer consumer = session.createConsumer(session.createQueue(QUEUE));
- int count = 0;
- while (consumer.receive(1000) != null) {
- count++;
+ assertQueueContainsMessages(100);
+ }
+
+ @Test
+ public void testWriteMessageWithCFProviderFn() throws Exception {
+
+ ArrayList<String> data = new ArrayList<>();
+ for (int i = 0; i < 100; i++) {
+ data.add("Message " + i);
}
- assertEquals(100, count);
+ pipeline
+ .apply(Create.of(data))
+ .apply(
+ JmsIO.<String>write()
+ .withConnectionFactoryProviderFn(
+ toSerializableFunction(commonJms::createConnectionFactory))
+ .withValueMapper(new TextMessageMapper())
+ .withRetryConfiguration(retryConfiguration)
+ .withQueue(QUEUE)
+ .withUsername(USERNAME)
+ .withPassword(PASSWORD));
+
+ pipeline.run();
+
+ assertQueueContainsMessages(100);
}
@Test
@@ -369,15 +463,7 @@ public class JmsIOTest {
pipeline.run();
- Connection connection = connectionFactory.createConnection(USERNAME, PASSWORD);
- connection.start();
- Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
- MessageConsumer consumer = session.createConsumer(session.createQueue(QUEUE));
- int count = 0;
- while (consumer.receive(1000) != null) {
- count++;
- }
- assertEquals(98, count);
+ assertQueueContainsMessages(98);
}
@Test
diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/MockNonSerializableConnectionFactory.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/MockNonSerializableConnectionFactory.java
new file mode 100644
index 00000000000..60ab20d3ef1
--- /dev/null
+++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/MockNonSerializableConnectionFactory.java
@@ -0,0 +1,34 @@
+/*
+ * 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.beam.sdk.io.jms;
+
+import javax.jms.Connection;
+import javax.jms.ConnectionFactory;
+import javax.jms.JMSException;
+
+public class MockNonSerializableConnectionFactory implements ConnectionFactory {
+ @Override
+ public Connection createConnection() throws JMSException {
+ return null;
+ }
+
+ @Override
+ public Connection createConnection(String userName, String password) throws JMSException {
+ return null;
+ }
+}