Skip to content

Commit

Permalink
UDT list support for Cassandra repositories (#72)
Browse files Browse the repository at this point in the history
  • Loading branch information
@Squiry
Squiry authored Feb 10, 2023
1 parent b0c299c commit 5c7b0c4
Show file tree
Hide file tree
Showing 33 changed files with 768 additions and 113 deletions.
3 changes: 3 additions & 0 deletions ...or-common/src/main/java/ru/tinkoff/kora/annotation/processor/common/CommonClassNames.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@

import com.squareup.javapoet.ClassName;

import java.util.List;


public class CommonClassNames {
public static final ClassName publisher = ClassName.get("org.reactivestreams", "Publisher");
Expand Down Expand Up @@ -36,4 +38,5 @@ public class CommonClassNames {
public static final ClassName refreshListener = ClassName.get("ru.tinkoff.kora.application.graph", "RefreshListener");

public static final ClassName koraGenerated = ClassName.get("ru.tinkoff.kora.common.annotation", "Generated");
public static final ClassName list = ClassName.get(List.class);
}
9 changes: 8 additions & 1 deletion ...in/java/ru/tinkoff/kora/database/annotation/processor/cassandra/CassandraNativeTypes.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
import javax.annotation.Nullable;
import java.math.BigDecimal;
import java.nio.ByteBuffer;
import java.time.Instant;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.util.List;
Expand Down Expand Up @@ -65,6 +66,11 @@ public class CassandraNativeTypes {
(rsName, i) -> CodeBlock.of("$L.getLocalDate($L)", rsName, i),
(stmt, var, i) -> CodeBlock.of("$L.setLocalDate($L, $L)", stmt, i, var)
);
var instant = CassandraNativeType.of(
TypeName.get(Instant.class),
(rsName, i) -> CodeBlock.of("$L.getInstant($L)", rsName, i),
(stmt, var, i) -> CodeBlock.of("$L.setInstant($L, $L)", stmt, i, var)
);

nativeTypes = List.of(
booleanPrimitive,
Expand All @@ -79,7 +85,8 @@ public class CassandraNativeTypes {
bigDecimal,
byteBuffer,
localDateTime,
localDate
localDate,
instant
);
}

Expand Down
2 changes: 2 additions & 0 deletions ...src/main/java/ru/tinkoff/kora/database/annotation/processor/cassandra/CassandraTypes.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,5 +30,7 @@ public class CassandraTypes {
public static final ClassName GETTABLE_BY_NAME = ClassName.get("com.datastax.oss.driver.api.core.data", "GettableByName");

public static final ClassName USER_DEFINED_TYPE = ClassName.get("com.datastax.oss.driver.api.core.type", "UserDefinedType");
public static final ClassName LIST_TYPE = ClassName.get("com.datastax.oss.driver.api.core.type", "ListType");
public static final ClassName UDT_VALUE = ClassName.get("com.datastax.oss.driver.api.core.data", "UdtValue");
public static final ClassName UDT_ANNOTATION = ClassName.get("ru.tinkoff.kora.database.cassandra", "UDT");
}
136 changes: 107 additions & 29 deletions ...kora/database/annotation/processor/cassandra/UserDefinedTypeResultExtractorGenerator.java
View file
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package ru.tinkoff.kora.database.annotation.processor.cassandra;

import com.squareup.javapoet.*;
import ru.tinkoff.kora.annotation.processor.common.CommonClassNames;
import ru.tinkoff.kora.annotation.processor.common.CommonUtils;
import ru.tinkoff.kora.database.annotation.processor.entity.DbEntity;

Expand All @@ -10,6 +11,7 @@
import javax.lang.model.util.Elements;
import javax.lang.model.util.Types;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Objects;

public class UserDefinedTypeResultExtractorGenerator {
Expand All @@ -24,6 +26,11 @@ public UserDefinedTypeResultExtractorGenerator(ProcessingEnvironment processingE
}

public void generate(TypeMirror type) {
this.generateMapper(type);
this.generateListMapper(type);
}

public void generateMapper(TypeMirror type) {
var element = types.asElement(type);
var typeName = TypeName.get(type);
var packageName = elements.getPackageOf(element);
Expand All @@ -32,6 +39,7 @@ public void generate(TypeMirror type) {
.addSuperinterface(ParameterizedTypeName.get(CassandraTypes.RESULT_COLUMN_MAPPER, typeName));
var constructor = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
var entity = Objects.requireNonNull(DbEntity.parseEntity(this.types, type));
this.addMappers(typeSpec, constructor, entity);

var apply = MethodSpec.methodBuilder("apply")
.addModifiers(Modifier.PUBLIC)
Expand All @@ -40,54 +48,124 @@ public void generate(TypeMirror type) {
.addParameter(int.class, "_index")
.returns(typeName);
apply.addStatement("var _object = _row.getUdtValue(_index)");
apply.addCode("\n");
for (var entityField : entity.entityFields()) {
var fieldName = entityField.element().getSimpleName().toString();
var nativeType = CassandraNativeTypes.findNativeType(TypeName.get(entityField.typeMirror()));
if (nativeType != null) {
apply.addStatement("var $N = $L", fieldName, nativeType.extract("_object", CodeBlock.of("_object.firstIndexOf($S)", entityField.columnName())));
} else {
var mapperName = "_" + fieldName + "_mapper";
// todo mapping annotation support?
var mapperType = ParameterizedTypeName.get(CassandraTypes.RESULT_COLUMN_MAPPER, TypeName.get(entityField.typeMirror()));
constructor.addParameter(mapperType, mapperName);
constructor.addStatement("this.$N = $N", mapperName, mapperName);
apply.beginControlFlow("if (_object == null)").addStatement("return null").endControlFlow();
apply.addStatement("var _type = ($T) _row.getType(_index)", CassandraTypes.USER_DEFINED_TYPE);
this.readIndexes(apply, entity);
this.readFields(apply, entity);
apply.addCode(this.buildEntity(entity));
apply.addStatement("return _result");

typeSpec.addField(mapperType, mapperName, Modifier.PRIVATE, Modifier.FINAL);
apply.addStatement("var $N = this.$N.apply(_object, _object.firstIndexOf($S))", fieldName, mapperName, entityField.columnName());
}
typeSpec.addMethod(apply.build());
typeSpec.addMethod(constructor.build());

try {
var javaFile = JavaFile.builder(packageName.getQualifiedName().toString(), typeSpec.build()).build();
CommonUtils.safeWriteTo(this.processingEnv, javaFile);
} catch (IOException e) {
throw new RuntimeException(e);
}
apply.addCode("\n");
}

public void generateListMapper(TypeMirror type) {
var element = types.asElement(type);
var typeName = TypeName.get(type);
var packageName = elements.getPackageOf(element);
var typeSpec = TypeSpec.classBuilder(CommonUtils.getOuterClassesAsPrefix(element) + element.getSimpleName() + "_List_CassandraRowColumnMapper")
.addModifiers(Modifier.PUBLIC, Modifier.FINAL)
.addSuperinterface(ParameterizedTypeName.get(CassandraTypes.RESULT_COLUMN_MAPPER, ParameterizedTypeName.get(CommonClassNames.list, typeName)));
var constructor = MethodSpec.constructorBuilder().addModifiers(Modifier.PUBLIC);
var entity = Objects.requireNonNull(DbEntity.parseEntity(this.types, type));
this.addMappers(typeSpec, constructor, entity);

var apply = MethodSpec.methodBuilder("apply")
.addModifiers(Modifier.PUBLIC)
.addAnnotation(Override.class)
.addParameter(CassandraTypes.GETTABLE_BY_NAME, "_row")
.addParameter(int.class, "_index")
.returns(ParameterizedTypeName.get(CommonClassNames.list, typeName));
apply.addStatement("var _list = _row.getList(_index, $T.class)", CassandraTypes.UDT_VALUE);
apply.beginControlFlow("if (_list == null)").addStatement("return null").endControlFlow();
apply.addStatement("var _listType = ($T) _row.getType(_index)", CassandraTypes.LIST_TYPE);
apply.addStatement("var _type = ($T) _listType.getElementType()", CassandraTypes.USER_DEFINED_TYPE);
this.readIndexes(apply, entity);
apply.addStatement("var _resultList = new $T<$T>(_list.size())", ArrayList.class, typeName);
apply.beginControlFlow("for (var _object : _list)");
this.readFields(apply, entity);
apply.addCode(this.buildEntity(entity));
apply.addStatement("_resultList.add(_result)");
apply.endControlFlow();
apply.addStatement("return _resultList");

typeSpec.addMethod(apply.build());
typeSpec.addMethod(constructor.build());

try {
var javaFile = JavaFile.builder(packageName.getQualifiedName().toString(), typeSpec.build()).build();
CommonUtils.safeWriteTo(this.processingEnv, javaFile);
} catch (IOException e) {
throw new RuntimeException(e);
}
}

private CodeBlock buildEntity(DbEntity entity) {
var b = CodeBlock.builder();
var typeName = TypeName.get(entity.typeMirror());
if (entity.entityType() == DbEntity.EntityType.RECORD) {
apply.addCode("return new $T(", typeName);
b.add("var _result = new $T(", typeName);
for (int i = 0; i < entity.entityFields().size(); i++) {
if (i > 0) {
apply.addCode(", ");
b.add(", ");
}
var entityField = entity.entityFields().get(i);
var fieldName = entityField.element().getSimpleName().toString();
apply.addCode("$N", fieldName);
b.add("$N", fieldName);
}
apply.addCode(");\n");
b.add(");\n");

} else {
apply.addStatement("var _result = new $T()", typeName);
b.addStatement("var _result = new $T()", typeName);
for (var entityField : entity.entityFields()) {
var fieldName = entityField.element().getSimpleName().toString();
apply.addStatement("_result.set$L($N)", CommonUtils.capitalize(fieldName), fieldName);
b.addStatement("_result.set$L($N)", CommonUtils.capitalize(fieldName), fieldName);
}
}
return b.build();
}

typeSpec.addMethod(apply.build());
typeSpec.addMethod(constructor.build());
private void readFields(MethodSpec.Builder apply, DbEntity entity) {
for (var entityField : entity.entityFields()) {
var fieldName = entityField.element().getSimpleName().toString();
var index = CodeBlock.of("$N", "_index_of_" + entityField.element().getSimpleName());
var nativeType = CassandraNativeTypes.findNativeType(TypeName.get(entityField.typeMirror()));
if (nativeType != null) {
apply.addStatement("var $N = $L", fieldName, nativeType.extract("_object", index));
} else {
var mapperName = "_" + fieldName + "_mapper";
apply.addStatement("var $N = this.$N.apply(_object, $L)", fieldName, mapperName, index);
}
}
apply.addCode("\n");
}

try {
var javaFile = JavaFile.builder(packageName.getQualifiedName().toString(), typeSpec.build()).build();
CommonUtils.safeWriteTo(this.processingEnv, javaFile);
} catch (IOException e) {
throw new RuntimeException(e);
private void addMappers(TypeSpec.Builder typeSpec, MethodSpec.Builder constructor, DbEntity entity) {
for (var entityField : entity.entityFields()) {
var nativeType = CassandraNativeTypes.findNativeType(TypeName.get(entityField.typeMirror()));
if (nativeType == null) {
var mapperName = "_" + entityField.element().getSimpleName() + "_mapper";
// todo mapping annotation support?
var mapperType = ParameterizedTypeName.get(CassandraTypes.RESULT_COLUMN_MAPPER, TypeName.get(entityField.typeMirror()));
constructor.addParameter(mapperType, mapperName);
constructor.addStatement("this.$N = $N", mapperName, mapperName);
typeSpec.addField(mapperType, mapperName, Modifier.PRIVATE, Modifier.FINAL);
}
}
}

private void readIndexes(MethodSpec.Builder apply, DbEntity entity) {
for (var entityField : entity.entityFields()) {
apply.addStatement("var $N = _type.firstIndexOf($S)", "_index_of_" + entityField.element().getSimpleName(), entityField.columnName());
}
apply.addCode("\n");
}
}

Loading

0 comments on commit 5c7b0c4

Please sign in to comment.