diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json index e0266d62f2e0..f1ba03a243ee 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Messaging_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } diff --git a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go index 3949d1af3248..2b502c679db9 100644 --- a/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go +++ b/sdks/go/pkg/beam/runners/prism/internal/engine/elementmanager.go @@ -922,6 +922,7 @@ func (em *ElementManager) PersistBundle(rb RunBundle, col2Coders map[string]PCol // ProcessingTime timers need to be scheduled into the processing time based queue. newHolds, ptRefreshes := em.triageTimers(d, inputInfo, stage) + // TODO(https://github.com/apache/beam/issues/39446) // Return unprocessed to this stage's pending // TODO sort out pending element watermark holds for process continuation residuals. unprocessedElements := reElementResiduals(residuals.Data, inputInfo, rb) diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/BeamGenericJmsConnectionFactory.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/BeamGenericJmsConnectionFactory.java new file mode 100644 index 000000000000..aee25273929e --- /dev/null +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/BeamGenericJmsConnectionFactory.java @@ -0,0 +1,40 @@ +/* + * 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 java.io.Serializable; +import javax.jms.ConnectionFactory; + +/** + * An interface for creating custom JMS {@link ConnectionFactory} instances. + * + *

Expansion service users connecting to JMS brokers other than the built-in supported ones + * (ActiveMQ, Qpid, IBM MQ) can implement this interface and specify their implementation class name + * in {@link ConnectionConfiguration}. + */ +@FunctionalInterface +public interface BeamGenericJmsConnectionFactory extends Serializable { + + /** + * Creates a {@link ConnectionFactory} using the given {@link ConnectionConfiguration}. + * + * @param config the JMS connection configuration + * @return configured JMS {@link ConnectionFactory} + */ + ConnectionFactory createConnectionFactory(ConnectionConfiguration config) throws Exception; +} diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/ConnectionConfiguration.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/ConnectionConfiguration.java new file mode 100644 index 000000000000..2a3af9a596c6 --- /dev/null +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/ConnectionConfiguration.java @@ -0,0 +1,280 @@ +/* + * 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 static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkArgument; + +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import java.lang.reflect.Method; +import java.util.List; +import javax.jms.ConnectionFactory; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Splitter; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** A POJO describing a JMS connection, used by SchemaTransformProvider. */ +@DefaultSchema(AutoValueSchema.class) +@AutoValue +public abstract class ConnectionConfiguration implements Serializable { + private static final Logger LOG = LoggerFactory.getLogger(ConnectionConfiguration.class); + + public static Builder builder() { + return new AutoValue_ConnectionConfiguration.Builder(); + } + + public static ConnectionConfiguration create( + String serverUri, @Nullable String connectionFactoryClassName) { + checkArgument(serverUri != null, "serverUri can not be null"); + return builder() + .setServerUri(serverUri) + .setConnectionFactoryClassName(connectionFactoryClassName) + .build(); + } + + public static ConnectionConfiguration create(String serverUri) { + return create(serverUri, null); + } + + @SchemaFieldDescription("The JMS broker URI.") + public abstract String getServerUri(); + + @SchemaFieldDescription("The JMS ConnectionFactory class name.") + public abstract @Nullable String getConnectionFactoryClassName(); + + @SchemaFieldDescription("The username to connect to the JMS broker.") + public abstract @Nullable String getUsername(); + + @SchemaFieldDescription("The password to connect to the JMS broker.") + public abstract @Nullable String getPassword(); + + public ConnectionConfiguration withUsername(String username) { + return toBuilder().setUsername(username).build(); + } + + public ConnectionConfiguration withPassword(String password) { + return toBuilder().setPassword(password).build(); + } + + public ConnectionConfiguration withConnectionFactoryClassName(String connectionFactoryClassName) { + return toBuilder().setConnectionFactoryClassName(connectionFactoryClassName).build(); + } + + abstract Builder toBuilder(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setServerUri(String serverUri); + + public abstract Builder setConnectionFactoryClassName( + @Nullable String connectionFactoryClassName); + + public abstract Builder setUsername(@Nullable String username); + + public abstract Builder setPassword(@Nullable String password); + + public abstract ConnectionConfiguration build(); + } + + public ConnectionFactory createConnectionFactory() { + String className = getConnectionFactoryClassName(); + // Default to ActiveMQ + if (className == null || className.isEmpty()) { + className = "org.apache.activemq.ActiveMQConnectionFactory"; + } + try { + Class clazz = Class.forName(className); + String uri = getServerUri(); + String username = getUsername(); + String password = getPassword(); + + if (className.contains("org.apache.activemq.ActiveMQConnectionFactory")) { + return createActiveMqConnectionFactory(clazz, uri, username, password); + } else if (className.contains("org.apache.qpid.jms")) { + return createQpidConnectionFactory(clazz, uri, username, password); + } else if (className.contains("com.ibm.mq")) { + return createIbmMqConnectionFactory(clazz, uri, username, password); + } else if (BeamGenericJmsConnectionFactory.class.isAssignableFrom(clazz)) { + BeamGenericJmsConnectionFactory factory = + (BeamGenericJmsConnectionFactory) clazz.getDeclaredConstructor().newInstance(); + return factory.createConnectionFactory(this); + } else { + try { + return createStandardConnectionFactory(clazz, uri, username, password); + } catch (Exception e) { + throw new IllegalArgumentException( + "Unable to instantiate JMS ConnectionFactory of class " + + className + + ". Must be a supported provider (ActiveMQ, Qpid, IBM MQ) or implement BeamGenericJmsConnectionFactory.", + e); + } + } + } catch (IllegalArgumentException e) { + throw e; + } catch (Exception e) { + throw new IllegalArgumentException( + "Unable to instantiate JMS ConnectionFactory of class " + className, e); + } + } + + private ConnectionFactory createActiveMqConnectionFactory( + Class clazz, String uri, @Nullable String username, @Nullable String password) + throws Exception { + return createStandardConnectionFactory(clazz, uri, username, password); + } + + private ConnectionFactory createQpidConnectionFactory( + Class clazz, String uri, @Nullable String username, @Nullable String password) + throws Exception { + return createStandardConnectionFactory(clazz, uri, username, password); + } + + private ConnectionFactory createStandardConnectionFactory( + Class clazz, String uri, @Nullable String username, @Nullable String password) + throws Exception { + if (username != null && password != null) { + try { + return (ConnectionFactory) + clazz + .getConstructor(String.class, String.class, String.class) + .newInstance(username, password, uri); + } catch (NoSuchMethodException e) { + // Fall through to 1-arg or 0-arg constructor + setters + } + } + ConnectionFactory cf; + try { + cf = (ConnectionFactory) clazz.getConstructor(String.class).newInstance(uri); + } catch (NoSuchMethodException e) { + cf = (ConnectionFactory) clazz.getConstructor().newInstance(); + } + + if (username != null && password != null) { + boolean setUsernameSuccess = + // ActiveMQ (capital N) + invokeMethodIfExists(cf, "setUserName", String.class, username) + // Qpid (lowercase n) + || invokeMethodIfExists(cf, "setUsername", String.class, username); + boolean setPasswordSuccess = invokeMethodIfExists(cf, "setPassword", String.class, password); + + if (!setUsernameSuccess || !setPasswordSuccess) { + LOG.warn("Unable to set username/password on JMS ConnectionFactory of class {}", clazz); + } + } + return cf; + } + + /** + * Instantiates and configures an IBM MQ {@link ConnectionFactory} for TCP client mode. + * + *

Sets transport type to client mode ({@code WMQ_CM_CLIENT = 1}), extracts host, port, + * channel, and queueManager parameters from the connection URI, and sets username/password if + * provided. + */ + private ConnectionFactory createIbmMqConnectionFactory( + Class clazz, String uri, @Nullable String username, @Nullable String password) + throws Exception { + ConnectionFactory cf = (ConnectionFactory) clazz.getConstructor().newInstance(); + // WMQ_CM_CLIENT = 1 (Network TCP connection mode) + invokeMethodIfExists(cf, "setTransportType", int.class, 1); + + if (!Strings.isNullOrEmpty(uri)) { + java.net.URI parsedUri = new java.net.URI(uri); + String host = parsedUri.getHost(); + int port = parsedUri.getPort(); + if (host != null) { + invokeMethodIfExists(cf, "setHostName", String.class, host); + } + if (port > 0) { + invokeMethodIfExists(cf, "setPort", int.class, port); + } + if (parsedUri.getQuery() != null) { + for (String param : Splitter.on('&').split(parsedUri.getQuery())) { + List pair = Splitter.on('=').splitToList(param); + if (pair.size() == 2) { + if ("channel".equalsIgnoreCase(pair.get(0))) { + invokeMethodIfExists(cf, "setChannel", String.class, pair.get(1)); + } else if ("queueManager".equalsIgnoreCase(pair.get(0))) { + invokeMethodIfExists(cf, "setQueueManager", String.class, pair.get(1)); + } + } + } + } + } + + if (username != null) { + // IBM MQ requires USER_AUTHENTICATION_MQCSP (XMSC_USER_AUTHENTICATION_MQCSP) set to true for + // username/password authentication + invokeTwoArgMethodIfExists( + cf, + "setBooleanProperty", + String.class, + "XMSC_USER_AUTHENTICATION_MQCSP", + boolean.class, + true); + + boolean setUsernameSuccess = + invokeTwoArgMethodIfExists( + cf, "setStringProperty", String.class, "XMSC_USERID", String.class, username); + + boolean setPasswordSuccess = true; + if (password != null) { + setPasswordSuccess = + invokeTwoArgMethodIfExists( + cf, "setStringProperty", String.class, "XMSC_PASSWORD", String.class, password); + } + + if (!setUsernameSuccess || !setPasswordSuccess) { + LOG.warn("Unable to set username/password on IBM MQ ConnectionFactory of class {}", clazz); + } + } + return cf; + } + + private static boolean invokeMethodIfExists( + Object target, String methodName, Class paramType, Object arg) { + try { + Method m = target.getClass().getMethod(methodName, paramType); + m.invoke(target, arg); + return true; + } catch (Exception e) { + return false; + } + } + + private static boolean invokeTwoArgMethodIfExists( + Object target, + String methodName, + Class p1Type, + Object arg1, + Class p2Type, + Object arg2) { + try { + Method m = target.getClass().getMethod(methodName, p1Type, p2Type); + m.invoke(target, arg1, arg2); + return true; + } catch (Exception e) { + return false; + } + } +} 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 8926d584a133..64c35a22f3ce 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 @@ -208,6 +208,21 @@ public static Write write() { return new AutoValue_JmsIO_Write.Builder().build(); } + /** @deprecated Use {@link org.apache.beam.sdk.io.jms.ConnectionConfiguration}. */ + @Deprecated + public abstract static class ConnectionConfiguration + extends org.apache.beam.sdk.io.jms.ConnectionConfiguration { + public static org.apache.beam.sdk.io.jms.ConnectionConfiguration create( + String serverUri, @Nullable String connectionFactoryClassName) { + return org.apache.beam.sdk.io.jms.ConnectionConfiguration.create( + serverUri, connectionFactoryClassName); + } + + public static org.apache.beam.sdk.io.jms.ConnectionConfiguration create(String serverUri) { + return org.apache.beam.sdk.io.jms.ConnectionConfiguration.create(serverUri); + } + } + public interface ConnectionFactoryContainer> { T withConnectionFactory(ConnectionFactory connectionFactory); @@ -367,6 +382,22 @@ public Read withConnectionFactoryProviderFn( return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build(); } + public Read withConnectionConfiguration( + org.apache.beam.sdk.io.jms.ConnectionConfiguration configuration) { + checkArgument(configuration != null, "configuration can not be null"); + Read read = + this.withConnectionFactoryProviderFn( + (SerializableFunction) + __ -> configuration.createConnectionFactory()); + if (configuration.getUsername() != null) { + read = read.withUsername(configuration.getUsername()); + } + if (configuration.getPassword() != null) { + read = read.withPassword(configuration.getPassword()); + } + return read; + } + /** * Specify the JMS queue destination name where to read messages from. The {@link JmsIO.Read} * acts as a consumer on the queue. @@ -1106,6 +1137,22 @@ public Write withConnectionFactoryProviderFn( return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build(); } + public Write withConnectionConfiguration( + org.apache.beam.sdk.io.jms.ConnectionConfiguration configuration) { + checkArgument(configuration != null, "configuration can not be null"); + Write write = + this.withConnectionFactoryProviderFn( + (SerializableFunction) + __ -> configuration.createConnectionFactory()); + if (configuration.getUsername() != null) { + write = write.withUsername(configuration.getUsername()); + } + if (configuration.getPassword() != null) { + write = write.withPassword(configuration.getPassword()); + } + return write; + } + /** * Specify the JMS queue destination name where to send messages to. The {@link JmsIO.Write} * acts as a producer on the queue. diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsReadSchemaTransformProvider.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsReadSchemaTransformProvider.java new file mode 100644 index 000000000000..c34f38f58e05 --- /dev/null +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsReadSchemaTransformProvider.java @@ -0,0 +1,227 @@ +/* + * 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 static org.apache.beam.sdk.io.jms.JmsReadSchemaTransformProvider.ReadConfiguration; + +import com.google.auto.service.AutoService; +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; +import org.apache.beam.sdk.schemas.transforms.SchemaTransform; +import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; +import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; +import org.joda.time.Duration; + +@AutoService(SchemaTransformProvider.class) +public class JmsReadSchemaTransformProvider + extends TypedSchemaTransformProvider { + @DefaultSchema(AutoValueSchema.class) + @AutoValue + public abstract static class ReadConfiguration implements Serializable { + public static Builder builder() { + return new AutoValue_JmsReadSchemaTransformProvider_ReadConfiguration.Builder(); + } + + @SchemaFieldDescription( + "Configuration options to set up the JMS connection.\nNote: if " + + "connection factory class name is set, spin up a persistent expansion service with " + + "the provider client JAR on the classpath:\n" + + "java -cp :\n" + + "org.apache.beam.sdk.expansion.service.ExpansionService \n" + + "and pass expansion_service='localhost:' to the transform. Currently, only " + + "ActiveMQ (org.apache.activemq.ActiveMQConnectionFactory) is embedded into the expansion service") + public abstract ConnectionConfiguration getConnectionConfiguration(); + + @SchemaFieldDescription( + "The JMS queue to read from. Exclusively one of queue or topic must be specified.") + @Nullable + public abstract String getQueue(); + + @SchemaFieldDescription( + "The JMS topic to read from. Exclusively one of queue or topic must be specified.") + @Nullable + public abstract String getTopic(); + + @SchemaFieldDescription( + "The max number of records to receive. Setting this will result in a bounded PCollection.") + @Nullable + public abstract Long getMaxNumRecords(); + + @SchemaFieldDescription( + "The maximum time for this source to read messages. Setting this will result in a bounded PCollection.") + @Nullable + public abstract Long getMaxReadTimeSeconds(); + + @SchemaFieldDescription("Close timeout for the JMS connection in seconds.") + @Nullable + public abstract Long getCloseTimeoutSeconds(); + + @SchemaFieldDescription( + "The JMS acknowledge mode: CLIENT_ACKNOWLEDGE, CLIENT_ACKNOWLEDGE_UNSAFE, or INDIVIDUAL_ACKNOWLEDGE.") + @Nullable + public abstract String getAcknowledgeMode(); + + @SchemaFieldDescription( + "The proprietary integer code for individual message acknowledgment when using INDIVIDUAL_ACKNOWLEDGE.") + @Nullable + public abstract Integer getIndividualAcknowledgeModeCode(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setConnectionConfiguration( + ConnectionConfiguration connectionConfiguration); + + public abstract Builder setQueue(String queue); + + public abstract Builder setTopic(String topic); + + public abstract Builder setMaxNumRecords(Long maxNumRecords); + + public abstract Builder setMaxReadTimeSeconds(Long maxReadTimeSeconds); + + public abstract Builder setCloseTimeoutSeconds(Long closeTimeoutSeconds); + + public abstract Builder setAcknowledgeMode(String acknowledgeMode); + + public abstract Builder setIndividualAcknowledgeModeCode( + Integer individualAcknowledgeModeCode); + + public abstract ReadConfiguration build(); + } + } + + @Override + public String identifier() { + return "beam:schematransform:org.apache.beam:jms_read:v1"; + } + + @Override + public String description() { + return "Reads messages from a JMS broker and outputs each message as a single `payload` " + + "(string) field.\n" + + "\n" + + "By default the read is unbounded (streaming): it keeps consuming messages from the " + + "specified queue or topic until the pipeline is stopped. Setting `maxNumRecords` and/or " + + "`maxReadTimeSeconds` bounds the read, producing a bounded (batch) PCollection."; + } + + @Override + public List outputCollectionNames() { + return Collections.singletonList("output"); + } + + @Override + protected SchemaTransform from(ReadConfiguration configuration) { + return new JmsReadSchemaTransform(configuration); + } + + private static class JmsReadSchemaTransform extends SchemaTransform { + private final ReadConfiguration config; + + JmsReadSchemaTransform(ReadConfiguration configuration) { + this.config = configuration; + } + + @Override + public PCollectionRowTuple expand(PCollectionRowTuple input) { + Preconditions.checkState( + input.getAll().isEmpty(), + "Expected zero input PCollections for this source, but found: %s", + input.getAll().keySet()); + + Preconditions.checkArgument( + (config.getQueue() != null) != (config.getTopic() != null), + "Exactly one of queue or topic must be specified."); + + JmsIO.Read readTransform = + JmsIO.read().withConnectionConfiguration(config.getConnectionConfiguration()); + + String queue = config.getQueue(); + if (queue != null) { + readTransform = readTransform.withQueue(queue); + } + String topic = config.getTopic(); + if (topic != null) { + readTransform = readTransform.withTopic(topic); + } + + Long maxRecords = config.getMaxNumRecords(); + Long maxReadTime = config.getMaxReadTimeSeconds(); + if (maxRecords != null) { + readTransform = readTransform.withMaxNumRecords(maxRecords); + } + if (maxReadTime != null) { + readTransform = readTransform.withMaxReadTime(Duration.standardSeconds(maxReadTime)); + } + + Long closeTimeout = config.getCloseTimeoutSeconds(); + if (closeTimeout != null) { + readTransform = readTransform.withCloseTimeout(Duration.standardSeconds(closeTimeout)); + } + + String ackMode = config.getAcknowledgeMode(); + if (ackMode != null) { + readTransform = + readTransform.withAcknowledgeMode(JmsIO.AcknowledgeMode.valueOf(ackMode.toUpperCase())); + } + + Integer indAckCode = config.getIndividualAcknowledgeModeCode(); + if (indAckCode != null) { + readTransform = readTransform.withIndividualAcknowledgeModeCode(indAckCode); + } + + Schema outputSchema = Schema.builder().addStringField("payload").build(); + + PCollection outputRows = + input + .getPipeline() + .apply(readTransform) + .apply( + "Wrap in Beam Rows", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element JmsRecord record, OutputReceiver outputReceiver) { + String payload = record.getPayload(); + if (payload == null) { + payload = ""; + } + outputReceiver.output( + Row.withSchema(outputSchema).addValue(payload).build()); + } + })) + .setRowSchema(outputSchema); + + return PCollectionRowTuple.of("output", outputRows); + } + } +} diff --git a/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsWriteSchemaTransformProvider.java b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsWriteSchemaTransformProvider.java new file mode 100644 index 000000000000..34316c9a8c1e --- /dev/null +++ b/sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsWriteSchemaTransformProvider.java @@ -0,0 +1,191 @@ +/* + * 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 static org.apache.beam.sdk.io.jms.JmsWriteSchemaTransformProvider.WriteConfiguration; + +import com.google.auto.service.AutoService; +import com.google.auto.value.AutoValue; +import java.io.Serializable; +import java.nio.charset.StandardCharsets; +import java.util.Collections; +import java.util.List; +import javax.annotation.Nullable; +import org.apache.beam.sdk.schemas.AutoValueSchema; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; +import org.apache.beam.sdk.schemas.transforms.SchemaTransform; +import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; +import org.apache.beam.sdk.schemas.transforms.TypedSchemaTransformProvider; +import org.apache.beam.sdk.transforms.DoFn; +import org.apache.beam.sdk.transforms.ParDo; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions; + +@AutoService(SchemaTransformProvider.class) +public class JmsWriteSchemaTransformProvider + extends TypedSchemaTransformProvider { + @DefaultSchema(AutoValueSchema.class) + @AutoValue + public abstract static class WriteConfiguration implements Serializable { + public static Builder builder() { + return new AutoValue_JmsWriteSchemaTransformProvider_WriteConfiguration.Builder(); + } + + @SchemaFieldDescription( + "Configuration options to set up the JMS connection.\nNote: if " + + "connection factory class name is set, spin up a persistent expansion service with " + + "the provider client JAR on the classpath:\n" + + "java -cp :\n" + + "org.apache.beam.sdk.expansion.service.ExpansionService \n" + + "and pass expansion_service='localhost:' to the transform. Currently, only " + + "ActiveMQ (org.apache.activemq.ActiveMQConnectionFactory) is embedded into the expansion service") + public abstract ConnectionConfiguration getConnectionConfiguration(); + + @SchemaFieldDescription( + "The JMS queue to write to. Exclusively one of queue or topic must be specified.") + @Nullable + public abstract String getQueue(); + + @SchemaFieldDescription( + "The JMS topic to write to. Exclusively one of queue or topic must be specified.") + @Nullable + public abstract String getTopic(); + + @AutoValue.Builder + public abstract static class Builder { + public abstract Builder setConnectionConfiguration( + ConnectionConfiguration connectionConfiguration); + + public abstract Builder setQueue(String queue); + + public abstract Builder setTopic(String topic); + + public abstract WriteConfiguration build(); + } + } + + @Override + public String identifier() { + return "beam:schematransform:org.apache.beam:jms_write:v1"; + } + + @Override + public String description() { + return "Publishes messages to a JMS broker. Expects an input PCollection of rows with a " + + "`payload` (string) or `bytes` field, each of which is published as one JMS TextMessage.\n" + + "\n" + + "Works with both bounded (batch) and unbounded (streaming) input PCollections."; + } + + @Override + public List inputCollectionNames() { + return Collections.singletonList("input"); + } + + @Override + protected SchemaTransform from(WriteConfiguration configuration) { + return new JmsWriteSchemaTransform(configuration); + } + + private static class JmsWriteSchemaTransform extends SchemaTransform { + private final WriteConfiguration config; + + JmsWriteSchemaTransform(WriteConfiguration configuration) { + this.config = configuration; + } + + @Override + public PCollectionRowTuple expand(PCollectionRowTuple input) { + PCollection inputRows = input.getSinglePCollection(); + + Preconditions.checkArgument( + (config.getQueue() != null) != (config.getTopic() != null), + "Exactly one of queue or topic must be specified."); + + int fieldIndex = -1; + boolean isBytes = false; + Schema schema = inputRows.getSchema(); + if (schema.hasField("payload") + && schema.getField("payload").getType().equals(Schema.FieldType.STRING)) { + fieldIndex = schema.indexOf("payload"); + } else if (schema.hasField("bytes") + && schema.getField("bytes").getType().equals(Schema.FieldType.BYTES)) { + fieldIndex = schema.indexOf("bytes"); + isBytes = true; + } else if (schema.getFieldCount() == 1 + && schema.getField(0).getType().equals(Schema.FieldType.STRING)) { + fieldIndex = 0; + } else if (schema.getFieldCount() == 1 + && schema.getField(0).getType().equals(Schema.FieldType.BYTES)) { + fieldIndex = 0; + isBytes = true; + } else { + throw new IllegalStateException( + String.format( + "Expected input Schema to have a 'payload' (STRING) or 'bytes' (BYTES) field, or" + + " a single string/bytes field, but received: %s", + schema)); + } + + JmsIO.Write writeTransform = + JmsIO.write() + .withConnectionConfiguration(config.getConnectionConfiguration()) + .withValueMapper(new TextMessageMapper()); + + String queue = config.getQueue(); + if (queue != null) { + writeTransform = writeTransform.withQueue(queue); + } + String topic = config.getTopic(); + if (topic != null) { + writeTransform = writeTransform.withTopic(topic); + } + + final int idx = fieldIndex; + final boolean bytesField = isBytes; + inputRows + .apply( + "Extract payload", + ParDo.of( + new DoFn() { + @ProcessElement + public void processElement( + @Element Row row, OutputReceiver outputReceiver) { + if (bytesField) { + byte[] bytes = + org.apache.beam.sdk.util.Preconditions.checkStateNotNull( + row.getBytes(idx)); + outputReceiver.output(new String(bytes, StandardCharsets.UTF_8)); + } else { + String payload = + org.apache.beam.sdk.util.Preconditions.checkStateNotNull( + row.getString(idx)); + outputReceiver.output(payload); + } + } + })) + .apply(writeTransform); + + return PCollectionRowTuple.empty(inputRows.getPipeline()); + } + } +} diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/ConnectionConfigurationTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/ConnectionConfigurationTest.java new file mode 100644 index 000000000000..de719b9208b7 --- /dev/null +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/ConnectionConfigurationTest.java @@ -0,0 +1,145 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import javax.jms.Connection; +import javax.jms.ConnectionFactory; +import javax.jms.JMSException; +import org.apache.activemq.ActiveMQConnectionFactory; +import org.apache.qpid.jms.JmsConnectionFactory; +import org.checkerframework.checker.nullness.qual.Nullable; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link ConnectionConfiguration}. */ +@RunWith(JUnit4.class) +public class ConnectionConfigurationTest { + + public static class CustomTestConnectionFactory implements BeamGenericJmsConnectionFactory { + @Override + public ConnectionFactory createConnectionFactory(ConnectionConfiguration config) { + return new DummyConnectionFactory( + config.getServerUri(), config.getUsername(), config.getPassword()); + } + } + + public static class DummyConnectionFactory implements ConnectionFactory { + private final String serverUri; + private final String username; + private final String password; + + public DummyConnectionFactory( + String serverUri, @Nullable String username, @Nullable String password) { + this.serverUri = serverUri; + this.username = username; + this.password = password; + } + + public String getServerUri() { + return serverUri; + } + + public String getUsername() { + return username; + } + + public String getPassword() { + return password; + } + + @Override + public Connection createConnection() throws JMSException { + return null; + } + + @Override + public Connection createConnection(String username, String password) throws JMSException { + return null; + } + + @Override + public javax.jms.JMSContext createContext() { + return null; + } + + @Override + public javax.jms.JMSContext createContext(int sessionMode) { + return null; + } + + @Override + public javax.jms.JMSContext createContext(String username, String password) { + return null; + } + + @Override + public javax.jms.JMSContext createContext(String username, String password, int sessionMode) { + return null; + } + } + + @Test + public void testDefaultActiveMQConnectionFactory() { + ConnectionConfiguration config = ConnectionConfiguration.create("vm://localhost"); + ConnectionFactory cf = config.createConnectionFactory(); + assertNotNull(cf); + assertTrue(cf instanceof ActiveMQConnectionFactory); + } + + @Test + public void testQpidConnectionFactory() { + ConnectionConfiguration config = + ConnectionConfiguration.create("amqp://localhost") + .withConnectionFactoryClassName("org.apache.qpid.jms.JmsConnectionFactory"); + ConnectionFactory cf = config.createConnectionFactory(); + assertNotNull(cf); + assertTrue(cf instanceof JmsConnectionFactory); + } + + @Test + public void testCustomBeamGenericJmsConnectionFactory() { + ConnectionConfiguration config = + ConnectionConfiguration.create("custom://localhost") + .withConnectionFactoryClassName(CustomTestConnectionFactory.class.getName()) + .withUsername("testUser") + .withPassword("testPass"); + ConnectionFactory cf = config.createConnectionFactory(); + assertNotNull(cf); + assertTrue(cf instanceof DummyConnectionFactory); + DummyConnectionFactory dummy = (DummyConnectionFactory) cf; + assertEquals("custom://localhost", dummy.getServerUri()); + assertEquals("testUser", dummy.getUsername()); + assertEquals("testPass", dummy.getPassword()); + } + + @Test + public void testUnsupportedConnectionFactoryClass() { + ConnectionConfiguration config = + ConnectionConfiguration.create("tcp://localhost") + .withConnectionFactoryClassName("java.lang.String"); + IllegalArgumentException exception = + assertThrows(IllegalArgumentException.class, config::createConnectionFactory); + assertTrue(exception.getMessage().contains("BeamGenericJmsConnectionFactory")); + } +} diff --git a/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsSchemaTransformProviderTest.java b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsSchemaTransformProviderTest.java new file mode 100644 index 000000000000..b354ef94a008 --- /dev/null +++ b/sdks/java/io/jms/src/test/java/org/apache/beam/sdk/io/jms/JmsSchemaTransformProviderTest.java @@ -0,0 +1,263 @@ +/* + * 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 static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import java.util.List; +import java.util.ServiceLoader; +import java.util.stream.Collectors; +import java.util.stream.StreamSupport; +import org.apache.beam.sdk.io.jms.JmsReadSchemaTransformProvider.ReadConfiguration; +import org.apache.beam.sdk.io.jms.JmsWriteSchemaTransformProvider.WriteConfiguration; +import org.apache.beam.sdk.schemas.Schema; +import org.apache.beam.sdk.schemas.transforms.SchemaTransform; +import org.apache.beam.sdk.schemas.transforms.SchemaTransformProvider; +import org.apache.beam.sdk.testing.TestPipeline; +import org.apache.beam.sdk.transforms.Create; +import org.apache.beam.sdk.values.PCollection; +import org.apache.beam.sdk.values.PCollectionRowTuple; +import org.apache.beam.sdk.values.Row; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Lists; +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.Sets; +import org.junit.Rule; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Tests for {@link JmsReadSchemaTransformProvider} and {@link JmsWriteSchemaTransformProvider}. */ +@RunWith(JUnit4.class) +public class JmsSchemaTransformProviderTest { + + @Rule + public final transient TestPipeline pipeline = + TestPipeline.create().enableAbandonedNodeEnforcement(false); + + @Test + public void testReadFindTransform() { + ServiceLoader serviceLoader = + ServiceLoader.load(SchemaTransformProvider.class); + List providers = + StreamSupport.stream(serviceLoader.spliterator(), false) + .filter(provider -> provider.getClass() == JmsReadSchemaTransformProvider.class) + .collect(Collectors.toList()); + SchemaTransformProvider jmsProvider = providers.get(0); + + assertEquals(Lists.newArrayList("output"), jmsProvider.outputCollectionNames()); + assertEquals(Lists.newArrayList(), jmsProvider.inputCollectionNames()); + assertEquals("beam:schematransform:org.apache.beam:jms_read:v1", jmsProvider.identifier()); + assertNotNull(jmsProvider.description()); + + assertEquals( + Sets.newHashSet( + "connection_configuration", + "queue", + "topic", + "max_num_records", + "max_read_time_seconds", + "close_timeout_seconds", + "acknowledge_mode", + "individual_acknowledge_mode_code"), + jmsProvider.configurationSchema().getFields().stream() + .map(Schema.Field::getName) + .collect(Collectors.toSet())); + } + + @Test + public void testReadBuildTransformWithQueue() { + ReadConfiguration readConfig = + ReadConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .setMaxNumRecords(100L) + .setMaxReadTimeSeconds(5L) + .setCloseTimeoutSeconds(10L) + .setAcknowledgeMode("CLIENT_ACKNOWLEDGE") + .setIndividualAcknowledgeModeCode(101) + .build(); + + SchemaTransform transform = new JmsReadSchemaTransformProvider().from(readConfig); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.empty(pipeline)); + + assertEquals(1, output.getAll().size()); + assertTrue(output.has("output")); + assertEquals( + Schema.builder().addStringField("payload").build(), output.get("output").getSchema()); + } + + @Test + public void testReadBuildTransformWithTopic() { + ReadConfiguration readConfig = + ReadConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setTopic("TEST_TOPIC") + .build(); + + SchemaTransform transform = new JmsReadSchemaTransformProvider().from(readConfig); + PCollectionRowTuple output = transform.expand(PCollectionRowTuple.empty(pipeline)); + + assertEquals(1, output.getAll().size()); + assertTrue(output.has("output")); + assertEquals( + Schema.builder().addStringField("payload").build(), output.get("output").getSchema()); + } + + @Test + public void testReadInvalidConfigurations() { + ReadConfiguration bothConfig = + ReadConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .setTopic("TEST_TOPIC") + .build(); + SchemaTransform bothTransform = new JmsReadSchemaTransformProvider().from(bothConfig); + assertThrows( + IllegalArgumentException.class, + () -> bothTransform.expand(PCollectionRowTuple.empty(pipeline))); + + ReadConfiguration neitherConfig = + ReadConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .build(); + SchemaTransform neitherTransform = new JmsReadSchemaTransformProvider().from(neitherConfig); + assertThrows( + IllegalArgumentException.class, + () -> neitherTransform.expand(PCollectionRowTuple.empty(pipeline))); + } + + @Test + public void testReadWithNonEmptyInputThrows() { + ReadConfiguration readConfig = + ReadConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .build(); + SchemaTransform transform = new JmsReadSchemaTransformProvider().from(readConfig); + + PCollection dummyInput = + pipeline.apply( + "CreateDummy", Create.empty(Schema.builder().addStringField("dummy").build())); + assertThrows( + IllegalStateException.class, + () -> transform.expand(PCollectionRowTuple.of("input", dummyInput))); + } + + @Test + public void testWriteFindTransform() { + ServiceLoader serviceLoader = + ServiceLoader.load(SchemaTransformProvider.class); + List providers = + StreamSupport.stream(serviceLoader.spliterator(), false) + .filter(provider -> provider.getClass() == JmsWriteSchemaTransformProvider.class) + .collect(Collectors.toList()); + SchemaTransformProvider jmsProvider = providers.get(0); + + assertEquals(Lists.newArrayList(), jmsProvider.outputCollectionNames()); + assertEquals(Lists.newArrayList("input"), jmsProvider.inputCollectionNames()); + assertEquals("beam:schematransform:org.apache.beam:jms_write:v1", jmsProvider.identifier()); + assertNotNull(jmsProvider.description()); + + assertEquals( + Sets.newHashSet("connection_configuration", "queue", "topic"), + jmsProvider.configurationSchema().getFields().stream() + .map(Schema.Field::getName) + .collect(Collectors.toSet())); + } + + @Test + public void testWriteBuildTransformWithQueueAndTopic() { + WriteConfiguration queueConfig = + WriteConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .build(); + SchemaTransform queueTransform = new JmsWriteSchemaTransformProvider().from(queueConfig); + Schema schema = Schema.builder().addStringField("payload").build(); + PCollection inputRows = pipeline.apply("CreateQueueRows", Create.empty(schema)); + PCollectionRowTuple output = queueTransform.expand(PCollectionRowTuple.of("input", inputRows)); + assertTrue(output.getAll().isEmpty()); + + WriteConfiguration topicConfig = + WriteConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setTopic("TEST_TOPIC") + .build(); + SchemaTransform topicTransform = new JmsWriteSchemaTransformProvider().from(topicConfig); + Schema bytesSchema = Schema.builder().addByteArrayField("bytes").build(); + PCollection bytesRows = pipeline.apply("CreateTopicRows", Create.empty(bytesSchema)); + PCollectionRowTuple topicOutput = + topicTransform.expand(PCollectionRowTuple.of("input", bytesRows)); + assertTrue(topicOutput.getAll().isEmpty()); + } + + @Test + public void testWriteInvalidConfigurations() { + WriteConfiguration bothConfig = + WriteConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .setTopic("TEST_TOPIC") + .build(); + SchemaTransform bothTransform = new JmsWriteSchemaTransformProvider().from(bothConfig); + Schema schema = Schema.builder().addStringField("payload").build(); + PCollection inputRows = pipeline.apply("CreateBothRows", Create.empty(schema)); + assertThrows( + IllegalArgumentException.class, + () -> bothTransform.expand(PCollectionRowTuple.of("input", inputRows))); + + WriteConfiguration neitherConfig = + WriteConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .build(); + SchemaTransform neitherTransform = new JmsWriteSchemaTransformProvider().from(neitherConfig); + PCollection inputRows2 = pipeline.apply("CreateNeitherRows", Create.empty(schema)); + assertThrows( + IllegalArgumentException.class, + () -> neitherTransform.expand(PCollectionRowTuple.of("input", inputRows2))); + } + + @Test + public void testWriteInvalidInputSchema() { + WriteConfiguration config = + WriteConfiguration.builder() + .setConnectionConfiguration( + JmsIO.ConnectionConfiguration.create("tcp://localhost:61616")) + .setQueue("TEST_QUEUE") + .build(); + SchemaTransform transform = new JmsWriteSchemaTransformProvider().from(config); + + Schema invalidSchema = Schema.builder().addInt32Field("id").addStringField("name").build(); + PCollection inputRows = + pipeline.apply("CreateInvalidSchemaRows", Create.empty(invalidSchema)); + assertThrows( + IllegalStateException.class, + () -> transform.expand(PCollectionRowTuple.of("input", inputRows))); + } +} diff --git a/sdks/java/io/messaging-expansion-service/build.gradle b/sdks/java/io/messaging-expansion-service/build.gradle index 6cdf67b86d6e..8693f4597bc1 100644 --- a/sdks/java/io/messaging-expansion-service/build.gradle +++ b/sdks/java/io/messaging-expansion-service/build.gradle @@ -42,6 +42,9 @@ dependencies { permitUnusedDeclared project(":sdks:java:expansion-service") // BEAM-11761 implementation project(":sdks:java:io:mqtt") permitUnusedDeclared project(":sdks:java:io:mqtt") // BEAM-11761 + implementation project(":sdks:java:io:jms") + permitUnusedDeclared project(":sdks:java:io:jms") // BEAM-11761 + runtimeOnly library.java.activemq_client runtimeOnly library.java.slf4j_jdk14 } diff --git a/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py b/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py new file mode 100644 index 000000000000..b0f64ab8f7fe --- /dev/null +++ b/sdks/python/apache_beam/io/external/xlang_jmsio_it_test.py @@ -0,0 +1,392 @@ +# +# 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. +# + +"""Integration tests for the cross-language JMS IO transforms +(ReadFromJms / WriteToJms), served by the messaging expansion service. + +Runs against ActiveMQ or IBM MQ brokers started once per test class via testcontainers. +""" + +import logging +import platform +import re +import threading +import time +import unittest + +import pytest + +import apache_beam as beam +from apache_beam.options.pipeline_options import PortableOptions +from apache_beam.options.pipeline_options import StandardOptions +from apache_beam.testing.test_pipeline import TestPipeline +from apache_beam.typehints.row_type import RowTypeConstraint + +# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports +try: + from apache_beam.io import ReadFromJms + from apache_beam.io import WriteToJms +except ImportError: + ReadFromJms = None + WriteToJms = None + +try: + from testcontainers.core.container import DockerContainer + from testcontainers.core.waiting_utils import wait_for_logs +except ImportError: + DockerContainer = None + +_LOGGER = logging.getLogger(__name__) +NUM_RECORDS = 100 +STRING_ROW = RowTypeConstraint.from_fields([('payload', str)]) + + +@pytest.mark.uses_messaging_java_expansion_service +@unittest.skipIf( + DockerContainer is None, 'testcontainers package is not installed') +@unittest.skipIf( + ReadFromJms is None or WriteToJms is None, + 'JMS cross-language wrappers are not generated') +@unittest.skipIf( + TestPipeline().get_pipeline_options().view_as(StandardOptions).runner + is None, + 'Do not run this test on precommit suites.') +@unittest.skipIf( + 'Dataflow' in ( + TestPipeline().get_pipeline_options().view_as(StandardOptions).runner or + ''), + 'The testcontainers broker is not reachable from Dataflow workers; ' + 'a Dataflow variant would need a remotely hosted JMS broker.') +class _BaseJmsIOTest(unittest.TestCase): + expansion_service = None + + def produce(self, source_queue, count): + raise NotImplementedError + + def browse_queue(self, sink_queue): + raise NotImplementedError + + def _connection_configuration(self, connection_param=None): + raise NotImplementedError + + def _run_streaming_test( + self, + source_queue, + sink_queue, + acknowledge_mode=None, + connection_param=None): + subscriber_result = {} + + def publish(): + self.produce(source_queue, remaining_records) + + stop_event = threading.Event() + + def subscribe(): + received_messages = [] + while len(received_messages) < NUM_RECORDS and not stop_event.is_set(): + time.sleep(5) + try: + messages = self.browse_queue(sink_queue) + if messages: + received_messages = messages + subscriber_result['received'] = received_messages + except Exception as e: + _LOGGER.warning('Error while browsing sink queue: %s', e) + break + _LOGGER.info('received %s messages', len(received_messages)) + + # TODO(https://github.com/apache/beam/issues/39446): Clean up + # pre-publishing Prism runner issue resolved + initial_records = 10 + remaining_records = NUM_RECORDS - initial_records + self.produce(source_queue, initial_records) + + publisher = threading.Thread(target=publish, daemon=True) + subscriber = threading.Thread(target=subscribe, daemon=True) + + p = TestPipeline(blocking=False) + p.get_pipeline_options().view_as(StandardOptions).streaming = True + p.not_use_test_runner_api = True + # Run pipeline without blocking + # TODO: Remove once subprocess cache leak fixed for pipeline running + # in LOOPBACK mode outside of with clause + p.__enter__() + try: + _ = ( + p + | 'ReadFromJms' >> ReadFromJms( + connection_configuration=self._connection_configuration( + connection_param), + queue=source_queue, + acknowledge_mode=acknowledge_mode, + expansion_service=self.expansion_service) + | + 'Passthrough' >> beam.Map(lambda row: beam.Row(payload=row.payload) + ).with_output_types(STRING_ROW) + | 'WriteToJms' >> WriteToJms( + connection_configuration=self._connection_configuration( + connection_param), + queue=sink_queue, + expansion_service=self.expansion_service)) + publisher.start() + result = p.run() + subscriber.start() + try: + subscriber.join(timeout=20) # 1.5 min + finally: + stop_event.set() + publisher.join() + try: + result.cancel() + except Exception: # pylint: disable=broad-except + _LOGGER.warning('Ignoring error while cancelling the pipeline.') + finally: + p._extra_context.__exit__(None, None, None) + + received = subscriber_result.get('received', []) + self.assertEqual(len(received), NUM_RECORDS) + # there are identical records + self.assertEqual(len(set(received)), NUM_RECORDS - initial_records) + + +class ActiveMQJmsIOTest(_BaseJmsIOTest): + @classmethod + def setUpClass(cls): + cls.start_jms_container(retries=3) + host = cls.broker.get_container_host_ip() + port = cls.broker.get_exposed_port(61616) + cls.server_uri = 'tcp://%s:%s' % (host, port) + + @classmethod + def tearDownClass(cls): + try: + cls.broker.stop() + except Exception: + logging.error('Could not stop the JMS broker container.') + + @classmethod + def start_jms_container(cls, retries): + for i in range(retries): + try: + cls.broker = DockerContainer( + 'apache/activemq-classic:5.18.3').with_exposed_ports(61616) + cls.broker.start() + wait_for_logs(cls.broker, '.*ActiveMQ .* started.*', timeout=30) + break + except Exception as e: + try: + cls.broker.stop() + except Exception: + pass + if i == retries - 1: + logging.error('Unable to initialize the JMS broker container.') + raise e + + def _connection_configuration(self, connection_param=None): + uri = self.server_uri + if connection_param: + uri += '?' + connection_param + return { + 'server_uri': uri, + 'connection_factory_class_name': 'org.apache.activemq.ActiveMQConnectionFactory' + } + + def produce(self, source_queue, count): + container = self.broker.get_wrapped_container() + exit_code, _ = container.exec_run([ + '/opt/apache-activemq/bin/activemq', + 'producer', + '--destination', + 'queue://' + source_queue, + '--messageCount', + str(count), + '--persistent', + 'false' + ]) + if exit_code == 0: + _LOGGER.info('published %s messages to %s', count, source_queue) + else: + _LOGGER.warning('publishing message returns exit code %s', exit_code) + + def browse_queue(self, sink_queue): + container = self.broker.get_wrapped_container() + exit_code, output = container.exec_run([ + '/opt/apache-activemq/bin/activemq', + 'browse', + sink_queue + ]) + if exit_code == 0 and output: + return [ + line.split('JMS_BODY_FIELD:JMSText = ')[-1].strip() + for line in output.decode('utf-8').splitlines() + if 'JMS_BODY_FIELD:JMSText = ' in line + ] + return [] + + def test_xlang_jms_write_read_queue_ind_ack(self): + self._run_streaming_test( + source_queue='xlang-jms-ind-source', + sink_queue='xlang-jms-ind-sink', + acknowledge_mode='INDIVIDUAL_ACKNOWLEDGE') + + def test_xlang_jms_write_read_queue(self): + self._run_streaming_test( + source_queue='xlang-jms-source', + sink_queue='xlang-jms-sink', + acknowledge_mode='CLIENT_ACKNOWLEDGE_UNSAFE', + connection_param='jms.prefetchPolicy.all=0') + + +class IbmMqJmsIOTest(_BaseJmsIOTest): + @classmethod + def setUpClass(cls): + cls.start_ibm_mq_container(retries=3) + + @classmethod + def tearDownClass(cls): + if getattr(cls, 'expansion_service_obj', None): + try: + cls.expansion_service_obj.__exit__(None, None, None) + except Exception: + pass + if getattr(cls, 'broker', None): + try: + cls.broker.stop() + except Exception: + logging.error('Could not stop the IBM MQ broker container.') + + @classmethod + def get_ibm_mq_image(cls): + arch = platform.machine().lower() + if 'arm' in arch or 'aarch64' in arch: + try: + import docker + client = docker.from_env() + for img in client.images.list(): + for tag in img.tags: + if ('ibm-mq' in tag.lower() or + 'ibm_mq' in tag.lower()) and 'arm64' in tag.lower(): + _LOGGER.info('Found local ARM64 IBM MQ image: %s', tag) + return tag + except Exception as e: + _LOGGER.warning('Failed to inspect local docker images: %s', e) + + raise RuntimeError( + 'Official IBM MQ docker images do not support ARM macOS (aarch64). ' + 'Please build an ARM64 image locally from ' + 'https://github.com/ibm-messaging/mq-container ' + 'and tag it with an "-arm64" suffix (e.g. localhost/ibm-mqadvanced-server-dev:9.4.0.0-arm64).' + ) + else: + return 'icr.io/ibm-messaging/mq:9.3.0.25-r1' + + @classmethod + def start_ibm_mq_container(cls, retries): + from apache_beam.transforms.external import BeamJarExpansionService + image_tag = cls.get_ibm_mq_image() + for i in range(retries): + try: + cls.broker = DockerContainer(image_tag).with_env( + 'LICENSE', 'accept').with_env('MQ_QMGR_NAME', 'QM1').with_env( + 'MQ_APP_PASSWORD', 'admin123').with_exposed_ports(1414) + cls.broker.start() + wait_for_logs( + cls.broker, '.*(MQQMNAME|Started queue manager).*', timeout=45) + host = cls.broker.get_container_host_ip() + port = cls.broker.get_exposed_port(1414) + cls.server_uri = 'tcp://%s:%s' % (host, port) + cls.expansion_service_obj = BeamJarExpansionService( + 'sdks:java:io:messaging-expansion-service:shadowJar', + classpath=[ + 'com.ibm.mq:com.ibm.mq.allclient:9.3.0.25', + 'org.json:json:20251224' + ]) + cls.expansion_service = cls.expansion_service_obj.__enter__() + break + except Exception as e: + if getattr(cls, 'broker', None): + try: + cls.broker.stop() + except Exception: + pass + if i == retries - 1: + logging.error('Unable to initialize the IBM MQ broker container.') + raise e + + def _connection_configuration(self, connection_param=None): + uri = self.server_uri + '?channel=DEV.APP.SVRCONN&queueManager=QM1' + if connection_param: + uri += '&' + connection_param + return { + 'server_uri': uri, + 'connection_factory_class_name': 'com.ibm.mq.jms.MQConnectionFactory', + 'username': 'app', + 'password': 'admin123' + } + + def produce(self, source_queue, count): + container = self.broker.get_wrapped_container() + cmd = ( + f'for i in $(seq 0 {count-1}); do ' + f'echo "test message: $i" | /opt/mqm/samp/bin/amqsput {source_queue} QM1 >/dev/null 2>&1; ' + f'done') + container.exec_run(['sh', '-c', cmd]) + + def browse_queue(self, sink_queue): + container = self.broker.get_wrapped_container() + exit_code, output = container.exec_run([ + '/opt/mqm/bin/dmpmqmsg', + '-m', + 'QM1', + '-i', + sink_queue, + '-f', + 'stdout', + '-d', + 'p' + ]) + + if exit_code == 0 and output: + content = output.decode('utf-8', errors='ignore') + messages = [] + current_msg = [] + for line in content.splitlines(): + # Example raw result: + # S "97491582 test message: 7" + # S "8" + if 'test message:' in line: + current_msg.append(line[1:].strip('"')) + elif re.match(r'^S "\d+"$', line): + current_msg.append(line[2:].strip('"') + " ") + if current_msg: + full_str = "".join(current_msg) + # extract all "test message: [number]" from the full string + messages = re.findall(r'test message: \d+', full_str) + return messages + return [] + + def test_xlang_jms_write_read_queue_client_ack(self): + self._run_streaming_test( + source_queue='DEV.QUEUE.1', + sink_queue='DEV.QUEUE.2', + acknowledge_mode='CLIENT_ACKNOWLEDGE') + + +if __name__ == '__main__': + logging.getLogger().setLevel(logging.INFO) + unittest.main() diff --git a/sdks/python/apache_beam/io/external/xlang_mqttio_it_test.py b/sdks/python/apache_beam/io/external/xlang_mqttio_it_test.py index 889096d1f22b..dad85a688ecb 100644 --- a/sdks/python/apache_beam/io/external/xlang_mqttio_it_test.py +++ b/sdks/python/apache_beam/io/external/xlang_mqttio_it_test.py @@ -214,46 +214,40 @@ def subscribe(): publisher.start() subscriber.start() - # MqttIO read is unbounded, so this pipeline runs in streaming mode and - # never terminates on its own. Amend the harness-provided pipeline options - # rather than discarding them: enable streaming, run non-blocking so the - # observe-then-cancel logic below can execute, and target the Prism portable - # runner. The latter is required because SwitchingDirectRunner disables its - # Prism delegation for pipelines containing external (cross-language) - # transforms (see runners/direct/direct_runner.py) and falls back to the - # BundleBasedDirectRunner, which cannot execute an unbounded read. - # The runner is instantiated during TestPipeline construction, so it must be - # passed to the constructor; the remaining harness-provided options are - # preserved and only amended (streaming + LOOPBACK environment) afterwards. - p = TestPipeline(runner='PrismRunner', blocking=False) + p = TestPipeline(blocking=False) p.get_pipeline_options().view_as(StandardOptions).streaming = True - p.get_pipeline_options().view_as( - PortableOptions).environment_type = 'LOOPBACK' p.not_use_test_runner_api = True - _ = ( - p - | 'ReadFromMqtt' >> ReadFromMqtt( - connection_configuration=self._connection_configuration( - source_topic, 'xlang-mqtt-streaming-read')) - | 'Passthrough' >> beam.Map( - lambda row: beam.Row(bytes=row.bytes)).with_output_types(BYTES_ROW) - | 'WriteToMqtt' >> WriteToMqtt( - connection_configuration=self._connection_configuration( - sink_topic, 'xlang-mqtt-streaming-write'))) - result = p.run() + # Run pipeline without blocking + # TODO: Remove once subprocess cache leak fixed for pipeline running + # in LOOPBACK mode outside of with clause + p.__enter__() try: - # The subscriber exits once NUM_RECORDS messages flowed through the - # streaming pipeline (or fails the assertions below on its timeout). - subscriber.join(timeout=200) - finally: - stop_publishing.set() - publisher.join() + _ = ( + p + | 'ReadFromMqtt' >> ReadFromMqtt( + connection_configuration=self._connection_configuration( + source_topic, 'xlang-mqtt-streaming-read')) + | 'Passthrough' >> beam.Map(lambda row: beam.Row(bytes=row.bytes)). + with_output_types(BYTES_ROW) + | 'WriteToMqtt' >> WriteToMqtt( + connection_configuration=self._connection_configuration( + sink_topic, 'xlang-mqtt-streaming-write'))) + result = p.run() try: - result.cancel() - except Exception: # pylint: disable=broad-except - # The unbounded pipeline never finishes on its own; cancellation - # after the assertion data was collected is best-effort. - logging.warning('Ignoring error while cancelling the pipeline.') + # The subscriber exits once NUM_RECORDS messages flowed through the + # streaming pipeline (or fails the assertions below on its timeout). + subscriber.join(timeout=200) + finally: + stop_publishing.set() + publisher.join() + try: + result.cancel() + except Exception: # pylint: disable=broad-except + # The unbounded pipeline never finishes on its own; cancellation + # after the assertion data was collected is best-effort. + logging.warning('Ignoring error while cancelling the pipeline.') + finally: + p._extra_context.__exit__(None, None, None) self.assertEqual(subscriber_result.get('exit_code'), 0) payloads = subscriber_result.get('output', b'').split() diff --git a/sdks/python/test-suites/direct/common.gradle b/sdks/python/test-suites/direct/common.gradle index 1dd15ecb09f9..10bff26afa5d 100644 --- a/sdks/python/test-suites/direct/common.gradle +++ b/sdks/python/test-suites/direct/common.gradle @@ -438,17 +438,28 @@ project.tasks.register("inferencePostCommitIT") { // Create cross-language tasks for running tests against Java expansion service(s) def gcpProject = project.findProperty('gcpProject') ?: 'apache-beam-testing' +def testDirectRunnerOpts = [ + "--runner=TestDirectRunner", + "--project=${gcpProject}", + "--temp_location=gs://temp-storage-for-end-to-end-tests/temp-it", +] + +def prismRunnerOpts = [ + "--runner=PrismRunner", + "--environment_type=LOOPBACK", + "--project=${gcpProject}", + "--temp_location=gs://temp-storage-for-end-to-end-tests/temp-it", +] + +def usePrismRunnerSuites = ["messagingCrossLanguage"] + project(":sdks:python:test-suites:xlang").ext.xlangTasks.each { taskMetadata -> createCrossLanguageUsingJavaExpansionTask( name: taskMetadata.name, expansionProjectPaths: taskMetadata.expansionProjectPaths, collectMarker: taskMetadata.collectMarker, numParallelTests: 1, - pythonPipelineOptions: [ - "--runner=TestDirectRunner", - "--project=${gcpProject}", - "--temp_location=gs://temp-storage-for-end-to-end-tests/temp-it", - ], + pythonPipelineOptions: usePrismRunnerSuites.contains(taskMetadata.name) ? prismRunnerOpts : testDirectRunnerOpts, pytestOptions: [ "--capture=no", // print stdout instantly "--timeout=4500", // timeout of whole command execution diff --git a/sdks/standard_expansion_services.yaml b/sdks/standard_expansion_services.yaml index d0c7f125c44b..cb617d3eba12 100644 --- a/sdks/standard_expansion_services.yaml +++ b/sdks/standard_expansion_services.yaml @@ -62,6 +62,10 @@ name: 'WriteToMqtt' 'beam:schematransform:org.apache.beam:mqtt_read:v1': name: 'ReadFromMqtt' + 'beam:schematransform:org.apache.beam:jms_write:v1': + name: 'WriteToJms' + 'beam:schematransform:org.apache.beam:jms_read:v1': + name: 'ReadFromJms' skip_transforms: # Core SchemaTransforms bundled via :sdks:java:expansion-service; already # generated from the Java IO expansion service above. diff --git a/sdks/standard_external_transforms.yaml b/sdks/standard_external_transforms.yaml index a3f7d5911dbb..b589448d64b5 100644 --- a/sdks/standard_external_transforms.yaml +++ b/sdks/standard_external_transforms.yaml @@ -19,7 +19,7 @@ # configuration in /sdks/standard_expansion_services.yaml. # Refer to gen_xlang_wrappers.py for more info. # -# Last updated on: 2026-06-11 +# Last updated on: 2026-07-22 - default_service: sdks:java:io:expansion-service:shadowJar description: '' @@ -214,6 +214,107 @@ type: str identifier: beam:schematransform:org.apache.beam:tfrecord_write:v1 name: TfrecordWrite +- default_service: sdks:java:io:messaging-expansion-service:shadowJar + description: 'Reads messages from a JMS broker and outputs each message as a single + `payload` (string) field. + + + By default the read is unbounded (streaming): it keeps consuming messages from + the specified queue or topic until the pipeline is stopped. Setting `maxNumRecords` + and/or `maxReadTimeSeconds` bounds the read, producing a bounded (batch) PCollection.' + destinations: + python: apache_beam/io + fields: + - description: 'The JMS acknowledge mode: CLIENT_ACKNOWLEDGE, CLIENT_ACKNOWLEDGE_UNSAFE, + or INDIVIDUAL_ACKNOWLEDGE.' + name: acknowledge_mode + nullable: true + type: str + - description: Close timeout for the JMS connection in seconds. + name: close_timeout_seconds + nullable: true + type: int64 + - description: 'Configuration options to set up the JMS connection. + + Note: if connection factory class name is set, spin up a persistent expansion + service with the provider client JAR on the classpath: + + java -cp : + + org.apache.beam.sdk.expansion.service.ExpansionService + + and pass expansion_service=''localhost:'' to the transform. Currently, + only ActiveMQ (org.apache.activemq.ActiveMQConnectionFactory) is embedded into + the expansion service' + name: connection_configuration + nullable: false + type: Row(connection_factory_class_name=typing.Optional[str], password=typing.Optional[str], + server_uri=, username=typing.Optional[str]) + - description: The proprietary integer code for individual message acknowledgment + when using INDIVIDUAL_ACKNOWLEDGE. + name: individual_acknowledge_mode_code + nullable: true + type: int32 + - description: The max number of records to receive. Setting this will result in + a bounded PCollection. + name: max_num_records + nullable: true + type: int64 + - description: The maximum time for this source to read messages. Setting this will + result in a bounded PCollection. + name: max_read_time_seconds + nullable: true + type: int64 + - description: The JMS queue to read from. Exclusively one of queue or topic must + be specified. + name: queue + nullable: true + type: str + - description: The JMS topic to read from. Exclusively one of queue or topic must + be specified. + name: topic + nullable: true + type: str + identifier: beam:schematransform:org.apache.beam:jms_read:v1 + name: ReadFromJms +- default_service: sdks:java:io:messaging-expansion-service:shadowJar + description: 'Publishes messages to a JMS broker. Expects an input PCollection of + rows with a `payload` (string) or `bytes` field, each of which is published as + one JMS TextMessage. + + + Works with both bounded (batch) and unbounded (streaming) input PCollections.' + destinations: + python: apache_beam/io + fields: + - description: 'Configuration options to set up the JMS connection. + + Note: if connection factory class name is set, spin up a persistent expansion + service with the provider client JAR on the classpath: + + java -cp : + + org.apache.beam.sdk.expansion.service.ExpansionService + + and pass expansion_service=''localhost:'' to the transform. Currently, + only ActiveMQ (org.apache.activemq.ActiveMQConnectionFactory) is embedded into + the expansion service' + name: connection_configuration + nullable: false + type: Row(connection_factory_class_name=typing.Optional[str], password=typing.Optional[str], + server_uri=, username=typing.Optional[str]) + - description: The JMS queue to write to. Exclusively one of queue or topic must + be specified. + name: queue + nullable: true + type: str + - description: The JMS topic to write to. Exclusively one of queue or topic must + be specified. + name: topic + nullable: true + type: str + identifier: beam:schematransform:org.apache.beam:jms_write:v1 + name: WriteToJms - default_service: sdks:java:io:messaging-expansion-service:shadowJar description: 'Reads messages from an MQTT broker and outputs each payload as a single `bytes` field.