[Gemini] Support JmsIO SchemaTransform and cross-lang

Author: Abacn

.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
 }

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)

sdks/java/io/jms/src/main/java/org/apache/beam/sdk/io/jms/JmsIO.java

@@ -56,6 +56,9 @@
 import org.apache.beam.sdk.metrics.Metrics;
 import org.apache.beam.sdk.options.ExecutorOptions;
 import org.apache.beam.sdk.options.PipelineOptions;
+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.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.PTransform;
 import org.apache.beam.sdk.transforms.ParDo;
@@ -73,6 +76,7 @@
 import org.apache.beam.sdk.values.TupleTagList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings;
 import org.checkerframework.checker.initialization.qual.Initialized;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.Duration;
@@ -208,6 +212,117 @@ public static <EventT> Write<EventT> write() {
     return new AutoValue_JmsIO_Write.Builder<EventT>().build();
   }
 
+  /** A POJO describing a JMS connection. */
+  @DefaultSchema(AutoValueSchema.class)
+  @AutoValue
+  public abstract static class ConnectionConfiguration implements Serializable {
+    public static Builder builder() {
+      return new AutoValue_JmsIO_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 (Strings.isNullOrEmpty(className)) {
+        className = "org.apache.activemq.ActiveMQConnectionFactory";
+      }
+      try {
+        Class<?> clazz = Class.forName(className);
+        String uri = getServerUri();
+        String username = getUsername();
+        String password = getPassword();
+
+        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 constructor + setters
+          }
+        }
+        try {
+          ConnectionFactory cf =
+              (ConnectionFactory) clazz.getConstructor(String.class).newInstance(uri);
+          if (username != null && password != null) {
+            try {
+              clazz.getMethod("setUserName", String.class).invoke(cf, username);
+              clazz.getMethod("setPassword", String.class).invoke(cf, password);
+            } catch (NoSuchMethodException e) {
+              try {
+                clazz.getMethod("setUsername", String.class).invoke(cf, username);
+                clazz.getMethod("setPassword", String.class).invoke(cf, password);
+              } catch (NoSuchMethodException e2) {
+                // ignore if setters not found
+              }
+            }
+          }
+          return cf;
+        } catch (NoSuchMethodException e) {
+          return (ConnectionFactory) clazz.getConstructor().newInstance();
+        }
+      } catch (Exception e) {
+        throw new IllegalArgumentException(
+            "Unable to instantiate JMS ConnectionFactory of class " + className, e);
+      }
+    }
+  }
+
   public interface ConnectionFactoryContainer<T extends ConnectionFactoryContainer<T>> {
 
     T withConnectionFactory(ConnectionFactory connectionFactory);
@@ -367,6 +482,21 @@ public Read<T> withConnectionFactoryProviderFn(
       return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build();
     }
 
+    public Read<T> withConnectionConfiguration(ConnectionConfiguration configuration) {
+      checkArgument(configuration != null, "configuration can not be null");
+      Read<T> read =
+          this.withConnectionFactoryProviderFn(
+              (SerializableFunction<Void, ? extends ConnectionFactory>)
+                  __ -> 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 +1236,21 @@ public Write<EventT> withConnectionFactoryProviderFn(
       return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build();
     }
 
+    public Write<EventT> withConnectionConfiguration(ConnectionConfiguration configuration) {
+      checkArgument(configuration != null, "configuration can not be null");
+      Write<EventT> write =
+          this.withConnectionFactoryProviderFn(
+              (SerializableFunction<Void, ? extends ConnectionFactory>)
+                  __ -> 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.