Configure DynamoDb TableSchema as a bean.

Fixes #674

Author: maciejwalkowiak

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

@@ -22,7 +22,6 @@
 import io.awspring.cloud.dynamodb.*;
 import java.io.IOException;
 import java.util.List;
-import java.util.Optional;
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.AutoConfiguration;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
@@ -48,6 +47,7 @@
  * {@link EnableAutoConfiguration Auto-configuration} for DynamoDB integration.
  *
  * @author Matej Nedic
+ * @author Maciej Walkowiak
  * @since 3.0.0
  */
 @AutoConfiguration
@@ -116,27 +116,24 @@ public DynamoDbEnhancedClient dynamoDbEnhancedClient(DynamoDbClient dynamoDbClie
 		return DynamoDbEnhancedClient.builder().dynamoDbClient(dynamoDbClient).build();
 	}
 
-	@ConditionalOnMissingBean(DynamoDbOperations.class)
+	@ConditionalOnMissingBean(DynamoDbTableNameResolver.class)
 	@Bean
-	public DynamoDbTemplate dynamoDBTemplate(DynamoDbEnhancedClient dynamoDbEnhancedClient,
-			Optional<DynamoDbTableSchemaResolver> tableSchemaResolver,
-			Optional<DynamoDbTableNameResolver> tableNameResolver, List<TableSchema<?>> tableSchemas) {
-		DynamoDbTableNameResolver tableNameRes = tableNameResolver.orElseGet(DefaultDynamoDbTableNameResolver::new);
-
-		DynamoDbTableSchemaResolver tableSchemaRes = tableSchemaResolver
-				.orElseGet(DefaultDynamoDbTableSchemaResolver::new);
-
-		registerTableSchemas(tableSchemas, tableNameRes, tableSchemaRes);
+	public DefaultDynamoDbTableResolver dynamoDbTableSchemaResolver(
+			DefaultDynamoDbTableNameResolver dynamoDbTableNameResolver, List<TableSchema<?>> tableSchemas) {
+		return new DefaultDynamoDbTableResolver(dynamoDbTableNameResolver, tableSchemas);
+	}
 
-		return new DynamoDbTemplate(dynamoDbEnhancedClient, tableSchemaRes, tableNameRes);
+	@ConditionalOnMissingBean(DynamoDbTableNameResolver.class)
+	@Bean
+	public DefaultDynamoDbTableNameResolver dynamoDbTableNameResolver() {
+		return new DefaultDynamoDbTableNameResolver();
 	}
 
-	private void registerTableSchemas(List<TableSchema<?>> tableSchemas, DynamoDbTableNameResolver tableNameRes,
-			DynamoDbTableSchemaResolver tableSchemaRes) {
-		tableSchemas.forEach(schema -> {
-			var tableName = tableNameRes.resolve(schema.itemType().rawClass());
-			tableSchemaRes.register(schema, tableName);
-		});
+	@ConditionalOnMissingBean(DynamoDbOperations.class)
+	@Bean
+	public DynamoDbTemplate dynamoDBTemplate(DynamoDbEnhancedClient dynamoDbEnhancedClient,
+			DynamoDbTableResolver tableSchemaResolver) {
+		return new DynamoDbTemplate(dynamoDbEnhancedClient, tableSchemaResolver);
 	}
 
 	static class MissingDaxUrlCondition extends NoneNestedConditions {

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

@@ -16,16 +16,15 @@
 package io.awspring.cloud.autoconfigure.dynamodb;
 
 import static org.assertj.core.api.Assertions.assertThat;
-import static org.mockito.Mockito.spy;
-import static org.mockito.Mockito.verify;
 
 import io.awspring.cloud.autoconfigure.ConfiguredAwsClient;
 import io.awspring.cloud.autoconfigure.core.AwsAutoConfiguration;
 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.DefaultDynamoDbTableResolver;
 import io.awspring.cloud.dynamodb.DynamoDbTableNameResolver;
-import io.awspring.cloud.dynamodb.DynamoDbTableSchemaResolver;
+import io.awspring.cloud.dynamodb.DynamoDbTableResolver;
 import io.awspring.cloud.dynamodb.DynamoDbTemplate;
 import java.net.URI;
 import java.time.Duration;
@@ -51,6 +50,7 @@
  * Tests for {@link DynamoDbAutoConfiguration}.
  *
  * @author Matej Nedic
+ * @author Maciej Walkowiak
  */
 class DynamoDbAutoConfigurationTest {
 
@@ -71,15 +71,14 @@ void dynamoDBAutoConfigurationIsDisabled() {
 		@Test
 		void customTableResolverResolverCanBeConfigured() {
 			contextRunner.withUserConfiguration(CustomDynamoDbConfiguration.class).run(context -> {
-				DynamoDbTableSchemaResolver dynamoDbTableSchemaResolver = context
-						.getBean(DynamoDbTableSchemaResolver.class);
+				DynamoDbTableResolver dynamoDbTableResolver = context.getBean(DynamoDbTableResolver.class);
 				DynamoDbTableNameResolver dynamoDBDynamoDbTableNameResolver = context
 						.getBean(DynamoDbTableNameResolver.class);
 
-				assertThat(dynamoDbTableSchemaResolver).isNotNull();
+				assertThat(dynamoDbTableResolver).isNotNull();
 				assertThat(dynamoDBDynamoDbTableNameResolver).isNotNull();
 
-				assertThat(dynamoDbTableSchemaResolver).isInstanceOf(CustomDynamoDBDynamoDbTableSchemaResolver.class);
+				assertThat(dynamoDbTableResolver).isInstanceOf(CustomDynamoDBDynamoDbTableResolver.class);
 				assertThat(dynamoDBDynamoDbTableNameResolver)
 						.isInstanceOf(CustomDynamoDBDynamoDbTableNameResolver.class);
 
@@ -130,29 +129,16 @@ void customDynamoDbClientConfigurer() {
 		}
 
 		@Test
-		void tableSchemaBeansGetRegistered() {
-
-			contextRunner.withUserConfiguration(TableSchemaConfiguration.class).run(context -> {
-				var schemaResolver = context.getBean(DynamoDbTableSchemaResolver.class);
-				var nameResolver = context.getBean(DynamoDbTableNameResolver.class);
-
-				assertThat(schemaResolver).isNotNull();
-				assertThat(nameResolver).isNotNull();
-
-				assertTableSchemaRegistered(context.getBean("firstTableSchema", TableSchema.class), nameResolver,
-						schemaResolver);
-				assertTableSchemaRegistered(context.getBean("secondTableSchema", TableSchema.class), nameResolver,
-						schemaResolver);
-
-			});
-		}
-
-		private static void assertTableSchemaRegistered(TableSchema<?> tableSchema,
-				DynamoDbTableNameResolver nameResolver, DynamoDbTableSchemaResolver schemaResolver) {
-			var tableName = verify(nameResolver).resolve(tableSchema.itemType().rawClass());
-			verify(schemaResolver).register(tableSchema, tableName);
+		void tableSchemaBeansRegistered() {
+			contextRunner.withUserConfiguration(DynamoDbAutoConfigurationTest.TableSchemaConfiguration.class)
+					.run(context -> {
+						DefaultDynamoDbTableResolver schemaResolver = context
+								.getBean(DefaultDynamoDbTableResolver.class);
+						TableSchema<TableSchemaConfiguration.Person> personTableSchema = schemaResolver
+								.resolveTableSchema(TableSchemaConfiguration.Person.class);
+						assertThat(context.getBean("personTableSchema")).isEqualTo(personTableSchema);
+					});
 		}
-
 	}
 
 	@Nested
@@ -252,16 +238,14 @@ void customTableResolverResolverCanBeConfigured() {
 					.withPropertyValues(
 							"spring.cloud.aws.dynamodb.dax.url:dax://something.dax-clusters.us-east-1.amazonaws.com")
 					.withUserConfiguration(CustomDynamoDbConfiguration.class).run(context -> {
-						DynamoDbTableSchemaResolver dynamoDbTableSchemaResolver = context
-								.getBean(DynamoDbTableSchemaResolver.class);
+						DynamoDbTableResolver dynamoDbTableResolver = context.getBean(DynamoDbTableResolver.class);
 						DynamoDbTableNameResolver dynamoDBDynamoDbTableNameResolver = context
 								.getBean(DynamoDbTableNameResolver.class);
 
-						assertThat(dynamoDbTableSchemaResolver).isNotNull();
+						assertThat(dynamoDbTableResolver).isNotNull();
 						assertThat(dynamoDBDynamoDbTableNameResolver).isNotNull();
 
-						assertThat(dynamoDbTableSchemaResolver)
-								.isInstanceOf(CustomDynamoDBDynamoDbTableSchemaResolver.class);
+						assertThat(dynamoDbTableResolver).isInstanceOf(CustomDynamoDBDynamoDbTableResolver.class);
 						assertThat(dynamoDBDynamoDbTableNameResolver)
 								.isInstanceOf(CustomDynamoDBDynamoDbTableNameResolver.class);
 
@@ -274,8 +258,8 @@ void customTableResolverResolverCanBeConfigured() {
 	static class CustomDynamoDbConfiguration {
 
 		@Bean
-		DynamoDbTableSchemaResolver tableSchemaResolver() {
-			return new CustomDynamoDBDynamoDbTableSchemaResolver();
+		DynamoDbTableResolver tableSchemaResolver() {
+			return new CustomDynamoDBDynamoDbTableResolver();
 		}
 
 		@Bean
@@ -285,10 +269,15 @@ DynamoDbTableNameResolver tableNameResolver() {
 
 	}
 
-	static class CustomDynamoDBDynamoDbTableSchemaResolver implements DynamoDbTableSchemaResolver {
+	static class CustomDynamoDBDynamoDbTableResolver implements DynamoDbTableResolver {
 
 		@Override
-		public <T> TableSchema resolve(Class<T> clazz, String tableName) {
+		public <T> TableSchema resolveTableSchema(Class<T> clazz) {
+			return null;
+		}
+
+		@Override
+		public String resolveTableName(Class<?> aClass) {
 			return null;
 		}
 
@@ -330,25 +319,12 @@ public SdkHttpClient httpClient() {
 	static class TableSchemaConfiguration {
 
 		@Bean
-		TableSchema<Object> firstTableSchema() {
-			return StaticTableSchema.builder(Object.class).build();
-		}
-
-		@Bean
-		TableSchema<String> secondTableSchema() {
-			return StaticTableSchema.builder(String.class).build();
+		TableSchema<Person> personTableSchema() {
+			return StaticTableSchema.builder(Person.class).build();
 		}
 
-		@Bean
-		DynamoDbTableSchemaResolver schemaResolver() {
-			return spy(CustomDynamoDBDynamoDbTableSchemaResolver.class);
+		static class Person {
 		}
-
-		@Bean
-		DynamoDbTableNameResolver tableNameResolver() {
-			return spy(CustomDynamoDBDynamoDbTableNameResolver.class);
-		}
-
 	}
 
 }

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

@@ -0,0 +1,60 @@
+/*
+ * Copyright 2013-2022 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.
+ * You may obtain a copy of the License at
+ *
+ *      https://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 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;
+
+/**
+ * Default implementation with simple cache for {@link TableSchema}.
+ *
+ * @author Matej Nedic
+ * @author Maciej Walkowiak
+ */
+public class DefaultDynamoDbTableResolver implements DynamoDbTableResolver {
+	private final DynamoDbTableNameResolver tableNameResolver;
+	private final List<TableSchema<?>> tableSchemas;
+	private final Map<String, TableSchema> tableSchemaCache = new ConcurrentHashMap<>();
+
+	DefaultDynamoDbTableResolver() {
+		this(new DefaultDynamoDbTableNameResolver(), Collections.emptyList());
+	}
+
+	public DefaultDynamoDbTableResolver(DynamoDbTableNameResolver tableNameResolver,
+			List<TableSchema<?>> tableSchemas) {
+		this.tableNameResolver = tableNameResolver;
+		this.tableSchemas = tableSchemas;
+	}
+
+	public <T> TableSchema<T> resolveTableSchema(Class<T> clazz) {
+		String tableName = tableNameResolver.resolve(clazz);
+		return tableSchemaCache.computeIfAbsent(tableName,
+				entityClassName -> tableSchemas.stream().filter(it -> it.itemType().rawClass().equals(clazz))
+						.findFirst().orElseGet(() -> TableSchema.fromBean(clazz)));
+	}
+
+	@Override
+	public String resolveTableName(Class<?> aClass) {
+		return tableNameResolver.resolve(aClass);
+	}
+
+	Map<String, TableSchema> getTableSchemaCache() {
+		return tableSchemaCache;
+	}
+}

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

@@ -18,31 +18,28 @@
 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 {
+public interface DynamoDbTableResolver {
 
 	/**
-	 * Resolving Class and TableName to {@link TableSchema} class.
+	 * Resolving {@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);
+	<T> TableSchema<T> resolveTableSchema(Class<T> clazz);
 
 	/**
-	 * Register manually a {@link TableSchema}
+	 * Resolves table name from {@link Class}.
 	 *
-	 * @param tableSchema - the schema to be registered
-	 * @param tableName - the table name
+	 * @param clazz - a class
+	 * @return a table name
 	 */
-	default <T> void register(TableSchema<T> tableSchema, String tableName) {
-	}
-
+	String resolveTableName(Class<?> clazz);
 }