DynamoDBRepositoryConfigExtension.java

/*
 * Copyright © 2018 spring-data-dynamodb (https://github.com/prasanna0586/spring-data-dynamodb)
 *
 * 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
 *
 *     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.socialsignin.spring.data.dynamodb.repository.config;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.socialsignin.spring.data.dynamodb.core.DynamoDBTemplate;
import org.socialsignin.spring.data.dynamodb.core.MarshallingMode;
import org.socialsignin.spring.data.dynamodb.mapping.DynamoDBMappingContext;
import org.socialsignin.spring.data.dynamodb.repository.DynamoDBCrudRepository;
import org.socialsignin.spring.data.dynamodb.repository.DynamoDBPagingAndSortingRepository;
import org.socialsignin.spring.data.dynamodb.repository.support.DynamoDBRepositoryFactoryBean;
import org.socialsignin.spring.data.dynamodb.repository.util.DynamoDBMappingContextProcessor;
import org.socialsignin.spring.data.dynamodb.repository.util.Entity2DynamoDBTableSynchronizer;
import org.springframework.beans.factory.support.AbstractBeanDefinition;
import org.springframework.beans.factory.support.BeanDefinitionBuilder;
import org.springframework.beans.factory.support.BeanDefinitionRegistry;
import org.springframework.core.annotation.AnnotationAttributes;
import org.springframework.data.config.ParsingUtils;
import org.springframework.data.repository.config.AnnotationRepositoryConfigurationSource;
import org.springframework.data.repository.config.RepositoryConfigurationExtensionSupport;
import org.springframework.data.repository.config.RepositoryConfigurationSource;
import org.springframework.data.repository.config.XmlRepositoryConfigurationSource;
import org.springframework.lang.NonNull;
import org.springframework.util.Assert;
import org.springframework.util.StringUtils;
import org.w3c.dom.Element;
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbBean;
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.DynamoDbImmutable;

import java.lang.annotation.Annotation;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

/**
 * Spring Data repository configuration extension for DynamoDB.
 *
 * Extends the base Spring Data configuration extension to support DynamoDB-specific
 * repository configuration including bean registration and post-processing.
 * @author Prasanna Kumar Ramachandran
 */
public class DynamoDBRepositoryConfigExtension extends RepositoryConfigurationExtensionSupport {

    /**
     * Default constructor.
     */
    public DynamoDBRepositoryConfigExtension() {
    }

    private static final Logger LOGGER = LoggerFactory.getLogger(DynamoDBRepositoryConfigExtension.class);

    private static final String DEFAULT_AMAZON_DYNAMO_DB_BEAN_NAME = "amazonDynamoDB";

    private static final String DYNAMO_DB_MAPPER_CONFIG_REF = "dynamodb-mapper-config-ref";

    private static final String DYNAMO_DB_OPERATIONS_REF = "dynamodb-operations-ref";

    private static final String AMAZON_DYNAMODB_REF = "amazon-dynamodb-ref";

    private static final String MAPPING_CONTEXT_REF = "mapping-context-ref";

    private BeanDefinitionRegistry registry;
    private String defaultDynamoDBMappingContext;
    @NonNull
    private MarshallingMode marshallingMode = MarshallingMode.SDK_V2_NATIVE;

    @NonNull
    @Override
    public String getRepositoryFactoryBeanClassName() {
        return DynamoDBRepositoryFactoryBean.class.getName();
    }

    @NonNull
    @Override
    protected Collection<Class<?>> getIdentifyingTypes() {
        return List.of(DynamoDBPagingAndSortingRepository.class, DynamoDBCrudRepository.class);
    }

    @NonNull
    @Override
    protected Collection<Class<? extends Annotation>> getIdentifyingAnnotations() {
        return List.of(DynamoDbBean.class, DynamoDbImmutable.class);
    }

