Fixes #363 Allow token() function call will all partition keys

doanduyhai · Jul 5, 2019 · 94a7d8b · 94a7d8b
1 parent 2a76af0
commit 94a7d8b
Show file tree

Hide file tree

Showing 23 changed files with 420 additions and 41 deletions.
diff --git a/achilles-core/buildDependentModules.sh b/achilles-core/buildDependentModules.sh
@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+cd ../achilles-model && \
+mvn clean install && \
+cd ../achilles-common && \
+mvn clean install && \
+cd ../achilles-core
diff --git a/...ava/info/archinnov/achilles/internals/codegen/function/FunctionParameterTypesCodeGen.java b/...ava/info/archinnov/achilles/internals/codegen/function/FunctionParameterTypesCodeGen.java
@@ -18,7 +18,6 @@
 
 import static com.squareup.javapoet.TypeName.BOOLEAN;
 import static info.archinnov.achilles.internals.parser.TypeUtils.*;
-import static java.lang.String.format;
 
 import java.util.List;
 import java.util.Set;
@@ -32,6 +31,8 @@
 
 public abstract class FunctionParameterTypesCodeGen {
 
+    public static String PARTITION_KEYS_TYPE = "PartitionKeys" + FUNCTION_TYPE_SUFFIX;
+
     public abstract List<TypeSpec> buildParameterTypesClasses(FunctionsContext functionContext);
 
     protected abstract void enhanceGeneratedType(TypeSpec.Builder builder, TypeName typeName);
@@ -44,28 +45,76 @@ protected List<TypeSpec> buildParameterTypesClassesInternal(FunctionsContext fun
                 .map(TypeName::box)
                 .collect(Collectors.toSet());
 
-        return uniqueTypeNames
-            .stream()
-            .map(typeName -> {
-                final TypeSpec.Builder builder = TypeSpec.classBuilder(TypeNameHelper.asString(typeName) + FUNCTION_TYPE_SUFFIX)
-                        .superclass(genericType(ABSTRACT_CQL_COMPATIBLE_TYPE, typeName))
-                        .addSuperinterface(FUNCTION_CALL)
+        List<TypeSpec> typeSpecs = uniqueTypeNames
+                .stream()
+                .map(typeName -> {
+                    final TypeSpec.Builder builder = TypeSpec.classBuilder(TypeNameHelper.asString(typeName) + FUNCTION_TYPE_SUFFIX)
+                            .superclass(genericType(ABSTRACT_CQL_COMPATIBLE_TYPE, typeName))
+                            .addSuperinterface(FUNCTION_CALL)
+                            .addModifiers(Modifier.PUBLIC)
+                            .addMethod(buildConstructor(typeName))
+                            .addMethod(buildIsFunctionCall());
+
+                    if (typeName.equals(LIST) || typeName.equals(SET) || typeName.equals(MAP)) {
+                        builder.addAnnotation(AnnotationSpec
+                                .builder(SuppressWarnings.class)
+                                .addMember("value", "$S", "rawtypes")
+                                .build());
+                    }
+
+                    enhanceGeneratedType(builder, typeName);
+                    return builder.build();
+                })
+                .collect(Collectors.toList());
+
+        /**
+         * Build here the special PartitionKey_Type
+         * useful for SystemFunctions.token(xxx_AchillesMeta.COLUMNS.PARTITION_KEYS, "tokens")
+         */
+        TypeSpec partitionKeyType = TypeSpec.classBuilder(PARTITION_KEYS_TYPE)
+                .superclass(genericType(ABSTRACT_CQL_COMPATIBLE_TYPE, STRING))
+                .addSuperinterface(FUNCTION_CALL)
+                .addModifiers(Modifier.PUBLIC)
+                /**
+                 * private final List<String> partitionKeys;
+                 */
+                .addField(FieldSpec
+                        .builder(ParameterizedTypeName.get(LIST, STRING), "partitionKeys", Modifier.PRIVATE, Modifier.FINAL)
+                        .build()
+                )
+                /**
+                 * public PartitionKeys_Type(List<String> partitionKeys) {
+                 *      this.partitionKeys = partitionKeys;
+                 * }
+                 */
+                .addMethod(MethodSpec.constructorBuilder()
                         .addModifiers(Modifier.PUBLIC)
-                        .addMethod(buildConstructor(typeName))
-                        .addMethod(buildIsFunctionCall());
-
-                if (typeName.equals(LIST) || typeName.equals(SET) || typeName.equals(MAP)) {
-                    builder.addAnnotation(AnnotationSpec
-                            .builder(SuppressWarnings.class)
-                            .addMember("value", "$S", "rawtypes")
-                            .build());
-                }
+                        .addParameter(ParameterizedTypeName.get(LIST, STRING), "partitionKeys", Modifier.FINAL)
+                        .addStatement("this.partitionKeys = partitionKeys")
+                        .build())
+                .addMethod(MethodSpec.methodBuilder("isFunctionCall")
+                        .addModifiers(Modifier.PUBLIC)
+                        .addAnnotation(Override.class)
+                        .returns(BOOLEAN)
+                        .addStatement("return true")
+                        .build())
+                /**
+                 * @Override
+                 *   public List<String> getValue() {
+                 *     return this.partitionKeys;
+                 *   }
+                 */
+                .addMethod(MethodSpec.methodBuilder("getValue")
+                        .addModifiers(Modifier.PUBLIC)
+                        .addAnnotation(Override.class)
+                        .returns(ParameterizedTypeName.get(LIST, STRING))
+                        .addStatement("return this.partitionKeys")
+                        .build())
+                .build();
 
-                enhanceGeneratedType(builder, typeName);
-                return builder.build();
-            })
-            .collect(Collectors.toList());
+        typeSpecs.add(partitionKeyType);
 
+        return typeSpecs;
     }
 
     protected MethodSpec buildConstructor(TypeName typeName) {

diff --git a/...ain/java/info/archinnov/achilles/internals/codegen/function/FunctionsRegistryCodeGen.java b/...ain/java/info/archinnov/achilles/internals/codegen/function/FunctionsRegistryCodeGen.java
@@ -18,6 +18,7 @@
 
 import static com.squareup.javapoet.TypeName.BOOLEAN;
 import static com.squareup.javapoet.TypeName.OBJECT;
+import static info.archinnov.achilles.internals.codegen.function.FunctionParameterTypesCodeGen.PARTITION_KEYS_TYPE;
 import static info.archinnov.achilles.internals.parser.TypeUtils.*;
 
 import java.util.ArrayList;
@@ -116,19 +117,8 @@ protected List<MethodSpec> buildAcceptAllMethodsForSystemFunction() {
                 .addMember("value", "$S", "rawtypes")
                 .build();
 
-        //Token function
-        final MethodSpec.Builder tokenFunctionBuilder = MethodSpec.methodBuilder("token")
-                .addTypeVariable(typeVariableName)
-                .addAnnotation(unchecked)
-                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
-                .addJavadoc("Call $S function with given parameters", "token")
-                .returns(LONG_TYPE)
-                .addParameter(typeVariableName, "input", Modifier.FINAL)
-                .addStatement("final $T<Object> params = new $T<>()", LIST, ARRAY_LIST)
-                .addStatement("$T.validateFalse(input.isFunctionCall(), $S)", VALIDATOR, "Invalid argument for 'token' function, it does not accept function call as argument, only simple column")
-                .addStatement("$T.validateFalse(input.hasLiteralValue(), $S)", VALIDATOR, "Invalid argument for 'token' function, it does not accept literal value as argument, only simple column")
-                .addStatement("params.add($T.column((String)$L.getValue()))", QUERY_BUILDER, "input");
-
+        //TODO To remove when upgrading to major version
+        //Legacy Token function
         final TypeSpec.Builder tokenAnonClassBuilder = TypeSpec.anonymousClassBuilder("$T.empty()", OPTIONAL)
                 .superclass(LONG_TYPE)
                 .addMethod(MethodSpec
@@ -153,7 +143,42 @@ protected List<MethodSpec> buildAcceptAllMethodsForSystemFunction() {
                         .addStatement("return params")
                         .build());
 
-        methods.add(tokenFunctionBuilder.addStatement("return $L", tokenAnonClassBuilder.build()).build());
+        final MethodSpec.Builder tokenFunctionBuilder = MethodSpec.methodBuilder("token")
+                .addTypeVariable(typeVariableName)
+                .addAnnotation(unchecked)
+                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
+                .addJavadoc("Call $S function with given parameters", "token")
+                .returns(LONG_TYPE)
+                .addParameter(typeVariableName, "input", Modifier.FINAL)
+                .addStatement("final $T<Object> params = new $T<>()", LIST, ARRAY_LIST)
+                .addStatement("$T.validateFalse(input.isFunctionCall(), $S)", VALIDATOR, "Invalid argument for 'token' function, it does not accept function call as argument, only simple column")
+                .addStatement("$T.validateFalse(input.hasLiteralValue(), $S)", VALIDATOR, "Invalid argument for 'token' function, it does not accept literal value as argument, only simple column")
+                .addStatement("params.add($T.column((String)$L.getValue()))", QUERY_BUILDER, "input")
+                .addStatement("return $L", tokenAnonClassBuilder.build());
+
+        methods.add(tokenFunctionBuilder.build());
+
+
+        //Type-safe token function
+        final MethodSpec.Builder typeSafeTokenFunctionBuilder = MethodSpec.methodBuilder("token")
+                .addAnnotation(unchecked)
+                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
+                .addJavadoc("Call $S function with given parameters", "token")
+                .returns(LONG_TYPE)
+                .addParameter(ClassName.get(FUNCTION_PACKAGE, PARTITION_KEYS_TYPE), "partitionKeys", Modifier.FINAL)
+                /**
+                 *  final List<Object> params = new ArrayList<>();
+                 *  for (String partitionKey : partitionKeys.getValue()) {
+                 *       params.add(QueryBuilder.column(partitionKey));
+                 *  }
+                 */
+                .addStatement("final $T<Object> params = new $T<>()", LIST, ARRAY_LIST)
+                .beginControlFlow("for ($T partitionKey: partitionKeys.getValue())", STRING)
+                    .addStatement("params.add($T.column(partitionKey))", QUERY_BUILDER)
+                .endControlFlow()
+                .addStatement("return $L", tokenAnonClassBuilder.build());
+
+        methods.add(typeSafeTokenFunctionBuilder.build());
 
 
         //writetime function

diff --git a/.../info/archinnov/achilles/internals/codegen/meta/EntityMetaColumnsForFunctionsCodeGen.java b/.../info/archinnov/achilles/internals/codegen/meta/EntityMetaColumnsForFunctionsCodeGen.java
@@ -16,20 +16,22 @@
 
 package info.archinnov.achilles.internals.codegen.meta;
 
+import static info.archinnov.achilles.internals.codegen.function.FunctionParameterTypesCodeGen.PARTITION_KEYS_TYPE;
+import static info.archinnov.achilles.internals.codegen.meta.EntityMetaCodeGen.PARTITION_KEY_SORTER;
 import static info.archinnov.achilles.internals.parser.TypeUtils.*;
 
 import java.util.List;
+import java.util.stream.Collectors;
 import javax.lang.model.element.Modifier;
 
-import com.squareup.javapoet.CodeBlock;
-import com.squareup.javapoet.FieldSpec;
-import com.squareup.javapoet.TypeName;
-import com.squareup.javapoet.TypeSpec;
+import com.squareup.javapoet.*;
 
 import info.archinnov.achilles.internals.metamodel.columns.ColumnType;
+import info.archinnov.achilles.internals.metamodel.columns.PartitionKeyInfo;
 import info.archinnov.achilles.internals.parser.FieldParser.FieldMetaSignature;
 import info.archinnov.achilles.internals.parser.TypeUtils;
 import info.archinnov.achilles.internals.strategy.naming.SnakeCaseNaming;
+import info.archinnov.achilles.type.tuples.Tuple2;
 
 public class EntityMetaColumnsForFunctionsCodeGen  {
 
@@ -45,6 +47,8 @@ public static final TypeSpec createColumnsClassForFunctionParam(List<FieldMetaSi
                 .filter(x -> x.context.columnType != ColumnType.COMPUTED)
                 .forEach(parsingResult -> builder.addField(buildField(parsingResult)));
 
+        builder.addField(buildPartitionKeysField(parsingResults));
+
         return builder.build();
     }
 
@@ -74,5 +78,33 @@ private static final FieldSpec buildField(FieldMetaSignature fieldMetaSignature)
                 .build();
     }
 
+    private static final FieldSpec buildPartitionKeysField(List<FieldMetaSignature> fieldMetaSignatures) {
+        List<String> cqlPartitionKeys = fieldMetaSignatures
+                .stream()
+                .filter(x -> x.context.columnType == ColumnType.PARTITION)
+                .map(x -> Tuple2.of(x.context.quotedCqlColumn, (PartitionKeyInfo) x.context.columnInfo))
+                .sorted(PARTITION_KEY_SORTER)
+                .map(x -> x._1())
+                .collect(Collectors.toList());
+
+        ClassName partitionKeysType = ClassName.get(FUNCTION_PACKAGE, PARTITION_KEYS_TYPE);
+        CodeBlock.Builder builder = CodeBlock.builder()
+                .add("new $T(new $T<$T>() {\n", partitionKeysType, ARRAY_LIST, STRING)
+                .add(" {\n");
+        for (String cqlPartitionKey : cqlPartitionKeys) {
+            builder.add(" add($S);", cqlPartitionKey);
+        }
+        CodeBlock initializer = builder.add(" }\n")
+                .add("  })\n")
+                .build();
+
+        return FieldSpec.builder(partitionKeysType, "PARTITION_KEYS", Modifier.PUBLIC, Modifier.FINAL)
+                .addJavadoc("<br/>\n")
+                .addJavadoc("Field to be used with <em>SystemFunctions.token(xxx_AchillesMeta.COLUMNS.PARTITION_KEYS, \"tokens\")</em> call\n")
+                .addJavadoc("<br/>\n")
+                .initializer(initializer)
+                .build();
+    }
+
 
 }
diff --git a/...esources/expected_code/entity_meta_builder/should_build_entity_with_clustering_column.txt b/...esources/expected_code/entity_meta_builder/should_build_entity_with_clustering_column.txt
@@ -9,6 +9,7 @@ import com.google.common.collect.HashBiMap;
 import com.google.common.reflect.TypeToken;
 import info.archinnov.achilles.generated.function.Date_Type;
 import info.archinnov.achilles.generated.function.Long_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.function.String_Type;
 import info.archinnov.achilles.generated.function.UUID_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithClusteringColumns_AchillesMeta.ColumnsForFunctions;
@@ -33,6 +34,7 @@ import java.lang.Override;
 import java.lang.String;
 import java.lang.SuppressWarnings;
 import java.lang.UnsupportedOperationException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -280,5 +282,16 @@ public final class TestEntityWithClusteringColumns_AchillesMeta extends Abstract
     }
       }
     ;
+
+    /**
+     * <br/>
+     * Field to be used with <em>SystemFunctions.token(xxx_AchillesMeta.COLUMNS.PARTITION_KEYS, "tokens")</em> call
+     * <br/>
+     */
+    public final PartitionKeys_Type PARTITION_KEYS = new PartitionKeys_Type(new ArrayList<String>() {
+     {
+     add("id"); }
+      })
+    ;
   }
 }
