-
Notifications
You must be signed in to change notification settings - Fork 4.2k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
[BEAM-7636] Migrate SqsIO to AWS SDK V2 for Java
- Loading branch information
1 parent
7a8a26b
commit 07d952f
Showing
17 changed files
with
1,166 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
55 changes: 55 additions & 0 deletions
55
...n-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/BasicSqsClientProvider.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import static org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Preconditions.checkArgument; | ||
|
||
import java.net.URI; | ||
import javax.annotation.Nullable; | ||
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; | ||
import software.amazon.awssdk.regions.Region; | ||
import software.amazon.awssdk.services.sqs.SqsClient; | ||
import software.amazon.awssdk.services.sqs.SqsClientBuilder; | ||
|
||
/** Basic implementation of {@link SqsClientProvider} used by default in {@link SqsIO}. */ | ||
class BasicSqsClientProvider implements SqsClientProvider { | ||
private final AwsCredentialsProvider awsCredentialsProvider; | ||
private final String region; | ||
@Nullable private final URI serviceEndpoint; | ||
|
||
BasicSqsClientProvider( | ||
AwsCredentialsProvider awsCredentialsProvider, String region, @Nullable URI serviceEndpoint) { | ||
checkArgument(awsCredentialsProvider != null, "awsCredentialsProvider can not be null"); | ||
checkArgument(region != null, "region can not be null"); | ||
this.awsCredentialsProvider = awsCredentialsProvider; | ||
this.region = region; | ||
this.serviceEndpoint = serviceEndpoint; | ||
} | ||
|
||
@Override | ||
public SqsClient getSqsClient() { | ||
SqsClientBuilder builder = | ||
SqsClient.builder().credentialsProvider(awsCredentialsProvider).region(Region.of(region)); | ||
|
||
if (serviceEndpoint != null) { | ||
builder.endpointOverride(serviceEndpoint); | ||
} | ||
|
||
return builder.build(); | ||
} | ||
} |
50 changes: 50 additions & 0 deletions
50
...a/io/amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/MessageCoder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,50 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.io.OutputStream; | ||
import java.io.Serializable; | ||
import org.apache.beam.sdk.coders.AtomicCoder; | ||
import org.apache.beam.sdk.coders.StringUtf8Coder; | ||
import software.amazon.awssdk.services.sqs.model.Message; | ||
|
||
/** Custom Coder for handling SendMessageRequest for using in Write. */ | ||
public class MessageCoder extends AtomicCoder<Message> implements Serializable { | ||
private static final MessageCoder INSTANCE = new MessageCoder(); | ||
|
||
private MessageCoder() {} | ||
|
||
static MessageCoder of() { | ||
return INSTANCE; | ||
} | ||
|
||
@Override | ||
public void encode(Message value, OutputStream outStream) throws IOException { | ||
StringUtf8Coder.of().encode(value.messageId(), outStream); | ||
StringUtf8Coder.of().encode(value.body(), outStream); | ||
} | ||
|
||
@Override | ||
public Message decode(InputStream inStream) throws IOException { | ||
final String messageId = StringUtf8Coder.of().decode(inStream); | ||
final String body = StringUtf8Coder.of().decode(inStream); | ||
return Message.builder().messageId(messageId).body(body).build(); | ||
} | ||
} |
37 changes: 37 additions & 0 deletions
37
...on-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/MessageCoderRegistrar.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import com.google.auto.service.AutoService; | ||
import java.util.List; | ||
import org.apache.beam.sdk.coders.CoderProvider; | ||
import org.apache.beam.sdk.coders.CoderProviderRegistrar; | ||
import org.apache.beam.sdk.coders.CoderProviders; | ||
import org.apache.beam.sdk.values.TypeDescriptor; | ||
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; | ||
import software.amazon.awssdk.services.sqs.model.Message; | ||
|
||
/** A {@link CoderProviderRegistrar} for standard types used with {@link SqsIO}. */ | ||
@AutoService(CoderProviderRegistrar.class) | ||
public class MessageCoderRegistrar implements CoderProviderRegistrar { | ||
@Override | ||
public List<CoderProvider> getCoderProviders() { | ||
return ImmutableList.of( | ||
CoderProviders.forCoder(TypeDescriptor.of(Message.class), MessageCoder.of())); | ||
} | ||
} |
51 changes: 51 additions & 0 deletions
51
...-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/SendMessageRequestCoder.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,51 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import java.io.IOException; | ||
import java.io.InputStream; | ||
import java.io.OutputStream; | ||
import java.io.Serializable; | ||
import org.apache.beam.sdk.coders.AtomicCoder; | ||
import org.apache.beam.sdk.coders.StringUtf8Coder; | ||
import software.amazon.awssdk.services.sqs.model.SendMessageRequest; | ||
|
||
/** Custom Coder for handling SendMessageRequest for using in Write. */ | ||
public class SendMessageRequestCoder extends AtomicCoder<SendMessageRequest> | ||
implements Serializable { | ||
private static final SendMessageRequestCoder INSTANCE = new SendMessageRequestCoder(); | ||
|
||
private SendMessageRequestCoder() {} | ||
|
||
static SendMessageRequestCoder of() { | ||
return INSTANCE; | ||
} | ||
|
||
@Override | ||
public void encode(SendMessageRequest value, OutputStream outStream) throws IOException { | ||
StringUtf8Coder.of().encode(value.queueUrl(), outStream); | ||
StringUtf8Coder.of().encode(value.messageBody(), outStream); | ||
} | ||
|
||
@Override | ||
public SendMessageRequest decode(InputStream inStream) throws IOException { | ||
final String queueUrl = StringUtf8Coder.of().decode(inStream); | ||
final String message = StringUtf8Coder.of().decode(inStream); | ||
return SendMessageRequest.builder().queueUrl(queueUrl).messageBody(message).build(); | ||
} | ||
} |
38 changes: 38 additions & 0 deletions
38
...ices2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/SendMessageRequestCoderRegistrar.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,38 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import com.google.auto.service.AutoService; | ||
import java.util.List; | ||
import org.apache.beam.sdk.coders.CoderProvider; | ||
import org.apache.beam.sdk.coders.CoderProviderRegistrar; | ||
import org.apache.beam.sdk.coders.CoderProviders; | ||
import org.apache.beam.sdk.values.TypeDescriptor; | ||
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; | ||
import software.amazon.awssdk.services.sqs.model.SendMessageRequest; | ||
|
||
/** A {@link CoderProviderRegistrar} for standard types used with {@link SqsIO}. */ | ||
@AutoService(CoderProviderRegistrar.class) | ||
public class SendMessageRequestCoderRegistrar implements CoderProviderRegistrar { | ||
@Override | ||
public List<CoderProvider> getCoderProviders() { | ||
return ImmutableList.of( | ||
CoderProviders.forCoder( | ||
TypeDescriptor.of(SendMessageRequest.class), SendMessageRequestCoder.of())); | ||
} | ||
} |
64 changes: 64 additions & 0 deletions
64
...amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/SqsCheckpointMark.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,64 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import java.io.Serializable; | ||
import java.util.Collection; | ||
import java.util.List; | ||
import java.util.Optional; | ||
import org.apache.beam.sdk.io.UnboundedSource; | ||
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.base.Objects; | ||
import org.apache.beam.vendor.guava.v26_0_jre.com.google.common.collect.ImmutableList; | ||
import software.amazon.awssdk.services.sqs.model.Message; | ||
|
||
class SqsCheckpointMark implements UnboundedSource.CheckpointMark, Serializable { | ||
|
||
private final List<Message> messagesToDelete; | ||
private final transient Optional<SqsUnboundedReader> reader; | ||
|
||
SqsCheckpointMark(SqsUnboundedReader reader, Collection<Message> messagesToDelete) { | ||
this.reader = Optional.of(reader); | ||
this.messagesToDelete = ImmutableList.copyOf(messagesToDelete); | ||
} | ||
|
||
@Override | ||
public void finalizeCheckpoint() { | ||
reader.ifPresent(r -> r.delete(messagesToDelete)); | ||
} | ||
|
||
List<Message> getMessagesToDelete() { | ||
return messagesToDelete; | ||
} | ||
|
||
@Override | ||
public boolean equals(Object o) { | ||
if (this == o) { | ||
return true; | ||
} | ||
if (o == null || getClass() != o.getClass()) { | ||
return false; | ||
} | ||
SqsCheckpointMark that = (SqsCheckpointMark) o; | ||
return Objects.equal(messagesToDelete, that.messagesToDelete); | ||
} | ||
|
||
@Override | ||
public int hashCode() { | ||
return Objects.hashCode(messagesToDelete); | ||
} | ||
} |
31 changes: 31 additions & 0 deletions
31
...amazon-web-services2/src/main/java/org/apache/beam/sdk/io/aws2/sqs/SqsClientProvider.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
/* | ||
* 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.aws2.sqs; | ||
|
||
import java.io.Serializable; | ||
import software.amazon.awssdk.services.sqs.SqsClient; | ||
|
||
/** | ||
* Provides instances of Sqs clients. | ||
* | ||
* <p>Please note, that any instance of {@link SqsClientProvider} must be {@link Serializable} to | ||
* ensure it can be sent to worker machines. | ||
*/ | ||
public interface SqsClientProvider extends Serializable { | ||
SqsClient getSqsClient(); | ||
} |
Oops, something went wrong.