DynamoDBEntityMetadataSupport.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.support;

import org.socialsignin.spring.data.dynamodb.core.DynamoDBOperations;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
import org.springframework.util.ReflectionUtils;
import org.springframework.util.StringUtils;
import software.amazon.awssdk.enhanced.dynamodb.AttributeConverter;
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.*;

import java.lang.annotation.Annotation;
import java.lang.reflect.Field;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.*;

/**
 * Base support class for DynamoDB entity metadata that handles extraction and caching of hash key information,
 * attribute name overrides, attribute converters, and global secondary index metadata from DynamoDB-annotated entities.
 * @param <T> the entity type
 * @param <ID> the ID type
 * @author Prasanna Kumar Ramachandran
 */
public class DynamoDBEntityMetadataSupport<T, ID> implements DynamoDBHashKeyExtractingEntityMetadata<T> {

    @NonNull
    private final Class<T> domainType;
    private boolean hasRangeKey;
    @Nullable
    private String hashKeyPropertyName;
    @NonNull
    private final List<String> globalIndexHashKeyPropertyNames;
    @NonNull
    private final List<String> globalIndexRangeKeyPropertyNames;

    private final String dynamoDBTableName;
    @NonNull
    private final Map<String, String[]> globalSecondaryIndexNames;

    @Override
    public String getDynamoDBTableName() {
        return dynamoDBTableName;
    }

    /**
     * Creates a new {@link DynamoDBEntityMetadataSupport} for the given domain type.
     *
     * @param domainType must not be {@literal null}.
     */
    public DynamoDBEntityMetadataSupport(@NonNull final Class<T> domainType) {
        this(domainType, null);
    }

    /**
     * Creates a new {@link DynamoDBEntityMetadataSupport} for the given domain type and dynamoDB mapper config.
     *
     * @param domainType must not be {@literal null}.
     * @param dynamoDBOperations dynamoDBOperations as populated from Spring Data DynamoDB Configuration
     */
    public DynamoDBEntityMetadataSupport(@NonNull final Class<T> domainType, @Nullable DynamoDBOperations dynamoDBOperations) {

        Assert.notNull(domainType, "Domain type must not be null!");
        this.domainType = domainType;

        DynamoDbBean table = this.domainType.getAnnotation(DynamoDbBean.class);
        DynamoDbImmutable immutableTable = this.domainType.getAnnotation(DynamoDbImmutable.class);
        Assert.isTrue(table != null || immutableTable != null, "Domain type must be annotated with @DynamoDbBean or @DynamoDbImmutable!");

        // In SDK v2, table name is typically inferred from class name or set via TableSchema
        // For now, use the class simple name as default
        String tableName = domainType.getSimpleName();

        if (dynamoDBOperations != null) {
            this.dynamoDBTableName = dynamoDBOperations.getOverriddenTableName(domainType, tableName);
        } else {
            this.dynamoDBTableName = tableName;
        }
        this.hashKeyPropertyName = null;
        this.globalSecondaryIndexNames = new HashMap<>();
        this.globalIndexHashKeyPropertyNames = new ArrayList<>();
        this.globalIndexRangeKeyPropertyNames = new ArrayList<>();
        ReflectionUtils.doWithMethods(domainType, method -> {
            if (method.getAnnotation(DynamoDbPartitionKey.class) != null) {
                hashKeyPropertyName = getPropertyNameForAccessorMethod(method);
            }
            if (method.getAnnotation(DynamoDbSortKey.class) != null) {
                hasRangeKey = true;
            }
            DynamoDbSecondarySortKey dynamoDBRangeKeyAnnotation = method.getAnnotation(DynamoDbSecondarySortKey.class);
            DynamoDbSecondaryPartitionKey dynamoDBHashKeyAnnotation = method.getAnnotation(DynamoDbSecondaryPartitionKey.class);

            if (dynamoDBRangeKeyAnnotation != null) {
                addGlobalSecondaryIndexNames(method, dynamoDBRangeKeyAnnotation);
            }
            if (dynamoDBHashKeyAnnotation != null) {
                addGlobalSecondaryIndexNames(method, dynamoDBHashKeyAnnotation);
            }
        });
        ReflectionUtils.doWithFields(domainType, field -> {
            if (field.getAnnotation(DynamoDbPartitionKey.class) != null) {
                hashKeyPropertyName = getPropertyNameForField(field);
            }
            if (field.getAnnotation(DynamoDbSortKey.class) != null) {
                hasRangeKey = true;
            }
            DynamoDbSecondarySortKey dynamoDBRangeKeyAnnotation = field.getAnnotation(DynamoDbSecondarySortKey.class);
            DynamoDbSecondaryPartitionKey dynamoDBHashKeyAnnotation = field.getAnnotation(DynamoDbSecondaryPartitionKey.class);

            if (dynamoDBRangeKeyAnnotation != null) {
                addGlobalSecondaryIndexNames(field, dynamoDBRangeKeyAnnotation);
            }
            if (dynamoDBHashKeyAnnotation != null) {
                addGlobalSecondaryIndexNames(field, dynamoDBHashKeyAnnotation);
            }
        });
        Assert.notNull(hashKeyPropertyName, "Unable to find hash key field or getter method on " + domainType + "!");
    }

