Configure DynamoDb TableSchema as a bean (#957)

Fixes #674

Co-authored-by: Gladîș Vladlen <gladis.vladlen+github@gmail.com>

docs/src/main/asciidoc/dynamodb.adoc

@@ -70,6 +70,18 @@ To use a custom implementation, declare a bean of type `DynamoDbTableNameResolve
 
 To resolve a table schema for an entity, `DynamoDbTemplate` uses a bean of type `DynamoDbTableSchemaResolver`. The default implementation caches `TableSchema` objects in internal map.
 To use custom implementation, declare a bean of type `DynamoDbTableSchemaResolver` and it will get injected into `DynamoDbTemplate` automatically during auto-configuration.
+To register a custom table schema for a DynamoDB entity a bean of type `TableSchema`  should be created:
+[source, java]
+----
+@Configuration
+public class MyTableSchemaConfiguration {
+
+    @Bean
+    public TableSchema<MyEntity> myEntityTableSchema() {
+        // create and return a TableSchema object for the MyEntity class
+    }
+}
+----
 
 IMPORTANT: Because of classloader related https://github.com/aws/aws-sdk-java-v2/issues/2604[issue] in AWS SDK DynamoDB Enhanced client, to use Spring Cloud AWS DynamoDB module together with Spring Boot DevTools you must create a custom table schema resolver and define schema using `StaticTableSchema`.
 

spring-cloud-aws-autoconfigure/src/main/java/io/awspring/cloud/autoconfigure/dynamodb/DynamoDbAutoConfiguration.java

@@ -21,7 +21,7 @@
 import io.awspring.cloud.autoconfigure.core.RegionProviderAutoConfiguration;
 import io.awspring.cloud.dynamodb.*;
 import java.io.IOException;
-import java.util.Optional;
+import java.util.List;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
@@ -37,6 +37,7 @@
 import org.springframework.context.annotation.Configuration;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
+import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
 import software.amazon.awssdk.regions.providers.AwsRegionProvider;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClientBuilder;
@@ -47,6 +48,7 @@
  *
  * @author Matej Nedic
  * @author Arun Patra
+ * @author Maciej Walkowiak
  * @since 3.0.0
  */
 @AutoConfiguration
@@ -115,17 +117,23 @@ public DynamoDbEnhancedClient dynamoDbEnhancedClient(DynamoDbClient dynamoDbClie
 		return DynamoDbEnhancedClient.builder().dynamoDbClient(dynamoDbClient).build();
 	}
 
+	@ConditionalOnMissingBean(DynamoDbTableSchemaResolver.class)
+	@Bean
+	public DefaultDynamoDbTableSchemaResolver dynamoDbTableSchemaResolver(List<TableSchema<?>> tableSchemas) {
+		return new DefaultDynamoDbTableSchemaResolver(tableSchemas);
+	}
+
+	@ConditionalOnMissingBean(DynamoDbTableNameResolver.class)
+	@Bean
+	public DefaultDynamoDbTableNameResolver dynamoDbTableNameResolver(DynamoDbProperties properties) {
+		return new DefaultDynamoDbTableNameResolver(properties.getTablePrefix());
+	}
+
 	@ConditionalOnMissingBean(DynamoDbOperations.class)
 	@Bean
-	public DynamoDbTemplate dynamoDBTemplate(DynamoDbProperties properties,
-			DynamoDbEnhancedClient dynamoDbEnhancedClient, Optional<DynamoDbTableSchemaResolver> tableSchemaResolver,
-			Optional<DynamoDbTableNameResolver> tableNameResolver) {
-		DynamoDbTableSchemaResolver tableSchemaRes = tableSchemaResolver
-				.orElseGet(DefaultDynamoDbTableSchemaResolver::new);
-
-		DynamoDbTableNameResolver tableNameRes = tableNameResolver
-				.orElseGet(() -> new DefaultDynamoDbTableNameResolver(properties.getTablePrefix()));
-		return new DynamoDbTemplate(dynamoDbEnhancedClient, tableSchemaRes, tableNameRes);
+	public DynamoDbTemplate dynamoDBTemplate(DynamoDbEnhancedClient dynamoDbEnhancedClient,
+			DynamoDbTableSchemaResolver tableSchemaResolver, DynamoDbTableNameResolver dynamoDbTableNameResolver) {
+		return new DynamoDbTemplate(dynamoDbEnhancedClient, tableSchemaResolver, dynamoDbTableNameResolver);
 	}
 
 	static class MissingDaxUrlCondition extends NoneNestedConditions {

spring-cloud-aws-autoconfigure/src/test/java/io/awspring/cloud/autoconfigure/dynamodb/DynamoDbAutoConfigurationTest.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * Copyright 2013-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 import io.awspring.cloud.autoconfigure.core.AwsClientCustomizer;
 import io.awspring.cloud.autoconfigure.core.CredentialsProviderAutoConfiguration;
 import io.awspring.cloud.autoconfigure.core.RegionProviderAutoConfiguration;
+import io.awspring.cloud.dynamodb.DefaultDynamoDbTableSchemaResolver;
 import io.awspring.cloud.dynamodb.DynamoDbTableNameResolver;
 import io.awspring.cloud.dynamodb.DynamoDbTableSchemaResolver;
 import io.awspring.cloud.dynamodb.DynamoDbTemplate;
@@ -38,6 +39,7 @@
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
+import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.apache.ApacheHttpClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -48,6 +50,7 @@
  * Tests for {@link DynamoDbAutoConfiguration}.
  *
  * @author Matej Nedic
+ * @author Maciej Walkowiak
  */
 class DynamoDbAutoConfigurationTest {
 
@@ -125,6 +128,18 @@ void customDynamoDbClientConfigurer() {
 						assertThat(dynamoDbClient.getSyncHttpClient()).isNotNull();
 					});
 		}
+
+		@Test
+		void tableSchemaBeansRegistered() {
+			contextRunner.withUserConfiguration(DynamoDbAutoConfigurationTest.TableSchemaConfiguration.class)
+					.run(context -> {
+						DefaultDynamoDbTableSchemaResolver schemaResolver = context
+								.getBean(DefaultDynamoDbTableSchemaResolver.class);
+						TableSchema<TableSchemaConfiguration.Person> personTableSchema = schemaResolver
+								.resolve(TableSchemaConfiguration.Person.class);
+						assertThat(context.getBean("personTableSchema")).isEqualTo(personTableSchema);
+					});
+		}
 	}
 
 	@Nested
@@ -254,15 +269,15 @@ DynamoDbTableSchemaResolver tableSchemaResolver() {
 		DynamoDbTableNameResolver tableNameResolver() {
 			return new CustomDynamoDBDynamoDbTableNameResolver();
 		}
+
 	}
 
 	static class CustomDynamoDBDynamoDbTableSchemaResolver implements DynamoDbTableSchemaResolver {
 
 		@Override
-		public <T> TableSchema resolve(Class<T> clazz, String tableName) {
+		public <T> TableSchema resolve(Class<T> clazz) {
 			return null;
 		}
-
 	}
 
 	static class CustomDynamoDBDynamoDbTableNameResolver implements DynamoDbTableNameResolver {
@@ -297,4 +312,16 @@ public SdkHttpClient httpClient() {
 
 	}
 
+	@Configuration(proxyBeanMethods = false)
+	static class TableSchemaConfiguration {
+
+		@Bean
+		TableSchema<Person> personTableSchema() {
+			return StaticTableSchema.builder(Person.class).build();
+		}
+
+		static class Person {
+		}
+	}
+
 }

spring-cloud-aws-dynamodb/src/main/java/io/awspring/cloud/dynamodb/DefaultDynamoDbTableSchemaResolver.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * Copyright 2013-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,6 +15,8 @@
  */
 package io.awspring.cloud.dynamodb;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
@@ -23,12 +25,27 @@
  * Default implementation with simple cache for {@link TableSchema}.
  *
  * @author Matej Nedic
+ * @author Maciej Walkowiak
  */
 public class DefaultDynamoDbTableSchemaResolver implements DynamoDbTableSchemaResolver {
-	private final Map<String, TableSchema> tableSchemaCache = new ConcurrentHashMap<>();
+	private final Map<Class<?>, TableSchema> tableSchemaCache = new ConcurrentHashMap<>();
+
+	public DefaultDynamoDbTableSchemaResolver() {
+		this(Collections.emptyList());
+	}
+
+	public DefaultDynamoDbTableSchemaResolver(List<TableSchema<?>> tableSchemas) {
+		for (TableSchema<?> tableSchema : tableSchemas) {
+			tableSchemaCache.put(tableSchema.itemType().rawClass(), tableSchema);
+		}
+	}
 
 	@Override
-	public <T> TableSchema<T> resolve(Class<T> clazz, String tableName) {
-		return tableSchemaCache.computeIfAbsent(tableName, entityClassName -> TableSchema.fromBean(clazz));
+	public <T> TableSchema<T> resolve(Class<T> clazz) {
+		return tableSchemaCache.computeIfAbsent(clazz, TableSchema::fromBean);
+	}
+
+	Map<Class<?>, TableSchema> getTableSchemaCache() {
+		return tableSchemaCache;
 	}
 }

spring-cloud-aws-dynamodb/src/main/java/io/awspring/cloud/dynamodb/DynamoDbTableSchemaResolver.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2013-2022 the original author or authors.
+ * Copyright 2013-2023 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,21 +18,33 @@
 import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
 
 /**
- * Resolving Class and TableName to {@link TableSchema} class. Should be cached since creating {@link TableSchema} is
- * expensive.
+ * Resolving table schema and table name from a class.
  *
  * @author Matej Nedic
+ * @author Maciej Walkowiak
  * @since 3.0
  */
 public interface DynamoDbTableSchemaResolver {
 
 	/**
-	 * Resolving Class and TableName to {@link TableSchema} class.
+	 * Resolves {@link TableSchema} from {@link Class}.
+	 *
+	 * @param clazz - the class from which table schema is resolved
+	 * @return table schema
+	 * @param <T> - type
+	 */
+	<T> TableSchema<T> resolve(Class<T> clazz);
+
+	/**
+	 * Resolves {@link TableSchema} from {@link Class}.
 	 *
 	 * @param clazz - the class from which table schema is resolved
 	 * @param tableName - the table name
 	 * @return table schema
 	 * @param <T> - type
 	 */
-	<T> TableSchema resolve(Class<T> clazz, String tableName);
+	@Deprecated
+	default <T> TableSchema resolve(Class<T> clazz, String tableName) {
+		return resolve(clazz);
+	}
 }

spring-cloud-aws-dynamodb/src/main/java/io/awspring/cloud/dynamodb/DynamoDbTemplate.java

@@ -15,6 +15,7 @@
  */
 package io.awspring.cloud.dynamodb;
 
+import java.util.Collections;
 import org.springframework.lang.Nullable;
 import org.springframework.util.Assert;
 import software.amazon.awssdk.enhanced.dynamodb.DynamoDbEnhancedClient;
@@ -29,29 +30,29 @@
  *
  * @author Matej Nedic
  * @author Arun Patra
+ * @author Maciej Walkowiak
  * @since 3.0
  */
 public class DynamoDbTemplate implements DynamoDbOperations {
 	private final DynamoDbEnhancedClient dynamoDbEnhancedClient;
 	private final DynamoDbTableSchemaResolver dynamoDbTableSchemaResolver;
-	private final DynamoDbTableNameResolver dynamoDbTableNameResolver;
+	private final DynamoDbTableNameResolver tableNameResolver;
 
 	public DynamoDbTemplate(@Nullable String tablePrefix, DynamoDbEnhancedClient dynamoDbEnhancedClient) {
-		this(dynamoDbEnhancedClient, new DefaultDynamoDbTableSchemaResolver(),
+		this(dynamoDbEnhancedClient, new DefaultDynamoDbTableSchemaResolver(Collections.emptyList()),
 				new DefaultDynamoDbTableNameResolver(tablePrefix));
 	}
 
 	public DynamoDbTemplate(DynamoDbEnhancedClient dynamoDbEnhancedClient) {
-		this(dynamoDbEnhancedClient, new DefaultDynamoDbTableSchemaResolver(),
-				new DefaultDynamoDbTableNameResolver(null));
+		this(dynamoDbEnhancedClient, new DefaultDynamoDbTableSchemaResolver(Collections.emptyList()),
+				new DefaultDynamoDbTableNameResolver());
 	}
 
 	public DynamoDbTemplate(DynamoDbEnhancedClient dynamoDbEnhancedClient,
-			DynamoDbTableSchemaResolver dynamoDbTableSchemaResolver,
-			DynamoDbTableNameResolver dynamoDbTableNameResolver) {
+			DynamoDbTableSchemaResolver dynamoDbTableSchemaResolver, DynamoDbTableNameResolver tableNameResolver) {
 		this.dynamoDbEnhancedClient = dynamoDbEnhancedClient;
 		this.dynamoDbTableSchemaResolver = dynamoDbTableSchemaResolver;
-		this.dynamoDbTableNameResolver = dynamoDbTableNameResolver;
+		this.tableNameResolver = tableNameResolver;
 	}
 
 	public <T> T save(T entity) {
@@ -114,15 +115,15 @@ public <T> PageIterable<T> query(QueryEnhancedRequest queryEnhancedRequest, Clas
 
 	private <T> DynamoDbTable<T> prepareTable(T entity) {
 		Assert.notNull(entity, "entity is required");
-		String tableName = dynamoDbTableNameResolver.resolve(entity.getClass());
-		return dynamoDbEnhancedClient.table(tableName,
-				dynamoDbTableSchemaResolver.resolve(entity.getClass(), tableName));
+		String tableName = tableNameResolver.resolve(entity.getClass());
+		return (DynamoDbTable<T>) dynamoDbEnhancedClient.table(tableName,
+				dynamoDbTableSchemaResolver.resolve(entity.getClass()));
 	}
 
 	private <T> DynamoDbTable<T> prepareTable(Class<T> clazz) {
 		Assert.notNull(clazz, "clazz is required");
-		String tableName = dynamoDbTableNameResolver.resolve(clazz);
-		return dynamoDbEnhancedClient.table(tableName, dynamoDbTableSchemaResolver.resolve(clazz, tableName));
+		String tableName = tableNameResolver.resolve(clazz);
+		return dynamoDbEnhancedClient.table(tableName, dynamoDbTableSchemaResolver.resolve(clazz));
 	}
 
 }