StaticTableSchemaGenerator.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.core;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.util.ReflectionUtils;
import software.amazon.awssdk.enhanced.dynamodb.AttributeConverter;
import software.amazon.awssdk.enhanced.dynamodb.EnhancedType;
import software.amazon.awssdk.enhanced.dynamodb.TableSchema;
import software.amazon.awssdk.enhanced.dynamodb.extensions.VersionedRecordExtension;
import software.amazon.awssdk.enhanced.dynamodb.extensions.annotations.DynamoDbVersionAttribute;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTags;
import software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema;
import software.amazon.awssdk.enhanced.dynamodb.mapper.annotations.*;

import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import java.lang.invoke.MethodType;
import java.lang.reflect.Method;
import java.lang.reflect.Modifier;
import java.lang.reflect.ParameterizedType;
import java.lang.reflect.Type;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Function;

/**
 * Generates {@link StaticTableSchema} instances from {@code @DynamoDbBean} annotated classes
 * using reflection and MethodHandles instead of LambdaMetafactory.
 *
 * <p>This generator is compatible with GraalVM native images because it uses
 * {@link MethodHandle#invoke} which works with ahead-of-time compilation,
 * unlike {@code TableSchema.fromBean()} which uses {@code LambdaMetafactory}.
 *
 * <p><b>Supported Annotations:</b></p>
 * <ul>
 *   <li>{@code @DynamoDbBean} - Required class-level annotation</li>
 *   <li>{@code @DynamoDbPartitionKey} - Marks the partition key attribute</li>
 *   <li>{@code @DynamoDbSortKey} - Marks the sort key attribute</li>
 *   <li>{@code @DynamoDbAttribute} - Customizes attribute name</li>
 *   <li>{@code @DynamoDbConvertedBy} - Specifies custom converter</li>
 *   <li>{@code @DynamoDbSecondaryPartitionKey} - GSI partition key</li>
 *   <li>{@code @DynamoDbSecondarySortKey} - GSI/LSI sort key</li>
 *   <li>{@code @DynamoDbVersionAttribute} - Optimistic locking version attribute</li>
 *   <li>{@code @DynamoDbIgnore} - Excludes property from mapping</li>
 * </ul>
 *
 * <p><b>Performance Note:</b></p>
 * <p>MethodHandle-based getters/setters are slightly slower than lambda-based ones
 * (generated by {@code TableSchema.fromBean()}), but the difference is typically
 * negligible for most applications. The tradeoff is GraalVM native image compatibility.
 *
 * @author Prasanna Kumar Ramachandran
 * @since 7.0.0
 * @see DynamoDbTableSchemaRegistry
 * @see software.amazon.awssdk.enhanced.dynamodb.mapper.StaticTableSchema
 */
public class StaticTableSchemaGenerator {

    private static final Logger LOGGER = LoggerFactory.getLogger(StaticTableSchemaGenerator.class);
    private static final MethodHandles.Lookup LOOKUP = MethodHandles.lookup();

    /**
     * Private constructor to prevent instantiation.
     */
    private StaticTableSchemaGenerator() {
    }

    /**
     * Generates a StaticTableSchema for the given domain class.
     *
     * @param <T>         the entity type
     * @param domainClass the domain class annotated with {@code @DynamoDbBean}
     * @return the generated TableSchema
     * @throws IllegalArgumentException if the class is not annotated with {@code @DynamoDbBean}
     * @throws RuntimeException         if schema generation fails
     */
    @NonNull
    public static <T> TableSchema<T> generateSchema(@NonNull Class<T> domainClass) {
        LOGGER.debug("Generating StaticTableSchema for class: {}", domainClass.getName());

        // Validate @DynamoDbBean annotation
        if (!domainClass.isAnnotationPresent(DynamoDbBean.class) &&
                !domainClass.isAnnotationPresent(DynamoDbImmutable.class)) {
            throw new IllegalArgumentException(
                    "Class " + domainClass.getName() + " must be annotated with @DynamoDbBean or @DynamoDbImmutable");
        }

        try {
            StaticTableSchema.Builder<T> builder = StaticTableSchema.builder(domainClass)
                    .newItemSupplier(createNewItemSupplier(domainClass));

            // Process all getter methods to find attributes
            List<AttributeInfo> attributes = discoverAttributes(domainClass);

            for (AttributeInfo attr : attributes) {
                addAttribute(builder, domainClass, attr);
            }

            TableSchema<T> schema = builder.build();
            LOGGER.debug("Successfully generated StaticTableSchema for class: {}", domainClass.getName());
            return schema;

        } catch (Exception e) {
            throw new RuntimeException("Failed to generate StaticTableSchema for " + domainClass.getName(), e);
        }
    }