    /**
     * Creates and returns the appropriate DynamoDBEntityInformation implementation based on whether
     * the entity has a range key.
     *
     * @return entity information for this entity type
     */
    @NonNull
    public DynamoDBEntityInformation<T, ID> getEntityInformation() {

        if (hasRangeKey) {
            DynamoDBHashAndRangeKeyExtractingEntityMetadataImpl<T, ID> metadata = new DynamoDBHashAndRangeKeyExtractingEntityMetadataImpl<>(
                    domainType);
            return new DynamoDBIdIsHashAndRangeKeyEntityInformationImpl<>(domainType, metadata);
        } else {
            return new DynamoDBIdIsHashKeyEntityInformationImpl<>(domainType, this);
        }
    }

    /*
     * (non-Javadoc)
     * @see org.springframework.data.repository.core.EntityMetadata#getJavaType()
     */
    @NonNull
    @Override
    public Class<T> getJavaType() {
        return domainType;
    }

    @Override
    public boolean isHashKeyProperty(String propertyName) {
        return hashKeyPropertyName != null && hashKeyPropertyName.equals(propertyName);
    }

    /**
     * Checks if the field for the given property is annotated with @Id.
     *
     * @param propertyName the property name
     * @return true if the field is annotated with @Id, false otherwise
     */
    protected boolean isFieldAnnotatedWith(@NonNull final String propertyName) {

        Field field = findField(propertyName);
        return field != null && field.getAnnotation((Class<? extends Annotation>) org.springframework.data.annotation.Id.class) != null;
    }

    @NonNull
    private String toGetMethodName(@NonNull String propertyName) {
        String methodName = propertyName.substring(0, 1).toUpperCase();
        if (propertyName.length() > 1) {
            methodName = methodName + propertyName.substring(1);
        }
        return "get" + methodName;
    }

    /**
     * Converts an accessor method name to its corresponding setter method name.
     *
     * @param method the accessor method
     * @return the setter method name, or null if conversion is not possible
     */
    @Nullable
    protected String toSetterMethodNameFromAccessorMethod(@NonNull Method method) {
        String accessorMethodName = method.getName();
        if (accessorMethodName.startsWith("get")) {
            return "set" + accessorMethodName.substring(3);
        } else if (accessorMethodName.startsWith("is")) {
            return "is" + accessorMethodName.substring(2);
        }
        return null;
    }

    @NonNull
    private String toIsMethodName(@NonNull String propertyName) {
        String methodName = propertyName.substring(0, 1).toUpperCase();
        if (propertyName.length() > 1) {
            methodName = methodName + propertyName.substring(1);
        }
        return "is" + methodName;
    }

