JmsIO { public static Read<JmsRecord> read() { return new AutoValue_JmsIO_Read.Builder<JmsRecord>() .setMaxNumRecords(Long.MAX_VALUE) .setCoder(SerializableCoder.of(JmsRecord.class)) .setMessageMapper( (MessageMapper<JmsRecord>) new MessageMapper<JmsRecord>() { @Override public JmsRecord mapMessage(Message message) throws Exception { TextMessage textMessage = (TextMessage) message; Map<String, Object> properties = new HashMap<>(); @SuppressWarnings("rawtypes") Enumeration propertyNames = textMessage.getPropertyNames(); while (propertyNames.hasMoreElements()) { String propertyName = (String) propertyNames.nextElement(); properties.put(propertyName, textMessage.getObjectProperty(propertyName)); } return new JmsRecord( textMessage.getJMSMessageID(), textMessage.getJMSTimestamp(), textMessage.getJMSCorrelationID(), textMessage.getJMSReplyTo(), textMessage.getJMSDestination(), textMessage.getJMSDeliveryMode(), textMessage.getJMSRedelivered(), textMessage.getJMSType(), textMessage.getJMSExpiration(), textMessage.getJMSPriority(), properties, textMessage.getText()); } }) .build(); } private  JmsIO(); static Read<JmsRecord> read(); static Read<T> readMessage(); static Write write();  }
@Test public void testAuthenticationRequired() { pipeline.apply(JmsIO.read().withConnectionFactory(connectionFactory).withQueue(QUEUE)); runPipelineExpectingJmsConnectException("User name [null] or password is invalid."); }
@Test public void testAuthenticationWithBadPassword() { pipeline.apply( JmsIO.read() .withConnectionFactory(connectionFactory) .withQueue(QUEUE) .withUsername(USERNAME) .withPassword("BAD")); runPipelineExpectingJmsConnectException("User name [" + USERNAME + "] or password is invalid."); }
@Test public void testReadMessages() throws Exception { 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.send(message); producer.close(); session.close(); connection.close(); PCollection<JmsRecord> output = pipeline.apply( JmsIO.read() .withConnectionFactory(connectionFactory) .withQueue(QUEUE) .withUsername(USERNAME) .withPassword(PASSWORD) .withMaxNumRecords(5)); PAssert.thatSingleton(output.apply("Count", Count.globally())).isEqualTo(5L); pipeline.run(); connection = connectionFactory.createConnection(USERNAME, PASSWORD); session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE); MessageConsumer consumer = session.createConsumer(session.createQueue(QUEUE)); Message msg = consumer.receiveNoWait(); assertNull(msg); }
@Test public void testSplitForQueue() throws Exception { JmsIO.Read read = JmsIO.read().withQueue(QUEUE); PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); int desiredNumSplits = 5; JmsIO.UnboundedJmsSource initialSource = new JmsIO.UnboundedJmsSource(read); List<JmsIO.UnboundedJmsSource> splits = initialSource.split(desiredNumSplits, pipelineOptions); assertEquals(desiredNumSplits, splits.size()); }
@Test public void testSplitForTopic() throws Exception { JmsIO.Read read = JmsIO.read().withTopic(TOPIC); PipelineOptions pipelineOptions = PipelineOptionsFactory.create(); int desiredNumSplits = 5; JmsIO.UnboundedJmsSource initialSource = new JmsIO.UnboundedJmsSource(read); List<JmsIO.UnboundedJmsSource> splits = initialSource.split(desiredNumSplits, pipelineOptions); assertEquals(1, splits.size()); }
@Test public void testCheckpointMark() throws Exception { Connection connection = connectionFactoryWithSyncAcksAndWithoutPrefetch.createConnection(USERNAME, PASSWORD); connection.start(); 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(); JmsIO.Read spec = JmsIO.read() .withConnectionFactory(connectionFactoryWithSyncAcksAndWithoutPrefetch) .withUsername(USERNAME) .withPassword(PASSWORD) .withQueue(QUEUE); JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec); JmsIO.UnboundedJmsReader reader = source.createReader(null, null); assertTrue(reader.start()); for (int i = 0; i < 3; i++) { assertTrue(reader.advance()); } assertEquals(10, count(QUEUE)); reader.getCheckpointMark().finalizeCheckpoint(); assertEquals(6, count(QUEUE)); for (int i = 0; i < 6; i++) { assertTrue(reader.advance()); } assertEquals(6, count(QUEUE)); reader.getCheckpointMark().finalizeCheckpoint(); assertEquals(0, count(QUEUE)); }
@Test public void testCheckpointMarkSafety() throws Exception { final int messagesToProcess = 100; Connection connection = connectionFactoryWithSyncAcksAndWithoutPrefetch.createConnection(USERNAME, PASSWORD); connection.start(); Session session = connection.createSession(false, Session.CLIENT_ACKNOWLEDGE); MessageProducer producer = session.createProducer(session.createQueue(QUEUE)); for (int i = 0; i < messagesToProcess; i++) { producer.send(session.createTextMessage("test " + i)); } producer.close(); session.close(); connection.close(); JmsIO.Read spec = JmsIO.read() .withConnectionFactory( withSlowAcks(connectionFactoryWithSyncAcksAndWithoutPrefetch, 10)) .withUsername(USERNAME) .withPassword(PASSWORD) .withQueue(QUEUE); JmsIO.UnboundedJmsSource source = new JmsIO.UnboundedJmsSource(spec); JmsIO.UnboundedJmsReader reader = source.createReader(null, null); assertTrue(reader.start()); for (int i = 0; i < (messagesToProcess / 2) - 1; i++) { assertTrue(reader.advance()); } assertEquals(messagesToProcess, count(QUEUE)); Thread runner = new Thread( () -> { try { for (int i = 0; i < messagesToProcess / 2; i++) { assertTrue(reader.advance()); } } catch (IOException ex) { throw new RuntimeException(ex); } }); runner.start(); reader.getCheckpointMark().finalizeCheckpoint(); runner.join(); }