    /**
     * Creates a supplier for new instances of the domain class.
     */
    @NonNull
    private static <T> java.util.function.Supplier<T> createNewItemSupplier(@NonNull Class<T> domainClass) {
        try {
            MethodHandle constructor = LOOKUP.findConstructor(domainClass, MethodType.methodType(void.class));
            return () -> {
                try {
                    @SuppressWarnings("unchecked")
                    T instance = (T) constructor.invoke();
                    return instance;
                } catch (Throwable t) {
                    throw new RuntimeException("Failed to create new instance of " + domainClass.getName(), t);
                }
            };
        } catch (NoSuchMethodException | IllegalAccessException e) {
            throw new RuntimeException("No accessible no-arg constructor found for " + domainClass.getName(), e);
        }
    }

    /**
     * Discovers all attributes from getter methods and fields.
     */
    @NonNull
    private static List<AttributeInfo> discoverAttributes(@NonNull Class<?> domainClass) {
        List<AttributeInfo> attributes = new ArrayList<>();
        Set<String> processedProperties = new HashSet<>();

        // Process getter methods
        for (Method method : domainClass.getMethods()) {
            if (isGetter(method) && !method.isAnnotationPresent(DynamoDbIgnore.class)) {
                String propertyName = getPropertyNameFromGetter(method);
                if (processedProperties.add(propertyName)) {
                    Method setter = findSetter(domainClass, propertyName, method.getReturnType());
                    if (setter != null) {
                        attributes.add(new AttributeInfo(propertyName, method, setter));
                    }
                }
            }
        }

        return attributes;
    }

    /**
     * Checks if a method is a getter.
     */
    private static boolean isGetter(@NonNull Method method) {
        if (Modifier.isStatic(method.getModifiers())) {
            return false;
        }
        if (method.getParameterCount() != 0) {
            return false;
        }
        if (method.getReturnType() == void.class) {
            return false;
        }
        String name = method.getName();
        if (name.startsWith("get") && name.length() > 3) {
            return !name.equals("getClass");
        }
        if (name.startsWith("is") && name.length() > 2) {
            return method.getReturnType() == boolean.class || method.getReturnType() == Boolean.class;
        }
        return false;
    }

    /**
     * Extracts property name from getter method name.
     */
    @NonNull
    private static String getPropertyNameFromGetter(@NonNull Method method) {
        String name = method.getName();
        String propertyName;
        if (name.startsWith("get")) {
            propertyName = name.substring(3);
        } else if (name.startsWith("is")) {
            propertyName = name.substring(2);
        } else {
            throw new IllegalArgumentException("Not a getter method: " + name);
        }
        return Character.toLowerCase(propertyName.charAt(0)) + propertyName.substring(1);
    }

    /**
     * Finds the setter method for a property.
     */
    @Nullable
    private static Method findSetter(@NonNull Class<?> domainClass, @NonNull String propertyName, @NonNull Class<?> propertyType) {
        String setterName = "set" + Character.toUpperCase(propertyName.charAt(0)) + propertyName.substring(1);
        return ReflectionUtils.findMethod(domainClass, setterName, propertyType);
    }

    /**
     * Adds an attribute to the schema builder.
     */
    @SuppressWarnings({"unchecked", "rawtypes"})
    private static <T> void addAttribute(
            @NonNull StaticTableSchema.Builder<T> builder,
            @NonNull Class<T> domainClass,
            @NonNull AttributeInfo attr) {

        Method getter = attr.getter();
        Method setter = attr.setter();
        String propertyName = attr.propertyName();
        Class<?> attributeType = getter.getReturnType();

        // Determine attribute name (may be overridden by @DynamoDbAttribute)
        String attributeName = propertyName;
        DynamoDbAttribute attrAnnotation = getter.getAnnotation(DynamoDbAttribute.class);
        if (attrAnnotation != null && !attrAnnotation.value().isEmpty()) {
            attributeName = attrAnnotation.value();
        }

        // Create getter function using MethodHandle
        Function<T, ?> getterFunction = createGetterFunction(domainClass, getter);

        // Create setter function using MethodHandle
        BiConsumer<T, ?> setterConsumer = createSetterConsumer(domainClass, setter, attributeType);

        // Determine tags (partition key, sort key, GSI keys, version)
        List<software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag> tags = determineTags(getter);

        // Get custom converter if specified
        AttributeConverter<?> converter = getCustomConverter(getter);

        // Build the attribute
        String finalAttributeName = attributeName;

        // Handle different attribute types
        EnhancedType enhancedType = getEnhancedType(getter);

        builder.addAttribute(enhancedType, a -> {
            a.name(finalAttributeName)
                    .getter((Function) getterFunction)
                    .setter((BiConsumer) setterConsumer);

            if (!tags.isEmpty()) {
                a.tags(tags.toArray(new software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag[0]));
            }

            if (converter != null) {
                a.attributeConverter((AttributeConverter) converter);
            }
        });

        LOGGER.trace("Added attribute '{}' (type: {}) to schema", finalAttributeName, attributeType.getSimpleName());
    }