    @Nullable
    private Method findMethod(@NonNull String propertyName) {
        Method method = ReflectionUtils.findMethod(domainType, toGetMethodName(propertyName));
        if (method == null) {
            method = ReflectionUtils.findMethod(domainType, toIsMethodName(propertyName));
        }
        return method;

    }

    @Nullable
    private Field findField(@NonNull String propertyName) {
        return ReflectionUtils.findField(domainType, propertyName);
    }

    @NonNull
    @Override
    public Optional<String> getOverriddenAttributeName(@NonNull final String propertyName) {

        Method method = findMethod(propertyName);
        if (method != null) {
            // In SDK v2, @DynamoDbAttribute is used to override attribute names
            if (method.getAnnotation(DynamoDbAttribute.class) != null
                    && StringUtils.hasText(method.getAnnotation(DynamoDbAttribute.class).value())) {
                return Optional.of(method.getAnnotation(DynamoDbAttribute.class).value());
            }
            // Note: SDK v2 key annotations don't support attribute name overrides like SDK v1
        }

        Field field = findField(propertyName);
        if (field != null) {
            // In SDK v2, @DynamoDbAttribute is used to override attribute names
            if (field.getAnnotation(DynamoDbAttribute.class) != null
                    && StringUtils.hasText(field.getAnnotation(DynamoDbAttribute.class).value())) {
                return Optional.of(field.getAnnotation(DynamoDbAttribute.class).value());
            }
            // Note: SDK v2 key annotations don't support attribute name overrides like SDK v1
        }
        return Optional.empty();

    }

    @Nullable
    @Override
    public AttributeConverter<?> getAttributeConverterForProperty(@NonNull final String propertyName) {
        // SDK v2 uses @DynamoDbConvertedBy annotation for custom converters
        DynamoDbConvertedBy annotation = null;

        Method method = findMethod(propertyName);
        if (method != null) {
            annotation = method.getAnnotation(DynamoDbConvertedBy.class);
        }

        if (annotation == null) {
            Field field = findField(propertyName);
            if (field != null) {
                annotation = field.getAnnotation(DynamoDbConvertedBy.class);
            }
        }

        if (annotation != null) {
            try {
                return annotation.value().getDeclaredConstructor().newInstance();
            } catch (InstantiationException | IllegalAccessException | IllegalArgumentException
                    | InvocationTargetException | NoSuchMethodException | SecurityException e) {
                throw new RuntimeException(e);
            }
        }

        // No custom converter annotation found, check if AWS SDK v2 has a default converter for this type
        // This allows enums and other types to work natively without requiring custom converters
        Class<?> propertyType = null;
        if (method != null) {
            propertyType = method.getReturnType();
        } else {
            Field field = findField(propertyName);
            if (field != null) {
                propertyType = field.getType();
            }
        }

        if (propertyType != null) {
            // AWS SDK v2's DefaultAttributeConverterProvider cannot create converters for raw collection types
            // (Set, List, Map) without generic type parameters. When we get the type via field.getType() or
            // method.getReturnType(), we lose the generic type information (e.g., Set<String> becomes Set).
            // DynamoDB natively supports these collection types, so no custom converter is needed.
            // Return null to indicate no converter is available - the value will be handled by DynamoDB's
            // native type support.
            if (java.util.Collection.class.isAssignableFrom(propertyType) ||
                java.util.Map.class.isAssignableFrom(propertyType)) {
                return null;
            }

            software.amazon.awssdk.enhanced.dynamodb.DefaultAttributeConverterProvider defaultProvider =
                software.amazon.awssdk.enhanced.dynamodb.DefaultAttributeConverterProvider.create();
            return defaultProvider.converterFor(software.amazon.awssdk.enhanced.dynamodb.EnhancedType.of(propertyType));
        }

        return null;
    }

