From 94a7d8bf7bfc8e99bfa8866850cbce98d55ab464 Mon Sep 17 00:00:00 2001
From: DuyHai DOAN <doanduyhai@gmail.com>
Date: Fri, 5 Jul 2019 23:23:37 +0200
Subject: [PATCH] Fixes #363 Allow token() function call will all partition
 keys

---
 achilles-core/buildDependentModules.sh        |  7 ++
 .../FunctionParameterTypesCodeGen.java        | 89 ++++++++++++++-----
 .../function/FunctionsRegistryCodeGen.java    | 53 ++++++++---
 .../EntityMetaColumnsForFunctionsCodeGen.java | 40 ++++++++-
 ...ld_build_entity_with_clustering_column.txt | 13 +++
 ...uild_entity_with_complex_counter_types.txt | 13 +++
 ...ould_build_entity_with_complex_indices.txt | 13 +++
 ...should_build_entity_with_complex_types.txt | 13 +++
 ...ld_entity_with_composite_partition_key.txt | 13 +++
 ...ould_build_entity_with_computed_column.txt | 13 +++
 ...hould_build_entity_with_counter_column.txt | 13 +++
 ...d_build_entity_with_custom_constructor.txt | 13 +++
 ...ustom_constructor_with_declared_fields.txt | 13 +++
 ...ild_entity_with_implicit_field_parsing.txt | 13 +++
 ...build_entity_with_simple_partition_key.txt | 13 +++
 ...d_build_entity_with_static_annotations.txt | 13 +++
 ...should_build_entity_with_static_column.txt | 13 +++
 ...uild_entity_with_static_counter_column.txt | 13 +++
 .../should_build_inherited_entity.txt         | 13 +++
 .../should_build_view_meta.txt                | 13 +++
 .../TestEntityWithCompositePartitionKey.java  | 35 +++++++-
 .../it/TestFunctionCallsSimpleEntity.java     |  4 +
 .../achilles/it/TestEntityForGroupByIT.java   | 25 ++++++
 23 files changed, 420 insertions(+), 41 deletions(-)
 create mode 100755 achilles-core/buildDependentModules.sh

diff --git a/achilles-core/buildDependentModules.sh b/achilles-core/buildDependentModules.sh
new file mode 100755
index 000000000..9650d34ec
--- /dev/null
+++ 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
\ No newline at end of file
diff --git a/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionParameterTypesCodeGen.java b/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionParameterTypesCodeGen.java
index 32ccef8fe..674415c07 100644
--- a/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionParameterTypesCodeGen.java
+++ b/achilles-core/src/main/java/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/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionsRegistryCodeGen.java b/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionsRegistryCodeGen.java
index 2a93539f0..6abff4e5a 100644
--- a/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/function/FunctionsRegistryCodeGen.java
+++ b/achilles-core/src/main/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/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/meta/EntityMetaColumnsForFunctionsCodeGen.java b/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/meta/EntityMetaColumnsForFunctionsCodeGen.java
index 3870bc2d6..bd1095f3c 100644
--- a/achilles-core/src/main/java/info/archinnov/achilles/internals/codegen/meta/EntityMetaColumnsForFunctionsCodeGen.java
+++ b/achilles-core/src/main/java/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/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_clustering_column.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_clustering_column.txt
index 938c8994f..b0749fef6 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_clustering_column.txt
+++ b/achilles-core/src/test/resources/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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_counter_types.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_counter_types.txt
index 74da2c437..dec296c32 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_counter_types.txt
+++ b/achilles-core/src/test/resources/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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_indices.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_indices.txt
index 7d2499a5b..32b901af9 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_indices.txt
+++ b/achilles-core/src/test/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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_types.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_types.txt
index fed57df3d..6717e6772 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_types.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_complex_types.txt
@@ -23,6 +23,7 @@ import info.archinnov.achilles.generated.function.Map_Integer_List_Integer_Type;
 import info.archinnov.achilles.generated.function.Map_Integer_List_Map_Integer_String_Type;
 import info.archinnov.achilles.generated.function.Map_Integer_TestUDT_Type;
 import info.archinnov.achilles.generated.function.Map_TestUDT_Map_Integer_Tuple3_Integer_Integer_ConsistencyLevel_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.function.Set_ConsistencyLevel_Type;
 import info.archinnov.achilles.generated.function.String_Type;
 import info.archinnov.achilles.generated.function.TestUDT_Type;
