|
56 | 56 | import org.apache.beam.sdk.metrics.Metrics; |
57 | 57 | import org.apache.beam.sdk.options.ExecutorOptions; |
58 | 58 | import org.apache.beam.sdk.options.PipelineOptions; |
| 59 | +import org.apache.beam.sdk.schemas.AutoValueSchema; |
| 60 | +import org.apache.beam.sdk.schemas.annotations.DefaultSchema; |
| 61 | +import org.apache.beam.sdk.schemas.annotations.SchemaFieldDescription; |
59 | 62 | import org.apache.beam.sdk.transforms.DoFn; |
60 | 63 | import org.apache.beam.sdk.transforms.PTransform; |
61 | 64 | import org.apache.beam.sdk.transforms.ParDo; |
|
73 | 76 | import org.apache.beam.sdk.values.TupleTagList; |
74 | 77 | import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting; |
75 | 78 | import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects; |
| 79 | +import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Strings; |
76 | 80 | import org.checkerframework.checker.initialization.qual.Initialized; |
77 | 81 | import org.checkerframework.checker.nullness.qual.Nullable; |
78 | 82 | import org.joda.time.Duration; |
@@ -208,6 +212,117 @@ public static <EventT> Write<EventT> write() { |
208 | 212 | return new AutoValue_JmsIO_Write.Builder<EventT>().build(); |
209 | 213 | } |
210 | 214 |
|
| 215 | + /** A POJO describing a JMS connection. */ |
| 216 | + @DefaultSchema(AutoValueSchema.class) |
| 217 | + @AutoValue |
| 218 | + public abstract static class ConnectionConfiguration implements Serializable { |
| 219 | + public static Builder builder() { |
| 220 | + return new AutoValue_JmsIO_ConnectionConfiguration.Builder(); |
| 221 | + } |
| 222 | + |
| 223 | + public static ConnectionConfiguration create( |
| 224 | + String serverUri, @Nullable String connectionFactoryClassName) { |
| 225 | + checkArgument(serverUri != null, "serverUri can not be null"); |
| 226 | + return builder() |
| 227 | + .setServerUri(serverUri) |
| 228 | + .setConnectionFactoryClassName(connectionFactoryClassName) |
| 229 | + .build(); |
| 230 | + } |
| 231 | + |
| 232 | + public static ConnectionConfiguration create(String serverUri) { |
| 233 | + return create(serverUri, null); |
| 234 | + } |
| 235 | + |
| 236 | + @SchemaFieldDescription("The JMS broker URI.") |
| 237 | + public abstract String getServerUri(); |
| 238 | + |
| 239 | + @SchemaFieldDescription("The JMS ConnectionFactory class name.") |
| 240 | + public abstract @Nullable String getConnectionFactoryClassName(); |
| 241 | + |
| 242 | + @SchemaFieldDescription("The username to connect to the JMS broker.") |
| 243 | + public abstract @Nullable String getUsername(); |
| 244 | + |
| 245 | + @SchemaFieldDescription("The password to connect to the JMS broker.") |
| 246 | + public abstract @Nullable String getPassword(); |
| 247 | + |
| 248 | + public ConnectionConfiguration withUsername(String username) { |
| 249 | + return toBuilder().setUsername(username).build(); |
| 250 | + } |
| 251 | + |
| 252 | + public ConnectionConfiguration withPassword(String password) { |
| 253 | + return toBuilder().setPassword(password).build(); |
| 254 | + } |
| 255 | + |
| 256 | + public ConnectionConfiguration withConnectionFactoryClassName( |
| 257 | + String connectionFactoryClassName) { |
| 258 | + return toBuilder().setConnectionFactoryClassName(connectionFactoryClassName).build(); |
| 259 | + } |
| 260 | + |
| 261 | + abstract Builder toBuilder(); |
| 262 | + |
| 263 | + @AutoValue.Builder |
| 264 | + public abstract static class Builder { |
| 265 | + public abstract Builder setServerUri(String serverUri); |
| 266 | + |
| 267 | + public abstract Builder setConnectionFactoryClassName( |
| 268 | + @Nullable String connectionFactoryClassName); |
| 269 | + |
| 270 | + public abstract Builder setUsername(@Nullable String username); |
| 271 | + |
| 272 | + public abstract Builder setPassword(@Nullable String password); |
| 273 | + |
| 274 | + public abstract ConnectionConfiguration build(); |
| 275 | + } |
| 276 | + |
| 277 | + public ConnectionFactory createConnectionFactory() { |
| 278 | + String className = getConnectionFactoryClassName(); |
| 279 | + // Default to ActiveMQ |
| 280 | + if (Strings.isNullOrEmpty(className)) { |
| 281 | + className = "org.apache.activemq.ActiveMQConnectionFactory"; |
| 282 | + } |
| 283 | + try { |
| 284 | + Class<?> clazz = Class.forName(className); |
| 285 | + String uri = getServerUri(); |
| 286 | + String username = getUsername(); |
| 287 | + String password = getPassword(); |
| 288 | + |
| 289 | + if (username != null && password != null) { |
| 290 | + try { |
| 291 | + return (ConnectionFactory) |
| 292 | + clazz |
| 293 | + .getConstructor(String.class, String.class, String.class) |
| 294 | + .newInstance(username, password, uri); |
| 295 | + } catch (NoSuchMethodException e) { |
| 296 | + // fall through to 1-arg constructor + setters |
| 297 | + } |
| 298 | + } |
| 299 | + try { |
| 300 | + ConnectionFactory cf = |
| 301 | + (ConnectionFactory) clazz.getConstructor(String.class).newInstance(uri); |
| 302 | + if (username != null && password != null) { |
| 303 | + try { |
| 304 | + clazz.getMethod("setUserName", String.class).invoke(cf, username); |
| 305 | + clazz.getMethod("setPassword", String.class).invoke(cf, password); |
| 306 | + } catch (NoSuchMethodException e) { |
| 307 | + try { |
| 308 | + clazz.getMethod("setUsername", String.class).invoke(cf, username); |
| 309 | + clazz.getMethod("setPassword", String.class).invoke(cf, password); |
| 310 | + } catch (NoSuchMethodException e2) { |
| 311 | + // ignore if setters not found |
| 312 | + } |
| 313 | + } |
| 314 | + } |
| 315 | + return cf; |
| 316 | + } catch (NoSuchMethodException e) { |
| 317 | + return (ConnectionFactory) clazz.getConstructor().newInstance(); |
| 318 | + } |
| 319 | + } catch (Exception e) { |
| 320 | + throw new IllegalArgumentException( |
| 321 | + "Unable to instantiate JMS ConnectionFactory of class " + className, e); |
| 322 | + } |
| 323 | + } |
| 324 | + } |
| 325 | + |
211 | 326 | public interface ConnectionFactoryContainer<T extends ConnectionFactoryContainer<T>> { |
212 | 327 |
|
213 | 328 | T withConnectionFactory(ConnectionFactory connectionFactory); |
@@ -367,6 +482,21 @@ public Read<T> withConnectionFactoryProviderFn( |
367 | 482 | return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build(); |
368 | 483 | } |
369 | 484 |
|
| 485 | + public Read<T> withConnectionConfiguration(ConnectionConfiguration configuration) { |
| 486 | + checkArgument(configuration != null, "configuration can not be null"); |
| 487 | + Read<T> read = |
| 488 | + this.withConnectionFactoryProviderFn( |
| 489 | + (SerializableFunction<Void, ? extends ConnectionFactory>) |
| 490 | + __ -> configuration.createConnectionFactory()); |
| 491 | + if (configuration.getUsername() != null) { |
| 492 | + read = read.withUsername(configuration.getUsername()); |
| 493 | + } |
| 494 | + if (configuration.getPassword() != null) { |
| 495 | + read = read.withPassword(configuration.getPassword()); |
| 496 | + } |
| 497 | + return read; |
| 498 | + } |
| 499 | + |
370 | 500 | /** |
371 | 501 | * Specify the JMS queue destination name where to read messages from. The {@link JmsIO.Read} |
372 | 502 | * acts as a consumer on the queue. |
@@ -1106,6 +1236,21 @@ public Write<EventT> withConnectionFactoryProviderFn( |
1106 | 1236 | return builder().setConnectionFactoryProviderFn(connectionFactoryProviderFn).build(); |
1107 | 1237 | } |
1108 | 1238 |
|
| 1239 | + public Write<EventT> withConnectionConfiguration(ConnectionConfiguration configuration) { |
| 1240 | + checkArgument(configuration != null, "configuration can not be null"); |
| 1241 | + Write<EventT> write = |
| 1242 | + this.withConnectionFactoryProviderFn( |
| 1243 | + (SerializableFunction<Void, ? extends ConnectionFactory>) |
| 1244 | + __ -> configuration.createConnectionFactory()); |
| 1245 | + if (configuration.getUsername() != null) { |
| 1246 | + write = write.withUsername(configuration.getUsername()); |
| 1247 | + } |
| 1248 | + if (configuration.getPassword() != null) { |
| 1249 | + write = write.withPassword(configuration.getPassword()); |
| 1250 | + } |
| 1251 | + return write; |
| 1252 | + } |
| 1253 | + |
1109 | 1254 | /** |
1110 | 1255 | * Specify the JMS queue destination name where to send messages to. The {@link JmsIO.Write} |
1111 | 1256 | * acts as a producer on the queue. |
|
0 commit comments