    /**
     * Extracts the property name from a getter/accessor method name.
     *
     * @param method the accessor method
     * @return the property name
     */
    @NonNull
    protected String getPropertyNameForAccessorMethod(@NonNull Method method) {
        String methodName = method.getName();
        String propertyName = null;
        if (methodName.startsWith("get")) {
            propertyName = methodName.substring(3);
        } else if (methodName.startsWith("is")) {
            propertyName = methodName.substring(2);
        }
        Assert.notNull(propertyName, "Hash or range key annotated accessor methods must start with 'get' or 'is'");

        String firstLetter = propertyName.substring(0, 1);
        String remainder = propertyName.substring(1);
        return firstLetter.toLowerCase() + remainder;
    }

    /**
     * Gets the property name from a field.
     *
     * @param field the field
     * @return the property name
     */
    @NonNull
    protected String getPropertyNameForField(@NonNull Field field) {
        return field.getName();
    }

    @Nullable
    @Override
    public String getHashKeyPropertyName() {
        return hashKeyPropertyName;
    }

    private void addGlobalSecondaryIndexNames(@NonNull Method method, @NonNull DynamoDbSecondarySortKey dynamoDBSecondarySortKey) {

        // SDK v2 uses indexNames() which returns array of index names for both GSI and LSI
        if (dynamoDBSecondarySortKey.indexNames() != null
                && dynamoDBSecondarySortKey.indexNames().length > 0) {
            String propertyName = getPropertyNameForAccessorMethod(method);

            globalSecondaryIndexNames.put(propertyName, dynamoDBSecondarySortKey.indexNames());
            globalIndexRangeKeyPropertyNames.add(propertyName);
        }

    }

    private void addGlobalSecondaryIndexNames(@NonNull Field field, @NonNull DynamoDbSecondarySortKey dynamoDBSecondarySortKey) {

        // SDK v2 uses indexNames() which returns array of index names for both GSI and LSI
        if (dynamoDBSecondarySortKey.indexNames() != null
                && dynamoDBSecondarySortKey.indexNames().length > 0) {
            String propertyName = getPropertyNameForField(field);

            globalSecondaryIndexNames.put(propertyName, dynamoDBSecondarySortKey.indexNames());
            globalIndexRangeKeyPropertyNames.add(propertyName);
        }

    }

    private void addGlobalSecondaryIndexNames(@NonNull Method method, @NonNull DynamoDbSecondaryPartitionKey dynamoDBSecondaryPartitionKey) {

        // SDK v2 uses indexNames() which returns array of index names for GSI
        if (dynamoDBSecondaryPartitionKey.indexNames() != null
                && dynamoDBSecondaryPartitionKey.indexNames().length > 0) {
            String propertyName = getPropertyNameForAccessorMethod(method);

            globalSecondaryIndexNames.put(propertyName, dynamoDBSecondaryPartitionKey.indexNames());
            globalIndexHashKeyPropertyNames.add(propertyName);
        }
    }

    private void addGlobalSecondaryIndexNames(@NonNull Field field, @NonNull DynamoDbSecondaryPartitionKey dynamoDBSecondaryPartitionKey) {

        // SDK v2 uses indexNames() which returns array of index names for GSI
        if (dynamoDBSecondaryPartitionKey.indexNames() != null
                && dynamoDBSecondaryPartitionKey.indexNames().length > 0) {
            String propertyName = getPropertyNameForField(field);

            globalSecondaryIndexNames.put(propertyName, dynamoDBSecondaryPartitionKey.indexNames());
            globalIndexHashKeyPropertyNames.add(propertyName);
        }
    }

    @NonNull
    @Override
    public Map<String, String[]> getGlobalSecondaryIndexNamesByPropertyName() {
        return globalSecondaryIndexNames;
    }

    @Override
    public boolean isGlobalIndexHashKeyProperty(String propertyName) {
        return globalIndexHashKeyPropertyNames.contains(propertyName);
    }

    @Override
    public boolean isGlobalIndexRangeKeyProperty(String propertyName) {
        return globalIndexRangeKeyPropertyNames.contains(propertyName);
    }

}