Forked from eintopf/CassandraSchemaGenerator.java
Last active
April 26, 2020 15:23
-
-
Save stuarthendren/6affdfe5a7aec4b7670e to your computer and use it in GitHub Desktop.
Since the Datastax driver currently does not support the Auto-Creation of table schema I have extended this class by https://gist.github.com/eintopf. It creates the create queries (strings). The execution has to be done in another step. The original problem originates from here: http://stackoverflow.com/questions/32953050/datastax-cassandra-java…
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import java.lang.reflect.Field; | |
import java.lang.reflect.Modifier; | |
import java.math.BigDecimal; | |
import java.math.BigInteger; | |
import java.net.InetAddress; | |
import java.nio.ByteBuffer; | |
import java.util.Date; | |
import java.util.HashMap; | |
import java.util.LinkedList; | |
import java.util.List; | |
import java.util.Map; | |
import java.util.UUID; | |
import com.datastax.driver.core.DataType; | |
import com.datastax.driver.mapping.annotations.ClusteringColumn; | |
import com.datastax.driver.mapping.annotations.Column; | |
import com.datastax.driver.mapping.annotations.PartitionKey; | |
import com.datastax.driver.mapping.annotations.Table; | |
import com.datastax.driver.mapping.annotations.UDT; | |
import com.google.common.base.Strings; | |
import com.google.common.collect.ImmutableList; | |
import com.google.gson.internal.Primitives; | |
public class CassandraSchemaGenerator { | |
private static final List<String> ignore = ImmutableList.of("serialVersionUID"); | |
private static final Map<Class<?>, DataType> dataTypes = new HashMap<>(); | |
static { | |
dataTypes.put(String.class, DataType.text()); | |
dataTypes.put(Long.class, DataType.bigint()); | |
dataTypes.put(ByteBuffer.class, DataType.blob()); | |
dataTypes.put(Boolean.class, DataType.cboolean()); | |
dataTypes.put(BigDecimal.class, DataType.decimal()); | |
dataTypes.put(Double.class, DataType.cdouble()); | |
dataTypes.put(Float.class, DataType.cfloat()); | |
dataTypes.put(InetAddress.class, DataType.inet()); | |
dataTypes.put(Integer.class, DataType.cint()); | |
dataTypes.put(Short.class, DataType.smallint()); | |
dataTypes.put(Date.class, DataType.timestamp()); | |
dataTypes.put(UUID.class, DataType.uuid()); | |
dataTypes.put(BigInteger.class, DataType.varint()); | |
} | |
public enum StructureType { | |
TABLE, TYPE | |
} | |
private Map<Class<?>, String> classToUdtName = new HashMap<>(); | |
private final String defaultKeyspace; | |
public CassandraSchemaGenerator() { | |
this(""); | |
} | |
public CassandraSchemaGenerator(String defaultKeyspace) { | |
this.defaultKeyspace = defaultKeyspace; | |
} | |
public String createTableQuery(Class<?> c) { | |
String requiredTypeQueries = createRequiredTypes(c); | |
return requiredTypeQueries + "\n" + createStructure(StructureType.TABLE, c); | |
} | |
private String createRequiredTypes(Class<?> c) { | |
StringBuilder stringBuilder = new StringBuilder(); | |
for (Class<?> clazz = c; clazz != null; clazz = clazz.getSuperclass()) { | |
for (Field f : clazz.getDeclaredFields()) { | |
Class<?> type = f.getType(); | |
if (classToUdtName.containsKey(type)) { | |
continue; | |
} | |
if (isCustomType(type)) { | |
stringBuilder.append(createTypeQuery(type)); | |
stringBuilder.append("\n"); | |
classToUdtName.put(type, extractSchemaName(type)); | |
} | |
} | |
} | |
return stringBuilder.toString(); | |
} | |
public String createTypeQuery(Class<?> c) { | |
return createStructure(StructureType.TYPE, c); | |
} | |
private String createStructure(StructureType structureType, Class<?> c) { | |
String schemaName = extractSchemaName(c); | |
if (schemaName == null) { | |
throw new RuntimeException("No table name defined for type / table: " + schemaName); | |
} | |
String structureName = structureType.name(); | |
StringBuilder stringBuilder = new StringBuilder("CREATE " + structureName + " "); | |
stringBuilder.append(schemaName); | |
stringBuilder.append(" (\n"); | |
int i = 0; | |
for (Class<?> clazz = c; clazz != null; clazz = clazz.getSuperclass()) { | |
for (Field f : clazz.getDeclaredFields()) { | |
if (ignore.contains(f.getName()) || Modifier.isStatic(f.getModifiers())) { | |
continue; | |
} | |
// Append comma on previous if not first field | |
if (i > 0) { | |
stringBuilder.append(",\n"); | |
} | |
String fieldName = getNameOfField(f); | |
// Append field name | |
stringBuilder.append("\t" + fieldName); | |
// Append Field Type | |
String fieldType = mapJavaToCassandraType(fieldName, f.getType()); | |
stringBuilder.append(" " + fieldType); | |
i++; | |
} | |
} | |
// Append Primary Key description | |
if (structureType.equals(StructureType.TABLE)) { | |
stringBuilder.append(",\n\t" + generatePrimaryKey(c)); | |
} | |
// Close create statement | |
stringBuilder.append("\n);"); | |
return stringBuilder.toString(); | |
} | |
private static String getNameOfField(Field f) { | |
Column nameAnnotation = f.getAnnotation(Column.class); | |
String fieldName = null; | |
if (nameAnnotation == null) { | |
fieldName = f.getName(); | |
} else { | |
fieldName = (nameAnnotation).name(); | |
} | |
return fieldName; | |
} | |
private static String generatePrimaryKey(Class<?> c) { | |
List<PartitionKey> partitionKeys = new LinkedList<>(); | |
Map<PartitionKey, String> partitionKeyToFieldName = new HashMap<>(); | |
List<String> clusterColumns = new LinkedList<>(); | |
for (Class<?> clazz = c; clazz != null; clazz = clazz.getSuperclass()) { | |
for (Field f : clazz.getDeclaredFields()) { | |
PartitionKey partitionKey = f.getAnnotation(PartitionKey.class); | |
ClusteringColumn clusterColumn = f.getAnnotation(ClusteringColumn.class); | |
if (partitionKey != null) { | |
partitionKeys.add(partitionKey); | |
partitionKeyToFieldName.put(partitionKey, getNameOfField(f)); | |
// partitonKeys. | |
} | |
if (clusterColumn != null) { | |
clusterColumns.add(getNameOfField(f)); | |
} | |
} | |
} | |
if (partitionKeys.size() == 0) { | |
throw new RuntimeException("Partition key annotations are required!"); | |
} | |
String partitonKey = partitionKeys.stream() | |
.sorted((p1, p2) -> Integer.compare(p1.value(), p2.value())) | |
.map((p) -> partitionKeyToFieldName.get(p)) | |
.reduce((p1, p2) -> p1 + ", " + p2).get(); | |
boolean multiPartitionKey = partitionKeys.size() > 1; | |
String result = "PRIMARY KEY ("; | |
if (multiPartitionKey) { | |
result += "("; | |
} | |
result += partitonKey; | |
if (multiPartitionKey) { | |
result += ")"; | |
} | |
if (clusterColumns.size() > 0) { | |
result += ", "; | |
result += clusterColumns.stream().reduce((c1, c2) -> c1 + ", " + c2).get(); | |
} | |
result += ")"; | |
return result; | |
} | |
private String extractSchemaName(Class<?> c) { | |
Table tableAnnotation = c.getAnnotation(Table.class); | |
UDT udtAnnotation = c.getAnnotation(UDT.class); | |
StringBuilder builder = new StringBuilder(); | |
if (tableAnnotation != null) { | |
String keyspace = tableAnnotation.keyspace(); | |
addKeyspace(builder, keyspace); | |
builder.append(tableAnnotation.name()); | |
} else if (udtAnnotation != null) { | |
String keyspace = udtAnnotation.keyspace(); | |
addKeyspace(builder, keyspace); | |
builder.append(udtAnnotation.name()); | |
} | |
return builder.toString(); | |
} | |
private void addKeyspace(StringBuilder builder, String keyspace) { | |
if (Strings.isNullOrEmpty(keyspace)) { | |
keyspace = defaultKeyspace; | |
} | |
builder.append(keyspace); | |
if (builder.length() > 0) { | |
builder.append("."); | |
} | |
} | |
private String mapJavaToCassandraType(String fieldName, Class<?> type) { | |
if (classToUdtName.containsKey(type)) { | |
return "frozen<" + classToUdtName.get(type) + ">"; | |
} | |
if (type.isArray()) { | |
Class<?> componentType = type.getComponentType(); | |
return DataType.list(getDataType(componentType)).toString(); | |
} | |
DataType dataType = getDataType(type); | |
return dataType.toString(); | |
} | |
private DataType getDataType(Class<?> type) { | |
if (type.isPrimitive()) { | |
type = Primitives.wrap(type); | |
} | |
DataType dataType = dataTypes.get(type); | |
return dataType; | |
} | |
private static boolean isCustomType(Class<?> type) { | |
String typeName = type.getName(); | |
return typeName.indexOf(".") != -1 && typeName.indexOf("java.") == -1; | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment