java: Provide an aws-advanced-jdbc-wrapper compatible example

[otaviocarvalho]

Describe the bug

Hi team,

I've been experimenting with Aurora DSQL but I couldn't produce a JDBC/AwsWrapperDataSource/Hikari compatible implementation yet.

I believe it is an important gap to cover, since a JDBC connector is a basic need for JVM-based applications that would be a significant part of the market for Aurora. For example, we are exploring migrating a few use cases from Aurora, which are currently Hikari based.

An Aurora compatible example similar to this one would go a long way: https://github.com/aws/aws-advanced-jdbc-wrapper/blob/main/examples/AWSDriverExample/src/main/java/software/amazon/DatasourceExample.java#L31

Expected Behavior

org.postgresql.util.PSQLException: FATAL: unable to accept connection, access denied Detail: Session Id: <session-id>

Current Behavior

Doesn't connect after providing some based information (similar to what would be needed for Aurora)

Reproduction Steps

@Override
    override fun startUp() {
        AwsCredentialsManager.setCustomHandler(
            AwsCredentialsProviderHandler { hostSpec: HostSpec?, props: Properties? ->
                return@AwsCredentialsProviderHandler CandidateRelationshipMatcherApplicationModule.getCredentialsProvider(config)
            },
        )

        dataSource = AwsWrapperDataSource()
        dataSource.setJdbcProtocol("jdbc:postgresql:")
        dataSource.setTargetDataSourceClassName("org.postgresql.ds.PGSimpleDataSource")

        dataSource.setServerName(config.dbClusterEndpoint)
        dataSource.setDatabase(config.dbName)
        dataSource.setServerPort("5432")

        val wrapperDataSourceProperties = Properties().apply {
            setProperty("wrapperLoggerLevel", "ALL")
            setProperty("useIamAuth", "true")
            setProperty("iamHost", config.dbClusterEndpoint)
            setProperty("iamPort", "5432")
            setProperty("iamRegion", config.awsRegion)
            setProperty("iamUser", config.dbIamUser)
            setProperty(PropertyDefinition.PLUGINS.name, "iam")
            setProperty(PropertyDefinition.USER.name, config.dbIamUser)
        }

        dataSource.setTargetDataSourceProperties(wrapperDataSourceProperties)

        try {
            dataSource.getConnection().use { conn ->
                conn.createStatement().use { statement ->
                    statement.executeQuery("select * from entity_lookup limit 1;").use { result ->
                        println("PostgresAWSRoleService: $result")
                    }
                }
            }
        } catch (e: Exception) {
            System.err.println("PostgresAWSRoleService: Failed to establish initial database connection: ${e.message}")
            throw e // Re-throw to indicate startup failure to AbstractIdleService
        }
        
        println("PostgresAWSRoleService: Service started. JDBI object is ready for use.")
    }

        // We are certain this one works, I am using the same configuration and keys to access Amazon Keyspaces
        fun getCredentialsProvider(config: AWSPostgresConfig): AwsCredentialsProvider {
            val baseProvider = StaticCredentialsProvider.create(
                AwsBasicCredentials.create(config.accessKeyId, config.secretAccessKey),
            )

            return StsAssumeRoleCredentialsProvider.builder()
                .asyncCredentialUpdateEnabled(true)
                .refreshRequest(
                    AssumeRoleRequest.builder()
                        .roleArn("arn:aws:iam::<account-id>:role/entity-platform")
                        .roleSessionName("entity-platform-session-${System.currentTimeMillis()}")
                        .build(),
                )
                .stsClient(
                    StsClient.builder()
                        .credentialsProvider(baseProvider)
                        .region(Region.of(config.awsRegion))
                        .build(),
                )
                .build()
        }

Possible Solution

Provide a JDBC compatible example for Aurora

Additional Information/Context

No response

Environment details (Version of the language and/or framework, OS environment)

OpenJDK Corretto-21.0.5.11.1, Mac Sequoia (15.5 (24F74)), aws-advanced-jdbc-wrapper 2.6.0

[gxjx-x]

Hi @otaviocarvalho feedback received! we're working on getting updated drivers with examples out. In the meantime we will publish a sample that allows you to use HikariCP and pgJDBC standalone. (I have put up PR 183 for this example)
aws-advanced-jdbc-wrapper won't directly work with DSQL and we will add support.

[wcmjunior]

Hi @otaviocarvalho . The PR mentioned above was just merged. Please see the standalone sample for HikariCP and pgJDBC here.

[otaviocarvalho]

Thanks for the swift work on this, @gxjx-x and @wcmjunior!

We will be looking forward to the aws-advanced-jdbc-wrapper implementation. The abstraction of features like AssumeRole authentication are very beneficial for our use case, using it via the library would be more convenient.

By the way, we happened to end up implementing a similar HikariCP extension internally to renew the token. Great minds think alike!

[github-actions]

This issue is now closed. Comments on closed issues are hard for our team to see.
If you need more assistance, please open a new issue that references this one.

[wcmjunior]

@otaviocarvalho we are currently working to add support for DSQL to the aws-advanced-jdbc-wrapper. Feel free to take a look at the PR.