    @Override
    public void postProcess(@NonNull BeanDefinitionBuilder builder, @NonNull AnnotationRepositoryConfigurationSource config) {
        AnnotationAttributes attributes = config.getAttributes();

        String repositoryBeanName = config.generateBeanName(builder.getBeanDefinition());

        postProcess(builder, attributes.getString("amazonDynamoDBRef"),
                attributes.getString("dynamoDBMapperRef"), attributes.getString("dynamoDBMapperConfigRef"),
                attributes.getString("dynamoDBOperationsRef"), attributes.getString("mappingContextRef"));
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.config. RepositoryConfigurationExtensionSupport
     * #postProcess(org.springframework.beans .factory.support.BeanDefinitionBuilder,
     * org.springframework.data.repository .config.XmlRepositoryConfigurationSource)
     */
    @Override
    public void postProcess(@NonNull BeanDefinitionBuilder builder, @NonNull XmlRepositoryConfigurationSource config) {
        Element element = config.getElement();

        ParsingUtils.setPropertyReference(builder, element, AMAZON_DYNAMODB_REF, "amazonDynamoDB");
        ParsingUtils.setPropertyReference(builder, element, DYNAMO_DB_MAPPER_CONFIG_REF, "dynamoDBMapperConfig");
        ParsingUtils.setPropertyReference(builder, element, DYNAMO_DB_OPERATIONS_REF, "dynamoDBOperations");

        String dynamoDBMappingContextRef = element.getAttribute(MAPPING_CONTEXT_REF);

        if (!StringUtils.hasText(dynamoDBMappingContextRef)) {
            // Register DynamoDBMappingContext only once if necessary
            if (defaultDynamoDBMappingContext == null) {
                defaultDynamoDBMappingContext = registerDynamoDBMappingContext(registry);
            }
            dynamoDBMappingContextRef = defaultDynamoDBMappingContext;

        }
        registerAndSetPostProcessingBeans(builder, registry, dynamoDBMappingContextRef);
    }

    private final Map<String, String> dynamoDBTemplateCache = new HashMap<>();

    private void postProcess(@NonNull BeanDefinitionBuilder builder, String amazonDynamoDBRef,
                             @NonNull String dynamoDBMapperRef, @NonNull String dynamoDBMapperConfigRef, @NonNull String dynamoDBOperationsRef,
                             String dynamoDBMappingContextRef) {

        // Ensure mapping context is created first
        if (!StringUtils.hasText(dynamoDBMappingContextRef)) {
            // Register DynamoDBMappingContext only once if necessary
            if (defaultDynamoDBMappingContext == null) {
                defaultDynamoDBMappingContext = registerDynamoDBMappingContext(registry);
            }
            dynamoDBMappingContextRef = defaultDynamoDBMappingContext;
        }

        if (StringUtils.hasText(dynamoDBOperationsRef)) {
            builder.addPropertyReference("dynamoDBOperations", dynamoDBOperationsRef);
            Assert.isTrue(!StringUtils.hasText(amazonDynamoDBRef),
                    "Cannot specify both amazonDynamoDB bean and dynamoDBOperationsBean in repository configuration");
            Assert.isTrue(!StringUtils.hasText(dynamoDBMapperConfigRef),
                    "Cannot specify both dynamoDBMapperConfigBean bean and dynamoDBOperationsBean in repository configuration");
        } else {

            if (!StringUtils.hasText(dynamoDBOperationsRef)) {

                String dynamoDBRef;
                if (StringUtils.hasText(amazonDynamoDBRef)) {
                    dynamoDBRef = amazonDynamoDBRef;
                } else {
                    dynamoDBRef = DEFAULT_AMAZON_DYNAMO_DB_BEAN_NAME;
                }

                final String finalDynamoDBMappingContextRef = dynamoDBMappingContextRef;
                dynamoDBOperationsRef = dynamoDBTemplateCache
                        .computeIfAbsent(getBeanNameWithModulePrefix("DynamoDBTemplate-" + dynamoDBRef), ref -> {
                            BeanDefinitionBuilder dynamoDBTemplateBuilder = BeanDefinitionBuilder
                                    .genericBeanDefinition(DynamoDBTemplate.class);
                            // DynamoDbClient, DynamoDBMapper, DynamoDBMapperConfig, DynamoDBMappingContext
                            dynamoDBTemplateBuilder.addConstructorArgReference(dynamoDBRef);

                            if (StringUtils.hasText(dynamoDBMapperRef)) {
                                dynamoDBTemplateBuilder.addConstructorArgReference(dynamoDBMapperRef);
                            } else {
                                dynamoDBTemplateBuilder.addConstructorArgReference(this.dynamoDBMapperName);
                            }

                            if (StringUtils.hasText(dynamoDBMapperConfigRef)) {
                                dynamoDBTemplateBuilder.addConstructorArgReference(dynamoDBMapperConfigRef);
                            } else {
                                dynamoDBTemplateBuilder.addConstructorArgReference(this.dynamoDBMapperConfigName);
                            }

                            // Add mapping context as fourth constructor argument
                            dynamoDBTemplateBuilder.addConstructorArgReference(finalDynamoDBMappingContextRef);

                            registry.registerBeanDefinition(ref, dynamoDBTemplateBuilder.getBeanDefinition());
                            return ref;
                        });
            }

            builder.addPropertyReference("dynamoDBOperations", dynamoDBOperationsRef);

        }

        builder.addPropertyReference("dynamoDBMappingContext", dynamoDBMappingContextRef);
        registerAndSetPostProcessingBeans(builder, registry, dynamoDBMappingContextRef);
    }

