SolaceIO write connector

CHANGES.md

@@ -70,6 +70,7 @@
 * [Managed Iceberg] Add support for TIMESTAMP, TIME, and DATE types ([#32688](https://github.com/apache/beam/pull/32688))
 * BigQuery CDC writes are now available in Python SDK, only supported when using StorageWrite API at least once mode ([#32527](https://github.com/apache/beam/issues/32527))
 * [Managed Iceberg] Allow updating table partition specs during pipeline runtime ([#32879](https://github.com/apache/beam/pull/32879))
+* Support for writing to [Solace messages queues](https://solace.com/) (`SolaceIO.Write`) added (Java) ([#31905](https://github.com/apache/beam/issues/31905)).
 
 ## New Features / Improvements
 

sdks/java/io/solace/build.gradle

@@ -53,6 +53,7 @@ dependencies {
     testImplementation library.java.junit
     testImplementation project(path: ":sdks:java:io:common")
     testImplementation project(path: ":sdks:java:testing:test-utils")
+    testImplementation project(path: ":sdks:java:core", configuration: "shadowTest")
     testRuntimeOnly library.java.slf4j_jdk14
     testImplementation library.java.testcontainers_solace
     testRuntimeOnly project(path: ":runners:direct-java", configuration: "shadow")

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/BasicAuthJcsmpSessionService.java

@@ -19,6 +19,7 @@
 
 import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 
+import com.google.auto.value.AutoValue;
 import com.solacesystems.jcsmp.ConsumerFlowProperties;
 import com.solacesystems.jcsmp.EndpointProperties;
 import com.solacesystems.jcsmp.FlowReceiver;
@@ -28,9 +29,15 @@
 import com.solacesystems.jcsmp.JCSMPProperties;
 import com.solacesystems.jcsmp.JCSMPSession;
 import com.solacesystems.jcsmp.Queue;
+import com.solacesystems.jcsmp.XMLMessageProducer;
 import java.io.IOException;
+import java.util.Objects;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ConcurrentLinkedQueue;
 import javax.annotation.Nullable;
 import org.apache.beam.sdk.io.solace.RetryCallableManager;
+import org.apache.beam.sdk.io.solace.SolaceIO.SubmissionMode;
+import org.apache.beam.sdk.io.solace.data.Solace.PublishResult;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableSet;
 
 /**
@@ -39,34 +46,50 @@
  * <p>This class provides a way to connect to a Solace broker and receive messages from a queue. The
  * connection is established using basic authentication.
  */
-public class BasicAuthJcsmpSessionService extends SessionService {
-  private final String queueName;
-  private final String host;
-  private final String username;
-  private final String password;
-  private final String vpnName;
-  @Nullable private JCSMPSession jcsmpSession;
-  @Nullable private MessageReceiver messageReceiver;
-  private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
+@AutoValue
+public abstract class BasicAuthJcsmpSessionService extends SessionService {
+
+  /** The name of the queue to receive messages from. */
+  public abstract @Nullable String queueName();
+
+  /** The host name or IP address of the Solace broker. Format: Host[:Port] */
+  public abstract String host();
+
+  /** The username to use for authentication. */
+  public abstract String username();
+
+  /** The password to use for authentication. */
+  public abstract String password();
+
+  /** The name of the VPN to connect to. */
+  public abstract String vpnName();
+
+  public static Builder builder() {
+    return new AutoValue_BasicAuthJcsmpSessionService.Builder().vpnName(DEFAULT_VPN_NAME);
+  }
+
+  @AutoValue.Builder
+  public abstract static class Builder {
+    public abstract Builder queueName(@Nullable String queueName);
+
+    public abstract Builder host(String host);
 
-  /**
-   * Creates a new {@link BasicAuthJcsmpSessionService} with the given parameters.
-   *
-   * @param queueName The name of the queue to receive messages from.
-   * @param host The host name or IP address of the Solace broker. Format: Host[:Port]
-   * @param username The username to use for authentication.
-   * @param password The password to use for authentication.
-   * @param vpnName The name of the VPN to connect to.
-   */
-  public BasicAuthJcsmpSessionService(
-      String queueName, String host, String username, String password, String vpnName) {
-    this.queueName = queueName;
-    this.host = host;
-    this.username = username;
-    this.password = password;
-    this.vpnName = vpnName;
+    public abstract Builder username(String username);
+
+    public abstract Builder password(String password);
+
+    public abstract Builder vpnName(String vpnName);
+
+    public abstract BasicAuthJcsmpSessionService build();
   }
 
+  @Nullable private transient JCSMPSession jcsmpSession;
+  @Nullable private transient MessageReceiver messageReceiver;
+  @Nullable private transient MessageProducer messageProducer;
+  private final java.util.Queue<PublishResult> publishedResultsQueue =
+      new ConcurrentLinkedQueue<>();
+  private final RetryCallableManager retryCallableManager = RetryCallableManager.create();
+
   @Override
   public void connect() {
     retryCallableManager.retryCallable(this::connectSession, ImmutableSet.of(JCSMPException.class));
@@ -79,6 +102,9 @@ public void close() {
           if (messageReceiver != null) {
             messageReceiver.close();
           }
+          if (messageProducer != null) {
+            messageProducer.close();
+          }
           if (!isClosed()) {
             checkStateNotNull(jcsmpSession).closeSession();
           }
@@ -88,24 +114,64 @@ public void close() {
   }
 
   @Override
-  public MessageReceiver createReceiver() {
-    this.messageReceiver =
-        retryCallableManager.retryCallable(
-            this::createFlowReceiver, ImmutableSet.of(JCSMPException.class));
+  public MessageReceiver getReceiver() {
+    if (this.messageReceiver == null) {
+      this.messageReceiver =
+          retryCallableManager.retryCallable(
+              this::createFlowReceiver, ImmutableSet.of(JCSMPException.class));
+    }
     return this.messageReceiver;
   }
 
+  @Override
+  public MessageProducer getInitializedProducer(SubmissionMode submissionMode) {
+    if (this.messageProducer == null || this.messageProducer.isClosed()) {
+      Callable<MessageProducer> create = () -> createXMLMessageProducer(submissionMode);
+      this.messageProducer =
+          retryCallableManager.retryCallable(create, ImmutableSet.of(JCSMPException.class));
+    }
+    return checkStateNotNull(this.messageProducer);
+  }
+
+  @Override
+  public java.util.Queue<PublishResult> getPublishedResultsQueue() {
+    return publishedResultsQueue;
+  }
+
   @Override
   public boolean isClosed() {
     return jcsmpSession == null || jcsmpSession.isClosed();
   }
 
+  private MessageProducer createXMLMessageProducer(SubmissionMode submissionMode)
+      throws JCSMPException, IOException {
+
+    if (isClosed()) {
+      connectWriteSession(submissionMode);
+    }
+
+    @SuppressWarnings("nullness")
+    Callable<XMLMessageProducer> initProducer =
+        () ->
+            Objects.requireNonNull(jcsmpSession)
+                .getMessageProducer(new PublishResultHandler(publishedResultsQueue));
+
+    XMLMessageProducer producer =
+        retryCallableManager.retryCallable(initProducer, ImmutableSet.of(JCSMPException.class));
+    if (producer == null) {
+      throw new IOException("SolaceIO.Write: Could not create producer, producer object is null");
+    }
+    return new SolaceMessageProducer(producer);
+  }
+
   private MessageReceiver createFlowReceiver() throws JCSMPException, IOException {
     if (isClosed()) {
       connectSession();
     }
 
-    Queue queue = JCSMPFactory.onlyInstance().createQueue(queueName);
+    Queue queue =
+        JCSMPFactory.onlyInstance()
+            .createQueue(checkStateNotNull(queueName(), "SolaceIO.Read: Queue is not set."));
 
     ConsumerFlowProperties flowProperties = new ConsumerFlowProperties();
     flowProperties.setEndpoint(queue);
@@ -118,7 +184,8 @@ private MessageReceiver createFlowReceiver() throws JCSMPException, IOException
           createFlowReceiver(jcsmpSession, flowProperties, endpointProperties));
     }
     throw new IOException(
-        "SolaceIO.Read: Could not create a receiver from the Jcsmp session: session object is null.");
+        "SolaceIO.Read: Could not create a receiver from the Jcsmp session: session object is"
+            + " null.");
   }
 
   // The `@SuppressWarning` is needed here, because the checkerframework reports an error for the
@@ -141,20 +208,33 @@ private int connectSession() throws JCSMPException {
     return 0;
   }
 
+  private int connectWriteSession(SubmissionMode mode) throws JCSMPException {
+    if (jcsmpSession == null) {
+      jcsmpSession = createWriteSessionObject(mode);
+    }
+    jcsmpSession.connect();
+    return 0;
+  }
+
   private JCSMPSession createSessionObject() throws InvalidPropertiesException {
     JCSMPProperties properties = initializeSessionProperties(new JCSMPProperties());
     return JCSMPFactory.onlyInstance().createSession(properties);
   }
 
+  private JCSMPSession createWriteSessionObject(SubmissionMode mode)
+      throws InvalidPropertiesException {
+    return JCSMPFactory.onlyInstance().createSession(initializeWriteSessionProperties(mode));
+  }
+
   @Override
   public JCSMPProperties initializeSessionProperties(JCSMPProperties baseProps) {
-    baseProps.setProperty(JCSMPProperties.VPN_NAME, vpnName);
+    baseProps.setProperty(JCSMPProperties.VPN_NAME, vpnName());
 
     baseProps.setProperty(
         JCSMPProperties.AUTHENTICATION_SCHEME, JCSMPProperties.AUTHENTICATION_SCHEME_BASIC);
-    baseProps.setProperty(JCSMPProperties.USERNAME, username);
-    baseProps.setProperty(JCSMPProperties.PASSWORD, password);
-    baseProps.setProperty(JCSMPProperties.HOST, host);
+    baseProps.setProperty(JCSMPProperties.USERNAME, username());
+    baseProps.setProperty(JCSMPProperties.PASSWORD, password());
+    baseProps.setProperty(JCSMPProperties.HOST, host());
     return baseProps;
   }
 }

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/BasicAuthJcsmpSessionServiceFactory.java

@@ -18,7 +18,6 @@
 package org.apache.beam.sdk.io.solace.broker;
 
 import static org.apache.beam.sdk.io.solace.broker.SessionService.DEFAULT_VPN_NAME;
-import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
 
 import com.google.auto.value.AutoValue;
 
@@ -31,12 +30,16 @@
  */
 @AutoValue
 public abstract class BasicAuthJcsmpSessionServiceFactory extends SessionServiceFactory {
+  /** The host name or IP address of the Solace broker. Format: Host[:Port] */
   public abstract String host();
 
+  /** The username to use for authentication. */
   public abstract String username();
 
+  /** The password to use for authentication. */
   public abstract String password();
 
+  /** The name of the VPN to connect to. */
   public abstract String vpnName();
 
   public static Builder builder() {
@@ -54,6 +57,7 @@ public abstract static class Builder {
 
     /** Set Solace username. */
     public abstract Builder username(String username);
+
     /** Set Solace password. */
     public abstract Builder password(String password);
 
@@ -65,11 +69,15 @@ public abstract static class Builder {
 
   @Override
   public SessionService create() {
-    return new BasicAuthJcsmpSessionService(
-        checkStateNotNull(queue, "SolaceIO.Read: Queue is not set.").getName(),
-        host(),
-        username(),
-        password(),
-        vpnName());
+    BasicAuthJcsmpSessionService.Builder builder = BasicAuthJcsmpSessionService.builder();
+    if (queue != null) {
+      builder = builder.queueName(queue.getName());
+    }
+    return builder
+        .host(host())
+        .username(username())
+        .password(password())
+        .vpnName(vpnName())
+        .build();
   }
 }

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/GCPSecretSessionServiceFactory.java

@@ -117,7 +117,7 @@ public abstract static class Builder {
 
   @Override
   public SessionService create() {
-    String password = null;
+    String password;
     try {
       password = retrieveSecret();
     } catch (IOException e) {

sdks/java/io/solace/src/main/java/org/apache/beam/sdk/io/solace/broker/MessageProducer.java

@@ -0,0 +1,61 @@
+/*
+ * 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.solace.broker;
+
+import com.solacesystems.jcsmp.DeliveryMode;
+import com.solacesystems.jcsmp.Destination;
+import java.util.List;
+import org.apache.beam.sdk.annotations.Internal;
+import org.apache.beam.sdk.io.solace.data.Solace;
+import org.apache.beam.sdk.transforms.SerializableFunction;
+
+/**
+ * Base class for publishing messages to a Solace broker.
+ *
+ * <p>Implementations of this interface are responsible for managing the connection to the broker
+ * and for publishing messages to the broker.
+ */
+@Internal
+public interface MessageProducer {
+
+  /** Publishes a message to the broker. */
+  void publishSingleMessage(
+      Solace.Record msg,
+      Destination topicOrQueue,
+      boolean useCorrelationKeyLatency,
+      DeliveryMode deliveryMode);
+
+  /**
+   * Publishes a batch of messages to the broker.
+   *
+   * <p>The size of the batch cannot exceed 50 messages, this is a limitation of the Solace API.
+   *
+   * <p>It returns the number of messages written.
+   */
+  int publishBatch(
+      List<Solace.Record> records,
+      boolean useCorrelationKeyLatency,
+      SerializableFunction<Solace.Record, Destination> destinationFn,
+      DeliveryMode deliveryMode);
+
+  /** Returns {@literal true} if the message producer is closed, {@literal false} otherwise. */
+  boolean isClosed();
+
+  /** Closes the message producer. */
+  void close();
+}