Cannot serialize custom CredentialProvider apache beam

Viewed 177

Im having trouble compiling my apache beam pipeline to a dataflow template. Im using a custom aws credential provider class to load basic aws credentials from GCP secret manager and see the following exception when building the template:

Caused by: com.fasterxml.jackson.databind.JsonMappingException: Unexpected IOException (of type java.io.IOException): Failed to serialize and deserialize property 'awsCredentialsProvider' with value 'org.example.iot.beam.providers.GCPSecretManagerAWSCredentialsProvider@69f90a92'
    at com.fasterxml.jackson.databind.JsonMappingException.fromUnexpectedIOE (JsonMappingException.java:336)
    at com.fasterxml.jackson.databind.ObjectMapper.writeValueAsBytes (ObjectMapper.java:3852)
    at org.apache.beam.runners.dataflow.DataflowRunner.run (DataflowRunner.java:1082)
    at org.apache.beam.runners.dataflow.DataflowRunner.run (DataflowRunner.java:196)
    at org.apache.beam.sdk.Pipeline.run (Pipeline.java:323)
    at org.apache.beam.sdk.Pipeline.run (Pipeline.java:309)
    at org.example.iot.beam.IOTSqsToPubSub.runPipeline (IOTSqsToPubSub.java:30)
    at org.example.iot.beam.IOTSqsToPubSub.main (IOTSqsToPubSub.java:22)
    at sun.reflect.NativeMethodAccessorImpl.invoke0 (Native Method)
    at sun.reflect.NativeMethodAccessorImpl.invoke (NativeMethodAccessorImpl.java:62)
    at sun.reflect.DelegatingMethodAccessorImpl.invoke (DelegatingMethodAccessorImpl.java:43)
    at java.lang.reflect.Method.invoke (Method.java:498)
    at org.codehaus.mojo.exec.ExecJavaMojo$1.run (ExecJavaMojo.java:282)
    at java.lang.Thread.run (Thread.java:748)

The thing that confuses me is that the credential providers, like StaticAWSCredentialsProvider for example, don't specify any json (de)serializers. I would expect my implementation of the credprovider to just work, because its kind of similar:

package org.example.iot.beam.providers;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.auth.BasicAWSCredentials;
import com.google.cloud.secretmanager.v1.SecretManagerServiceClient;
import com.google.cloud.secretmanager.v1.SecretVersionName;

import java.io.IOException;

public class GCPSecretManagerAWSCredentialsProvider implements AWSCredentialsProvider {

    private static final String LATEST_VERSION = "latest";
    private final String secretKeySecretName;
    private final String keyIdSecretName;
    private final String projectId;
    private final String version;
    private AWSCredentials credentials;

    public GCPSecretManagerAWSCredentialsProvider(String secretKeySecretName, String keyIdSecretName, String projectId) {
        this(secretKeySecretName, keyIdSecretName, projectId, LATEST_VERSION);
    }

    public GCPSecretManagerAWSCredentialsProvider(String secretKeySecretName, String keyIdSecretName, String projectId, String version) {
        this.secretKeySecretName = secretKeySecretName;
        this.keyIdSecretName = keyIdSecretName;
        this.projectId = projectId;
        this.version = version;
    }

    @Override
    public AWSCredentials getCredentials() {
        if (this.credentials == null) {
            this.loadCredentialsFromSecretManager();
        }
        return this.credentials;
    }

    public void loadCredentialsFromSecretManager() {
        try (SecretManagerServiceClient client = SecretManagerServiceClient.create()) {
            SecretVersionName keyIdSecretVersionName = SecretVersionName.of(this.projectId, this.keyIdSecretName, this.version);
            String awsKeyId = client.accessSecretVersion(keyIdSecretVersionName).getPayload().getData().toStringUtf8();
            SecretVersionName secretKeySecretVersionName = SecretVersionName.of(this.projectId, this.secretKeySecretName, this.version);
            String awsSecretKey = client.accessSecretVersion(secretKeySecretVersionName).getPayload().getData().toStringUtf8();
            this.credentials = new BasicAWSCredentials(awsKeyId, awsSecretKey);
        } catch (IOException e) {
            throw new RuntimeException(
                    "Error while instantiating secret manager client, check if default authentication chain is available on the machine", e);
        }
    }

    @Override
    public void refresh() {
    }

}

Should I implement a custom (de)serializer for my credential provider, or is there an easier way to get this working?

2 Answers

The exception looks like you are intending to use the credentials provider via pipeline options with SqsIO.read(), right?

For this to work you need a Json de/serializer that is aware of your specific custom provider. There's a Jackson Module (AwsModule.java) that defines a polymorphic de/serializer for AWSCredentialsProvider such as AWSStaticCredentialsProvider and some others. Unfortunately that one is not extensible and you can't simply provide your own de/serializer for GCPSecretManagerAWSCredentialsProvider to make things work.

AwsModule is automatically loaded using the Java ServiceLoader (SPI) and I'm not aware of a way to override / replace the de/serializer for AWSCredentialsProvider with a custom one matching your needs once the module is registered.

The only hack I can think of is to shadow the entire AwsModule (so using the same fully qualified class name) with a module that contains only your custom de/serializer. But it might be a bit of trial error to figure out if anything else in AwsModule is needed...

But I'm wondering, can't you just use the AWSStaticCredentialsProvider and initialise credentials accordingly? Looking at your code above it doesn't look like you have to refresh credentials:

static BasicAWSCredentials loadCredentialsFromSecretManager() {
   ...
}


public static void main(String[] args) {
  AwsOptions options = PipelineOptionsFactory.fromArgs(args).withValidation().as(AwsOptions.class);
  options.setAwsCredentialsProvider(
    new AWSStaticCredentialsProvider(loadCredentialsFromSecretManager())); 

  run(options);
}

I may guess that it's caused by the fact that your GCPSecretManagerAWSCredentialsProvider implements com.amazonaws.auth.AWSCredentialsProvider which is not Serializable. On the other hand, Beam AWS SDK classes, like StaticAwsClientsProvider and similar ones, implement an own AwsClientsProvider interface which is Serializable. So, you may want to try to do the similar thing for your own provider.

Related