diff --git a/...rces/expected_code/entity_meta_builder/should_build_entity_with_complex_counter_types.txt b/...rces/expected_code/entity_meta_builder/should_build_entity_with_complex_counter_types.txt
@@ -8,6 +8,7 @@ import com.google.common.collect.BiMap;
 import com.google.common.collect.HashBiMap;
 import com.google.common.reflect.TypeToken;
 import info.archinnov.achilles.generated.function.Long_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.function.String_Type;
 import info.archinnov.achilles.generated.function.UUID_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithComplexCounters_AchillesMeta.ColumnsForFunctions;
@@ -33,6 +34,7 @@ import java.lang.Override;
 import java.lang.String;
 import java.lang.SuppressWarnings;
 import java.lang.UnsupportedOperationException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Optional;
@@ -308,5 +310,16 @@ public final class TestEntityWithComplexCounters_AchillesMeta extends AbstractEn
     }
       }
     ;
+
+    /**
+     * <br/>
+     * Field to be used with <em>SystemFunctions.token(xxx_AchillesMeta.COLUMNS.PARTITION_KEYS, "tokens")</em> call
+     * <br/>
+     */
+    public final PartitionKeys_Type PARTITION_KEYS = new PartitionKeys_Type(new ArrayList<String>() {
+     {
+     add("id"); }
+      })
+    ;
   }
 }