    /**
     * Registers and sets post-processing beans for DynamoDB repository configuration.
     * @param builder the bean definition builder
     * @param registry the bean definition registry
     * @param dynamoDBMappingContextRef the DynamoDB mapping context reference
     */
    protected void registerAndSetPostProcessingBeans(@NonNull BeanDefinitionBuilder builder, @NonNull BeanDefinitionRegistry registry,
                                                     @NonNull String dynamoDBMappingContextRef) {
        String tableSynchronizerName = registerEntity2DynamoDBTableSynchronizer(registry, dynamoDBMappingContextRef);
        builder.addPropertyReference("entity2DynamoDBTableSynchronizer", tableSynchronizerName);

        String dynamoDBMappingContextProcessorName = registerDynamoDBMappingContextProcessor(registry,
                dynamoDBMappingContextRef);
        builder.addPropertyReference("dynamoDBMappingContextProcessor", dynamoDBMappingContextProcessorName);
    }

    private final Map<String, String> entity2DynamoDBTableSynchronizerCache = new ConcurrentHashMap<>();

    @NonNull
    private String registerEntity2DynamoDBTableSynchronizer(@NonNull BeanDefinitionRegistry registry,
                                                            String dynamoDBMappingContextRef) {

        return entity2DynamoDBTableSynchronizerCache.computeIfAbsent(dynamoDBMappingContextRef, ref -> {
            BeanDefinitionBuilder entity2DynamoDBTableSynchronizerBuilder = BeanDefinitionBuilder
                    .genericBeanDefinition(Entity2DynamoDBTableSynchronizer.class);
            entity2DynamoDBTableSynchronizerBuilder.setAutowireMode(AbstractBeanDefinition.AUTOWIRE_CONSTRUCTOR);
            String tableSynchronizerName = getBeanNameWithModulePrefix(
                    "Entity2DynamoDBTableSynchronizer-" + dynamoDBMappingContextRef);
            registry.registerBeanDefinition(tableSynchronizerName,
                    entity2DynamoDBTableSynchronizerBuilder.getBeanDefinition());

            return tableSynchronizerName;
        });
    }

    private final Map<String, String> dynamoDBMappingContextProcessorCache = new ConcurrentHashMap<>();

    private String dynamoDBMapperName;

    private String dynamoDBMapperConfigName;

    @NonNull
    private String registerDynamoDBMappingContextProcessor(@NonNull BeanDefinitionRegistry registry,
                                                           @NonNull String dynamoDBMappingContextRef) {

        return dynamoDBMappingContextProcessorCache.computeIfAbsent(dynamoDBMappingContextRef, ref -> {
            BeanDefinitionBuilder dynamoDBMappingContextProcessorBuilder = BeanDefinitionBuilder
                    .genericBeanDefinition(DynamoDBMappingContextProcessor.class);
            dynamoDBMappingContextProcessorBuilder.addConstructorArgReference(dynamoDBMappingContextRef);

            String dynamoDBMappingContextProcessorRef = getBeanNameWithModulePrefix(
                    "DynamoDBMappingContextProcessor-" + dynamoDBMappingContextRef);
            registry.registerBeanDefinition(dynamoDBMappingContextProcessorRef,
                    dynamoDBMappingContextProcessorBuilder.getBeanDefinition());

            return dynamoDBMappingContextProcessorRef;

        });
    }