@@ -75,6 +76,7 @@ import java.lang.String;
 import java.lang.SuppressWarnings;
 import java.lang.UnsupportedOperationException;
 import java.nio.ByteBuffer;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -896,5 +898,16 @@ public final class TestEntityWithComplexTypes_AchillesMeta extends AbstractEntit
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_composite_partition_key.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_composite_partition_key.txt
index 23af0e471..782020525 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_composite_partition_key.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_composite_partition_key.txt
@@ -7,6 +7,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.TestEntityWithCompositePartitionKey_AchillesMeta.ColumnsForFunctions;
@@ -30,6 +31,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;
@@ -247,5 +249,16 @@ public final class TestEntityWithCompositePartitionKey_AchillesMeta extends Abst
     }
       }
     ;
+
+    /**
+     * <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"); add("uuid"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_computed_column.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_computed_column.txt
index cc993f91b..63249cff1 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_computed_column.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_computed_column.txt
@@ -7,6 +7,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.meta.entity.TestEntityWithComputedColumn_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -31,6 +32,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;
@@ -229,5 +231,16 @@ public final class TestEntityWithComputedColumn_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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_counter_column.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_counter_column.txt
index f12c15857..30a2f9a4e 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_counter_column.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_counter_column.txt
@@ -7,6 +7,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.meta.entity.TestEntityWithCounterColumn_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
 import info.archinnov.achilles.internals.codec.FallThroughCodec;
@@ -28,6 +29,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;
@@ -215,5 +217,16 @@ public final class TestEntityWithCounterColumn_AchillesMeta extends AbstractEnti
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor.txt
index f51237a19..211254bc6 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor.txt
@@ -10,6 +10,7 @@ import com.google.common.reflect.TypeToken;
 import info.archinnov.achilles.generated.function.Date_Type;
 import info.archinnov.achilles.generated.function.Double_Type;
 import info.archinnov.achilles.generated.function.Long_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithCustomConstructor_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
 import info.archinnov.achilles.internals.codec.FallThroughCodec;
@@ -32,6 +33,7 @@ import java.lang.Long;
 import java.lang.Override;
 import java.lang.String;
 import java.lang.SuppressWarnings;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -252,5 +254,16 @@ public final class TestEntityWithCustomConstructor_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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor_with_declared_fields.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor_with_declared_fields.txt
index 455a722bc..981a63715 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor_with_declared_fields.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_custom_constructor_with_declared_fields.txt
@@ -10,6 +10,7 @@ import com.google.common.reflect.TypeToken;
 import info.archinnov.achilles.generated.function.Date_Type;
 import info.archinnov.achilles.generated.function.Double_Type;
 import info.archinnov.achilles.generated.function.Long_Type;
+import info.archinnov.achilles.generated.function.PartitionKeys_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithCustomConstructorAndDeclaredFields_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
 import info.archinnov.achilles.internals.codec.FallThroughCodec;
@@ -32,6 +33,7 @@ import java.lang.Long;
 import java.lang.Override;
 import java.lang.String;
 import java.lang.SuppressWarnings;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.List;
@@ -252,5 +254,16 @@ public final class TestEntityWithCustomConstructorAndDeclaredFields_AchillesMeta
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_implicit_field_parsing.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_implicit_field_parsing.txt
index f82e0abee..f394b0bf5 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_implicit_field_parsing.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_implicit_field_parsing.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.TestUDTImplicitFieldParsing_Type;
 import info.archinnov.achilles.generated.function.UUID_Type;
@@ -36,6 +37,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;
@@ -281,5 +283,16 @@ public final class TestEntityWithImplicitFieldParsing_AchillesMeta extends Abstr
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_simple_partition_key.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_simple_partition_key.txt
index e229d77f8..45fa19aa8 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_simple_partition_key.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_simple_partition_key.txt
@@ -7,6 +7,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.meta.entity.TestEntityWithSimplePartitionKey_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -29,6 +30,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;
@@ -216,5 +218,16 @@ public final class TestEntityWithSimplePartitionKey_AchillesMeta extends Abstrac
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_annotations.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_annotations.txt
index 682ffeee0..bb39bc438 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_annotations.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_annotations.txt
@@ -7,6 +7,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.meta.entity.TestEntityWithStaticAnnotations_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -29,6 +30,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;
@@ -245,5 +247,16 @@ public final class TestEntityWithStaticAnnotations_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("partition_key"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_column.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_column.txt
index 64ab7863a..682550549 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_column.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_column.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.TestEntityWithStaticColumn_AchillesMeta.ColumnsForFunctions;
@@ -32,6 +33,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;
@@ -278,5 +280,16 @@ public final class TestEntityWithStaticColumn_AchillesMeta extends AbstractEntit
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_counter_column.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_counter_column.txt
index a44d5e7c3..0938c3776 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_counter_column.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_entity_with_static_counter_column.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.UUID_Type;
 import info.archinnov.achilles.generated.meta.entity.TestEntityWithStaticCounterColumn_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -31,6 +32,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;
@@ -248,5 +250,16 @@ public final class TestEntityWithStaticCounterColumn_AchillesMeta extends Abstra
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_inherited_entity.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_inherited_entity.txt
index b319f096a..7c93ad921 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_inherited_entity.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_inherited_entity.txt
@@ -7,6 +7,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.meta.entity.TestEntityAsChild_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -29,6 +30,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;
@@ -245,5 +247,16 @@ public final class TestEntityAsChild_AchillesMeta extends AbstractEntityProperty
     }
       }
     ;
+
+    /**
+     * <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"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_view_meta.txt b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_view_meta.txt
index a42665aa8..6d38a14bf 100644
--- a/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_view_meta.txt
+++ b/achilles-core/src/test/resources/expected_code/entity_meta_builder/should_build_view_meta.txt
@@ -10,6 +10,7 @@ import com.google.common.reflect.TypeToken;
 import info.archinnov.achilles.generated.function.Date_Type;
 import info.archinnov.achilles.generated.function.Double_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.meta.entity.TestViewSensorByType_AchillesMeta.ColumnsForFunctions;
 import info.archinnov.achilles.internals.apt.annotations.AchillesMeta;
@@ -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;
@@ -248,5 +250,16 @@ public final class TestViewSensorByType_AchillesMeta extends AbstractViewPropert
     }
       }
     ;
+
+    /**
+     * <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("type"); }
+      })
+    ;
   }
 }
\ No newline at end of file
diff --git a/integration-test-2_1/src/test/java/info/archinnov/achilles/it/TestEntityWithCompositePartitionKey.java b/integration-test-2_1/src/test/java/info/archinnov/achilles/it/TestEntityWithCompositePartitionKey.java
index f03f62af9..30a2d3b6a 100644
--- a/integration-test-2_1/src/test/java/info/archinnov/achilles/it/TestEntityWithCompositePartitionKey.java
+++ b/integration-test-2_1/src/test/java/info/archinnov/achilles/it/TestEntityWithCompositePartitionKey.java
@@ -17,6 +17,8 @@
 package info.archinnov.achilles.it;
 
 import static info.archinnov.achilles.embedded.CassandraEmbeddedConfigParameters.DEFAULT_CASSANDRA_EMBEDDED_KEYSPACE_NAME;
+import static info.archinnov.achilles.generated.function.SystemFunctions.token;
+import static info.archinnov.achilles.generated.meta.entity.EntityWithCompositePartitionKey_AchillesMeta.COLUMNS;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
 import static org.assertj.core.api.Assertions.assertThat;
@@ -35,15 +37,13 @@
 
 import info.archinnov.achilles.generated.ManagerFactory;
 import info.archinnov.achilles.generated.ManagerFactoryBuilder;
-import info.archinnov.achilles.generated.function.SystemFunctions;
 import info.archinnov.achilles.generated.manager.EntityWithCompositePartitionKey_Manager;
 import info.archinnov.achilles.internals.entities.EntityWithCompositePartitionKey;
-import info.archinnov.achilles.internals.metamodel.functions.FunctionCall;
 import info.archinnov.achilles.junit.AchillesTestResource;
 import info.archinnov.achilles.junit.AchillesTestResourceBuilder;
 import info.archinnov.achilles.script.ScriptExecutor;
+import info.archinnov.achilles.type.TypedMap;
 import info.archinnov.achilles.type.tuples.Tuple2;
-import info.archinnov.achilles.type.tuples.Tuple3;
 
 public class TestEntityWithCompositePartitionKey {
 
@@ -199,4 +199,33 @@ public void should_dsl_select_with_IN_clause() throws Exception {
                 .collect(toList()))
                 .containsExactly("val1-1", "val2-1", "val2-3");
     }
+
+    @Test
+    public void should_dsl_select_using_token_function() throws Exception {
+        //Given
+        final long id = RandomUtils.nextLong(0L, Long.MAX_VALUE);
+        final UUID uuid = new UUID(1L, 1L);
+        final EntityWithCompositePartitionKey entity = new EntityWithCompositePartitionKey(id, uuid, "val");
+        manager.crud().insert(entity).execute();
+
+
+        //When
+        TypedMap typedMap = manager
+                .dsl()
+                .select()
+                .id()
+                .uuid()
+                .function(token(COLUMNS.PARTITION_KEYS), "tokens")
+                .fromBaseTable()
+                .where()
+                .id().Eq(id)
+                .uuid().Eq(uuid)
+                .getTypedMap();
+
+        //Then
+        assertThat(typedMap).isNotNull();
+        assertThat(typedMap).isNotEmpty();
+        assertThat(typedMap.<Long>getTyped("tokens")).isNotNull();
+
+    }
 }
diff --git a/integration-test-2_2/src/test/java/info/archinnov/achilles/it/TestFunctionCallsSimpleEntity.java b/integration-test-2_2/src/test/java/info/archinnov/achilles/it/TestFunctionCallsSimpleEntity.java
index 2340cef34..e9dea08e0 100644
--- a/integration-test-2_2/src/test/java/info/archinnov/achilles/it/TestFunctionCallsSimpleEntity.java
+++ b/integration-test-2_2/src/test/java/info/archinnov/achilles/it/TestFunctionCallsSimpleEntity.java
@@ -96,6 +96,8 @@ public void should_dsl_with_system_function_call() throws Exception {
                 .dsl()
                 .select()
                 .id()
+                .function(token(SimpleEntity_AchillesMeta.COLUMNS.ID), "tokens")
+                .function(token(SimpleEntity_AchillesMeta.COLUMNS.PARTITION_KEYS), "partitionTokens")
                 .function(toUnixTimestamp(SimpleEntity_AchillesMeta.COLUMNS.DATE), "dateAsLong")
                 .function(writetime(SimpleEntity_AchillesMeta.COLUMNS.VALUE), "writetimeOfValue")
                 .fromBaseTable()
@@ -107,6 +109,8 @@ public void should_dsl_with_system_function_call() throws Exception {
         //Then
         assertThat(typedMap).isNotNull();
         assertThat(typedMap).isNotEmpty();
+        assertThat(typedMap.<Long>getTyped("tokens")).isNotNull();
+        assertThat(typedMap.<Long>getTyped("partitiontokens")).isNotNull();
         assertThat(typedMap.<Long>getTyped("dateaslong")).isGreaterThan(0L);
         assertThat(typedMap.<Long>getTyped("writetimeofvalue")).isGreaterThan(date.getTime());
     }
diff --git a/integration-test-3_10/src/test/java/info/archinnov/achilles/it/TestEntityForGroupByIT.java b/integration-test-3_10/src/test/java/info/archinnov/achilles/it/TestEntityForGroupByIT.java
index 748db267d..4c139f196 100644
--- a/integration-test-3_10/src/test/java/info/archinnov/achilles/it/TestEntityForGroupByIT.java
+++ b/integration-test-3_10/src/test/java/info/archinnov/achilles/it/TestEntityForGroupByIT.java
@@ -304,4 +304,29 @@ public void should_select_sum_group_by_partition_keys_and_one_clustering() throw
 
 
     }
+
+    @Test
+    public void should_select_using_token_function() throws Exception {
+        //Given
+        final Long id = RandomUtils.nextLong(0L, Long.MAX_VALUE);
+        scriptExecutor.executeScriptTemplate("EntityForGroupBy/insert_multi_partitions.cql", ImmutableMap.of("id", id));
+
+        //When
+        TypedMap typedMap = manager
+                .dsl()
+                .select()
+                .id()
+                .function(SystemFunctions.token(COLUMNS.PARTITION_KEYS), "tokens")
+                .fromBaseTable()
+                .without_WHERE_Clause()
+                .limit(1)
+                .getTypedMap();
+
+        //Then
+        assertThat(typedMap).isNotNull();
+        assertThat(typedMap).isNotEmpty();
+        assertThat(typedMap.<Long>getTyped("tokens")).isNotNull();
+
+    }
+
 }