diff --git a/.../resources/expected_code/entity_meta_builder/should_build_entity_with_complex_indices.txt b/.../resources/expected_code/entity_meta_builder/should_build_entity_with_complex_indices.txt
@@ -10,6 +10,7 @@ import info.archinnov.achilles.generated.function.List_String_Type;
 import info.archinnov.achilles.generated.function.Long_Type;
 import info.archinnov.achilles.generated.function.Map_Integer_String_Type;
 import info.archinnov.achilles.generated.function.Map_String_String_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.function.Set_String_Type;
 import info.archinnov.achilles.generated.function.String_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithComplexIndices_AchillesMeta.ColumnsForFunctions;
@@ -37,6 +38,7 @@ import java.lang.Override;
 import java.lang.String;
 import java.lang.SuppressWarnings;
 import java.lang.UnsupportedOperationException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
@@ -342,5 +344,16 @@ public final class TestEntityWithComplexIndices_AchillesMeta extends AbstractEnt
     }
       }
     ;
+
+    /**
+     * <br/>
+     * Field to be used with <em>SystemFunctions.token(xxx_AchillesMeta.COLUMNS.PARTITION_KEYS, "tokens")</em> call
+     * <br/>
+     */
+    public final PartitionKeys_Type PARTITION_KEYS = new PartitionKeys_Type(new ArrayList<String>() {
+     {
+     add("id"); }
+      })
+    ;
   }
 }