    @NonNull
    private String registerDynamoDBMappingContext(@NonNull BeanDefinitionRegistry registry) {

        // Use the standard bean name to ensure compatibility with @EnableDynamoDBAuditing
        String dynamoDBMappingContextRef = org.socialsignin.spring.data.dynamodb.config.BeanNames.MAPPING_CONTEXT_BEAN_NAME;

        // Check if the bean already exists (might be provided by user or @EnableDynamoDBAuditing)
        if (!registry.containsBeanDefinition(dynamoDBMappingContextRef)) {
            BeanDefinitionBuilder dynamoDBMappingContextBuilder = BeanDefinitionBuilder
                    .genericBeanDefinition(DynamoDBMappingContext.class);

            // Pass marshalling mode as constructor argument
            dynamoDBMappingContextBuilder.addConstructorArgValue(marshallingMode);

            LOGGER.debug("Registering DynamoDBMappingContext bean <{}> with marshalling mode <{}>",
                    dynamoDBMappingContextRef, marshallingMode);

            registry.registerBeanDefinition(dynamoDBMappingContextRef, dynamoDBMappingContextBuilder.getBeanDefinition());
        } else {
            LOGGER.debug("DynamoDBMappingContext bean <{}> already exists, reusing it", dynamoDBMappingContextRef);
        }

        return dynamoDBMappingContextRef;
    }

    @Override
    public void registerBeansForRoot(@NonNull BeanDefinitionRegistry registry,
                                     @NonNull RepositoryConfigurationSource configurationSource) {
        super.registerBeansForRoot(registry, configurationSource);

        // Store for later to be used by #postProcess, too
        this.registry = registry;

        // Read marshalling mode from configuration
        if (configurationSource instanceof AnnotationRepositoryConfigurationSource) {
            AnnotationAttributes attributes = ((AnnotationRepositoryConfigurationSource) configurationSource).getAttributes();
            this.marshallingMode = attributes.getEnum("marshallingMode");
            LOGGER.debug("Read marshalling mode from @EnableDynamoDBRepositories: {}", this.marshallingMode);
        } else {
            LOGGER.debug("ConfigurationSource is not AnnotationRepositoryConfigurationSource, using default marshalling mode: {}", this.marshallingMode);
        }

        this.dynamoDBMapperConfigName = getBeanNameWithModulePrefix("DynamoDBMapperConfig");
        Optional<String> dynamoDBMapperConfigRef = configurationSource.getAttribute("dynamoDBMapperConfigRef");

        if (dynamoDBMapperConfigRef.isEmpty()) {
            BeanDefinitionBuilder dynamoDBMapperConfigBuiilder = BeanDefinitionBuilder
                    .genericBeanDefinition(DynamoDBMapperConfigFactory.class);
            registry.registerBeanDefinition(this.dynamoDBMapperConfigName,
                    dynamoDBMapperConfigBuiilder.getBeanDefinition());
        }

        Optional<String> dynamoDBMapperRef = configurationSource.getAttribute("dynamoDBMapperRef");
        if (dynamoDBMapperRef.isEmpty()) {
            this.dynamoDBMapperName = getBeanNameWithModulePrefix("DynamoDBMapper");
            BeanDefinitionBuilder dynamoDBMapperBuilder = BeanDefinitionBuilder
                    .genericBeanDefinition(DynamoDBMapperFactory.class);
            registry.registerBeanDefinition(this.dynamoDBMapperName, dynamoDBMapperBuilder.getBeanDefinition());
        }
    }

    /**
     * Generates a bean name with the DynamoDB module prefix.
     * @param baseBeanName the base bean name
     * @return the prefixed bean name
     */
    @NonNull
    protected String getBeanNameWithModulePrefix(String baseBeanName) {
        return String.format("%s-%s", getModulePrefix(), baseBeanName);
    }

    @NonNull
    @Override
    protected String getModulePrefix() {
        return "dynamoDB";
    }

}