Examples with GettableByName com.datastax.oss.driver.api.core.data.GettableByName used on opensource projects

Example 1 with GettableByName

use of com.datastax.oss.driver.api.core.data.GettableByName in project java-driver by datastax.

the class DaoGetEntityMethodGenerator method generate.

@Override
public Optional<MethodSpec> generate() {
    // Validate the parameter: there must be exactly one, of type GettableByName or a ResultSet:
    if (methodElement.getParameters().size() != 1) {
        context.getMessager().error(methodElement, "Wrong number of parameters: %s methods must have exactly one", GetEntity.class.getSimpleName());
        return Optional.empty();
    }
    VariableElement parameterElement = methodElement.getParameters().get(0);
    String parameterName = parameterElement.getSimpleName().toString();
    TypeMirror parameterType = parameterElement.asType();
    boolean parameterIsGettable = context.getClassUtils().implementsGettableByName(parameterType);
    boolean parameterIsResultSet = context.getClassUtils().isSame(parameterType, ResultSet.class);
    boolean parameterIsAsyncResultSet = context.getClassUtils().isSame(parameterType, AsyncResultSet.class);
    if (!parameterIsGettable && !parameterIsResultSet && !parameterIsAsyncResultSet) {
        context.getMessager().error(methodElement, "Invalid parameter type: %s methods must take a %s, %s or %s", GetEntity.class.getSimpleName(), GettableByName.class.getSimpleName(), ResultSet.class.getSimpleName(), AsyncResultSet.class.getSimpleName());
        return Optional.empty();
    }
    // Validate the return type. Make sure it matches the parameter type
    Transformation transformation = null;
    TypeMirror returnType = methodElement.getReturnType();
    TypeElement entityElement = EntityUtils.asEntityElement(returnType, typeParameters);
    if (entityElement != null) {
        transformation = parameterIsGettable ? Transformation.NONE : Transformation.ONE;
    } else if (returnType.getKind() == TypeKind.DECLARED) {
        Element element = ((DeclaredType) returnType).asElement();
        if (context.getClassUtils().isSame(element, PagingIterable.class)) {
            if (!parameterIsResultSet) {
                context.getMessager().error(methodElement, "Invalid return type: %s methods returning %s must have an argument of type %s", GetEntity.class.getSimpleName(), PagingIterable.class.getSimpleName(), ResultSet.class.getSimpleName());
                return Optional.empty();
            }
            entityElement = EntityUtils.typeArgumentAsEntityElement(returnType, typeParameters);
            transformation = Transformation.MAP;
        } else if (context.getClassUtils().isSame(element, Stream.class)) {
            if (!parameterIsResultSet) {
                context.getMessager().error(methodElement, "Invalid return type: %s methods returning %s must have an argument of type %s", GetEntity.class.getSimpleName(), Stream.class.getSimpleName(), ResultSet.class.getSimpleName());
                return Optional.empty();
            }
            entityElement = EntityUtils.typeArgumentAsEntityElement(returnType, typeParameters);
            transformation = Transformation.STREAM;
        } else if (context.getClassUtils().isSame(element, MappedAsyncPagingIterable.class)) {
            if (!parameterIsAsyncResultSet) {
                context.getMessager().error(methodElement, "Invalid return type: %s methods returning %s must have an argument of type %s", GetEntity.class.getSimpleName(), MappedAsyncPagingIterable.class.getSimpleName(), AsyncResultSet.class.getSimpleName());
                return Optional.empty();
            }
            entityElement = EntityUtils.typeArgumentAsEntityElement(returnType, typeParameters);
            transformation = Transformation.MAP;
        }
    }
    if (entityElement == null) {
        context.getMessager().error(methodElement, "Invalid return type: " + "%s methods must return a %s-annotated class, or a %s, a %s or %s thereof", GetEntity.class.getSimpleName(), Entity.class.getSimpleName(), PagingIterable.class.getSimpleName(), Stream.class.getSimpleName(), MappedAsyncPagingIterable.class.getSimpleName());
        return Optional.empty();
    }
    // Generate the implementation:
    String helperFieldName = enclosingClass.addEntityHelperField(ClassName.get(entityElement));
    MethodSpec.Builder overridingMethodBuilder = GeneratedCodePatterns.override(methodElement, typeParameters);
    switch(transformation) {
        case NONE:
            overridingMethodBuilder.addStatement("return $L.get($L, $L)", helperFieldName, parameterName, lenient);
            break;
        case ONE:
            overridingMethodBuilder.addStatement("$T row = $L.one()", Row.class, parameterName).addStatement("return (row == null) ? null : $L.get(row, $L)", helperFieldName, lenient);
            break;
        case MAP:
            overridingMethodBuilder.addStatement("return $L.map(row -> $L.get(row, $L))", parameterName, helperFieldName, lenient);
            break;
        case STREAM:
            overridingMethodBuilder.addStatement("return $T.stream($L.map(row -> $L.get(row, $L)).spliterator(), false)", StreamSupport.class, parameterName, helperFieldName, lenient);
            break;
    }
    return Optional.of(overridingMethodBuilder.build());
}