    /**
     * Gets the EnhancedType for a getter method, handling generic types.
     */
    @NonNull
    @SuppressWarnings("rawtypes")
    private static EnhancedType getEnhancedType(@NonNull Method getter) {
        Class<?> returnType = getter.getReturnType();
        Type genericReturnType = getter.getGenericReturnType();

        // Handle parameterized types (List<String>, Set<Integer>, Map<String, Object>, etc.)
        if (genericReturnType instanceof ParameterizedType parameterizedType) {
            if (List.class.isAssignableFrom(returnType)) {
                Type[] typeArgs = parameterizedType.getActualTypeArguments();
                if (typeArgs.length > 0 && typeArgs[0] instanceof Class) {
                    return EnhancedType.listOf((Class<?>) typeArgs[0]);
                }
            } else if (Set.class.isAssignableFrom(returnType)) {
                Type[] typeArgs = parameterizedType.getActualTypeArguments();
                if (typeArgs.length > 0 && typeArgs[0] instanceof Class) {
                    return EnhancedType.setOf((Class<?>) typeArgs[0]);
                }
            } else if (Map.class.isAssignableFrom(returnType)) {
                Type[] typeArgs = parameterizedType.getActualTypeArguments();
                if (typeArgs.length == 2 && typeArgs[0] instanceof Class && typeArgs[1] instanceof Class) {
                    return EnhancedType.mapOf((Class<?>) typeArgs[0], (Class<?>) typeArgs[1]);
                }
            }
        }

        return EnhancedType.of(returnType);
    }

    /**
     * Creates a getter function using MethodHandle.
     */
    @NonNull
    private static <T> Function<T, Object> createGetterFunction(@NonNull Class<T> domainClass, @NonNull Method getter) {
        try {
            MethodHandle handle = LOOKUP.unreflect(getter);
            return instance -> {
                try {
                    return handle.invoke(instance);
                } catch (Throwable t) {
                    throw new RuntimeException("Failed to invoke getter " + getter.getName(), t);
                }
            };
        } catch (IllegalAccessException e) {
            throw new RuntimeException("Failed to create getter for " + getter.getName(), e);
        }
    }

    /**
     * Creates a setter consumer using MethodHandle.
     */
    @NonNull
    private static <T> BiConsumer<T, Object> createSetterConsumer(
            @NonNull Class<T> domainClass,
            @NonNull Method setter,
            @NonNull Class<?> attributeType) {
        try {
            MethodHandle handle = LOOKUP.unreflect(setter);
            return (instance, value) -> {
                try {
                    handle.invoke(instance, value);
                } catch (Throwable t) {
                    throw new RuntimeException("Failed to invoke setter " + setter.getName(), t);
                }
            };
        } catch (IllegalAccessException e) {
            throw new RuntimeException("Failed to create setter for " + setter.getName(), e);
        }
    }

    /**
     * Determines the attribute tags based on annotations.
     */
    @NonNull
    private static List<software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag> determineTags(@NonNull Method getter) {
        List<software.amazon.awssdk.enhanced.dynamodb.mapper.StaticAttributeTag> tags = new ArrayList<>();

        if (getter.isAnnotationPresent(DynamoDbPartitionKey.class)) {
            tags.add(StaticAttributeTags.primaryPartitionKey());
        }

        if (getter.isAnnotationPresent(DynamoDbSortKey.class)) {
            tags.add(StaticAttributeTags.primarySortKey());
        }

        DynamoDbSecondaryPartitionKey secondaryPartitionKey = getter.getAnnotation(DynamoDbSecondaryPartitionKey.class);
        if (secondaryPartitionKey != null) {
            for (String indexName : secondaryPartitionKey.indexNames()) {
                tags.add(StaticAttributeTags.secondaryPartitionKey(indexName));
            }
        }

        DynamoDbSecondarySortKey secondarySortKey = getter.getAnnotation(DynamoDbSecondarySortKey.class);
        if (secondarySortKey != null) {
            for (String indexName : secondarySortKey.indexNames()) {
                tags.add(StaticAttributeTags.secondarySortKey(indexName));
            }
        }

        // Handle @DynamoDbVersionAttribute for optimistic locking
        if (getter.isAnnotationPresent(DynamoDbVersionAttribute.class)) {
            tags.add(VersionedRecordExtension.AttributeTags.versionAttribute());
        }

        return tags;
    }

    /**
     * Gets a custom converter from @DynamoDbConvertedBy annotation.
     */
    @Nullable
    private static AttributeConverter<?> getCustomConverter(@NonNull Method getter) {
        DynamoDbConvertedBy convertedBy = getter.getAnnotation(DynamoDbConvertedBy.class);
        if (convertedBy != null) {
            try {
                return convertedBy.value().getDeclaredConstructor().newInstance();
            } catch (Exception e) {
                throw new RuntimeException("Failed to instantiate converter: " + convertedBy.value().getName(), e);
            }
        }
        return null;
    }

    /**
     * Holds information about an attribute.
     */
    private record AttributeInfo(String propertyName, Method getter, Method setter) {
    }
}