Allow configuration to be serialized

apache · Dec 20, 2024 · a4aa6f1 · a4aa6f1
1 parent 3f01f70
commit a4aa6f1
Showing 1 changed file with 24 additions and 11 deletions.
diff --git a/...ervices2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisTransformRegistrar.java b/...ervices2/src/main/java/org/apache/beam/sdk/io/aws2/kinesis/KinesisTransformRegistrar.java
@@ -62,7 +62,7 @@ private abstract static class CrossLanguageConfiguration {
     String awsAccessKey;
     String awsSecretKey;
     Region region;
-    @Nullable URI serviceEndpoint;
+    @Nullable String serviceEndpoint;
 
     public void setStreamName(String streamName) {
       this.streamName = streamName;
@@ -81,14 +81,7 @@ public void setRegion(String region) {
     }
 
     public void setServiceEndpoint(@Nullable String serviceEndpoint) {
-      if (serviceEndpoint != null) {
-        try {
-          this.serviceEndpoint = new URI(serviceEndpoint);
-        } catch (URISyntaxException ex) {
-          throw new RuntimeException(
-              String.format("Service endpoint must be URI format, got: %s", serviceEndpoint));
-        }
-      }
+      this.serviceEndpoint = serviceEndpoint;
     }
   }
 
@@ -111,14 +104,24 @@ public PTransform<PCollection<byte[]>, KinesisIO.Write.Result> buildExternal(
           AwsBasicCredentials.create(configuration.awsAccessKey, configuration.awsSecretKey);
       StaticCredentialsProvider provider = StaticCredentialsProvider.create(creds);
       SerializableFunction<byte[], byte[]> serializer = v -> v;
+      @Nullable URI endpoint;
+      if (configuration.serviceEndpoint != null) {
+        try {
+          endpoint = configuration.serviceEndpoint;
+        }
+        catch (URISyntaxException ex) {
+          throw new RuntimeException(
+              String.format("Service endpoint must be URI format, got: %s", serviceEndpoint));
+        }
+      }
       KinesisIO.Write<byte[]> writeTransform =
           KinesisIO.<byte[]>write()
               .withStreamName(configuration.streamName)
               .withClientConfiguration(
                   ClientConfiguration.builder()
                       .credentialsProvider(provider)
                       .region(configuration.region)
-                      .endpoint(configuration.serviceEndpoint)
+                      .endpoint(endpoint)
                       .build())
               .withPartitioner(p -> configuration.partitionKey)
               .withSerializer(serializer);
@@ -211,14 +214,24 @@ public PTransform<PBegin, PCollection<byte[]>> buildExternal(
       AwsBasicCredentials creds =
           AwsBasicCredentials.create(configuration.awsAccessKey, configuration.awsSecretKey);
       StaticCredentialsProvider provider = StaticCredentialsProvider.create(creds);
+      @Nullable URI endpoint;
+      if (configuration.serviceEndpoint != null) {
+        try {
+          endpoint = configuration.serviceEndpoint;
+        }
+        catch (URISyntaxException ex) {
+          throw new RuntimeException(
+              String.format("Service endpoint must be URI format, got: %s", serviceEndpoint));
+        }
+      }
       KinesisIO.Read readTransform =
           KinesisIO.read()
               .withStreamName(configuration.streamName)
               .withClientConfiguration(
                   ClientConfiguration.builder()
                       .credentialsProvider(provider)
                       .region(configuration.region)
-                      .endpoint(configuration.serviceEndpoint)
+                      .endpoint(endpoint)
                       .build());
 
       if (configuration.maxNumRecords != null) {