org.apache.gora.cassandra.serializers.AvroSerializer.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.gora.cassandra.serializers.AvroSerializer.java

Source

/*
 *  Licensed to the Apache Software Foundation (ASF) under one or more
 *  contributor license agreements.  See the NOTICE file distributed with
 *  this work for additional information regarding copyright ownership.
 *  The ASF licenses this file to You under the Apache License, Version 2.0
 *  (the "License"); you may not use this file except in compliance with
 *  the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package org.apache.gora.cassandra.serializers;

import com.datastax.driver.core.AbstractGettableData;
import com.datastax.driver.core.ColumnDefinitions;
import com.datastax.driver.core.ConsistencyLevel;
import com.datastax.driver.core.DataType;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.Row;
import com.datastax.driver.core.SimpleStatement;
import com.datastax.driver.core.UDTValue;
import com.datastax.driver.core.UserType;
import org.apache.avro.Schema;
import org.apache.avro.specific.SpecificData;
import org.apache.commons.lang.ArrayUtils;
import org.apache.gora.cassandra.bean.CassandraKey;
import org.apache.gora.cassandra.bean.Field;
import org.apache.gora.cassandra.query.CassandraResultSet;
import org.apache.gora.cassandra.store.CassandraClient;
import org.apache.gora.cassandra.store.CassandraMapping;
import org.apache.gora.persistency.Persistent;
import org.apache.gora.persistency.impl.PersistentBase;
import org.apache.gora.query.Query;
import org.apache.gora.query.Result;
import org.apache.gora.store.DataStore;
import org.apache.gora.util.GoraException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

/**
 * This class contains the operations relates to Avro Serialization.
 */
class AvroSerializer<K, T extends PersistentBase> extends CassandraSerializer {

    private static final Logger LOG = LoggerFactory.getLogger(AvroSerializer.class);

    private DataStore<K, T> cassandraDataStore;

    private Schema persistentSchema;

    AvroSerializer(CassandraClient cassandraClient, DataStore<K, T> dataStore, CassandraMapping mapping)
            throws GoraException {
        super(cassandraClient, dataStore.getKeyClass(), dataStore.getPersistentClass(), mapping);
        if (PersistentBase.class.isAssignableFrom(dataStore.getPersistentClass())) {
            persistentSchema = ((PersistentBase) dataStore.getBeanFactory().getCachedPersistent()).getSchema();
        } else {
            throw new GoraException("Unsupported persistent class, couldn't able to find the Avro schema.");
        }
        this.cassandraDataStore = dataStore;
        try {
            analyzePersistent();
        } catch (Exception e) {
            throw new GoraException("Error occurred while analyzing the persistent class, :" + e.getMessage());
        }
    }

    /**
     * {@inheritDoc}
     *
     * @throws Exception
     */
    protected void analyzePersistent() throws Exception {
        userDefineTypeMaps = new HashMap<>();
        for (Field field : mapping.getFieldList()) {
            String fieldType = field.getType();
            if (fieldType.contains("frozen")) {
                String udtType = fieldType.substring(fieldType.indexOf("<") + 1, fieldType.indexOf(">"));
                if (PersistentBase.class.isAssignableFrom(persistentClass)) {
                    Schema fieldSchema = persistentSchema.getField(field.getFieldName()).schema();
                    if (fieldSchema.getType().equals(Schema.Type.UNION)) {
                        for (Schema currentSchema : fieldSchema.getTypes()) {
                            if (currentSchema.getType().equals(Schema.Type.RECORD)) {
                                fieldSchema = currentSchema;
                                break;
                            }
                        }
                    }
                    String createQuery = CassandraQueryFactory.getCreateUDTTypeForAvro(mapping, udtType,
                            fieldSchema);
                    userDefineTypeMaps.put(udtType, createQuery);
                } else {
                    throw new RuntimeException(
                            "Unsupported Class for User Define Types, Please use PersistentBase class. field : "
                                    + udtType);
                }
            }
        }
    }

    /**
     * {@inheritDoc}
     *
     * @param query
     * @return
     */
    @Override
    public boolean updateByQuery(Query query) {
        List<Object> objectArrayList = new ArrayList<>();
        String cqlQuery = CassandraQueryFactory.getUpdateByQueryForAvro(mapping, query, objectArrayList,
                persistentSchema);
        ResultSet results;
        SimpleStatement statement;
        if (objectArrayList.size() == 0) {
            statement = new SimpleStatement(cqlQuery);
        } else {
            statement = new SimpleStatement(cqlQuery, objectArrayList.toArray());
        }
        if (writeConsistencyLevel != null) {
            statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
        }
        results = client.getSession().execute(statement);
        return results.wasApplied();
    }

    /**
     * {@inheritDoc}
     *
     * @param key
     * @param fields
     * @return
     */
    @Override
    public Persistent get(Object key, String[] fields) throws GoraException {
        try {
            if (fields == null) {
                fields = getFields();
            }
            ArrayList<String> cassandraKeys = new ArrayList<>();
            ArrayList<Object> cassandraValues = new ArrayList<>();
            AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
            String cqlQuery = CassandraQueryFactory.getSelectObjectWithFieldsQuery(mapping, fields, cassandraKeys);
            SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
            if (readConsistencyLevel != null) {
                statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
            }
            ResultSet resultSet = this.client.getSession().execute(statement);
            Iterator<Row> iterator = resultSet.iterator();
            ColumnDefinitions definitions = resultSet.getColumnDefinitions();
            T obj = null;
            if (iterator.hasNext()) {
                obj = cassandraDataStore.newPersistent();
                AbstractGettableData row = (AbstractGettableData) iterator.next();
                populateValuesToPersistent(row, definitions, obj, fields);
            }
            return obj;
        } catch (GoraException e) {
            throw e;
        } catch (Exception e) {
            throw new GoraException(e);
        }
    }

    /**
     * {@inheritDoc}
     *
     * @param key
     * @param persistent
     */
    @Override
    public void put(Object key, Persistent persistent) throws GoraException {
        try {
            if (persistent instanceof PersistentBase) {
                if (persistent.isDirty()) {
                    PersistentBase persistentBase = (PersistentBase) persistent;
                    ArrayList<String> fields = new ArrayList<>();
                    ArrayList<Object> values = new ArrayList<>();
                    AvroCassandraUtils.processKeys(mapping, key, fields, values);
                    for (Schema.Field f : persistentBase.getSchema().getFields()) {
                        String fieldName = f.name();
                        Field field = mapping.getFieldFromFieldName(fieldName);
                        if (field == null) {
                            LOG.debug("Ignoring {} adding field, {} field can't find in {} mapping",
                                    new Object[] { fieldName, fieldName, persistentClass });
                            continue;
                        }
                        if (persistent.isDirty(f.pos()) || mapping.getInlinedDefinedPartitionKey()
                                .equals(mapping.getFieldFromFieldName(fieldName))) {
                            Object value = persistentBase.get(f.pos());
                            String fieldType = field.getType();
                            if (fieldType.contains("frozen")) {
                                fieldType = fieldType.substring(fieldType.indexOf("<") + 1, fieldType.indexOf(">"));
                                UserType userType = client.getSession().getCluster().getMetadata()
                                        .getKeyspace(mapping.getKeySpace().getName()).getUserType(fieldType);
                                UDTValue udtValue = userType.newValue();
                                Schema udtSchema = f.schema();
                                if (udtSchema.getType().equals(Schema.Type.UNION)) {
                                    for (Schema schema : udtSchema.getTypes()) {
                                        if (schema.getType().equals(Schema.Type.RECORD)) {
                                            udtSchema = schema;
                                            break;
                                        }
                                    }
                                }
                                PersistentBase udtObjectBase = (PersistentBase) value;
                                for (Schema.Field udtField : udtSchema.getFields()) {
                                    Object udtFieldValue = AvroCassandraUtils.getFieldValueFromAvroBean(
                                            udtField.schema(), udtField.schema().getType(),
                                            udtObjectBase.get(udtField.name()), field);
                                    if (udtField.schema().getType().equals(Schema.Type.MAP)) {
                                        udtValue.setMap(udtField.name(), (Map) udtFieldValue);
                                    } else if (udtField.schema().getType().equals(Schema.Type.ARRAY)) {
                                        udtValue.setList(udtField.name(), (List) udtFieldValue);
                                    } else {
                                        udtValue.set(udtField.name(), udtFieldValue,
                                                (Class) udtFieldValue.getClass());
                                    }
                                }
                                value = udtValue;
                            } else {
                                value = AvroCassandraUtils.getFieldValueFromAvroBean(f.schema(),
                                        f.schema().getType(), value, field);
                            }
                            values.add(value);
                            fields.add(fieldName);
                        }
                    }
                    String cqlQuery = CassandraQueryFactory.getInsertDataQuery(mapping, fields);
                    SimpleStatement statement = new SimpleStatement(cqlQuery, values.toArray());
                    if (writeConsistencyLevel != null) {
                        statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
                    }
                    client.getSession().execute(statement);
                } else {
                    LOG.info("Ignored putting persistent bean {} in the store as it is neither "
                            + "new, neither dirty.", new Object[] { persistent });
                }
            } else {
                LOG.error("{} Persistent bean isn't extended by {} .",
                        new Object[] { this.persistentClass, PersistentBase.class });
            }
        } catch (Exception e) {
            throw new GoraException(e);
        }
    }

    /**
     * {@inheritDoc}
     *
     * @param key
     * @return
     */
    @Override
    public Persistent get(Object key) throws GoraException {
        try {
            ArrayList<String> cassandraKeys = new ArrayList<>();
            ArrayList<Object> cassandraValues = new ArrayList<>();
            AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
            String cqlQuery = CassandraQueryFactory.getSelectObjectQuery(mapping, cassandraKeys);
            SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
            if (readConsistencyLevel != null) {
                statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
            }
            ResultSet resultSet = client.getSession().execute(statement);
            Iterator<Row> iterator = resultSet.iterator();
            ColumnDefinitions definitions = resultSet.getColumnDefinitions();
            T obj = null;
            if (iterator.hasNext()) {
                obj = cassandraDataStore.newPersistent();
                AbstractGettableData row = (AbstractGettableData) iterator.next();
                populateValuesToPersistent(row, definitions, obj, mapping.getFieldNames());
            }
            return obj;
        } catch (Exception e) {
            throw new GoraException(e);
        }
    }

    /**
     * This method wraps result set data in to DataEntry and creates a list of DataEntry.
     **/
    private void populateValuesToPersistent(AbstractGettableData row, ColumnDefinitions columnDefinitions,
            PersistentBase base, String[] fields) {
        Object paramValue;
        for (String fieldName : fields) {
            Schema.Field avroField = base.getSchema().getField(fieldName);
            Field field = mapping.getFieldFromFieldName(fieldName);
            //to ignore unspecified fields in the mapping
            if (field == null || avroField == null) {
                continue;
            }
            Schema fieldSchema = avroField.schema();
            String columnName = field.getColumnName();
            paramValue = getValue(row, columnDefinitions.getType(columnName), columnName, fieldSchema);
            Object value = AvroCassandraUtils.getAvroFieldValue(paramValue, fieldSchema);
            base.put(avroField.pos(), value);
        }
    }

    private Object getValue(AbstractGettableData row, DataType columnType, String columnName, Schema schema) {
        Object paramValue;
        String dataType;
        switch (columnType.getName()) {
        case ASCII:
            paramValue = row.getString(columnName);
            break;
        case BIGINT:
            paramValue = row.isNull(columnName) ? null : row.getLong(columnName);
            break;
        case BLOB:
            paramValue = row.isNull(columnName) ? null : row.getBytes(columnName);
            break;
        case BOOLEAN:
            paramValue = row.isNull(columnName) ? null : row.getBool(columnName);
            break;
        case COUNTER:
            paramValue = row.isNull(columnName) ? null : row.getLong(columnName);
            break;
        case DECIMAL:
            paramValue = row.isNull(columnName) ? null : row.getDecimal(columnName);
            break;
        case DOUBLE:
            paramValue = row.isNull(columnName) ? null : row.getDouble(columnName);
            break;
        case FLOAT:
            paramValue = row.isNull(columnName) ? null : row.getFloat(columnName);
            break;
        case INET:
            paramValue = row.isNull(columnName) ? null : row.getInet(columnName).toString();
            break;
        case INT:
            paramValue = row.isNull(columnName) ? null : row.getInt(columnName);
            break;
        case TEXT:
            paramValue = row.getString(columnName);
            break;
        case TIMESTAMP:
            paramValue = row.isNull(columnName) ? null : row.getDate(columnName);
            break;
        case UUID:
            paramValue = row.isNull(columnName) ? null : row.getUUID(columnName);
            break;
        case VARCHAR:
            paramValue = row.getString(columnName);
            break;
        case VARINT:
            paramValue = row.isNull(columnName) ? null : row.getVarint(columnName);
            break;
        case TIMEUUID:
            paramValue = row.isNull(columnName) ? null : row.getUUID(columnName);
            break;
        case LIST:
            dataType = columnType.getTypeArguments().get(0).toString();
            paramValue = row.isNull(columnName) ? null
                    : row.getList(columnName, AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
            break;
        case SET:
            dataType = columnType.getTypeArguments().get(0).toString();
            paramValue = row.isNull(columnName) ? null
                    : row.getList(columnName, AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
            break;
        case MAP:
            dataType = columnType.getTypeArguments().get(1).toString();
            // Avro supports only String for keys
            paramValue = row.isNull(columnName) ? null
                    : row.getMap(columnName, String.class,
                            AvroCassandraUtils.getRelevantClassForCassandraDataType(dataType));
            break;
        case UDT:
            paramValue = row.isNull(columnName) ? null : row.getUDTValue(columnName);
            if (paramValue != null) {
                try {
                    PersistentBase udtObject = (PersistentBase) SpecificData
                            .newInstance(Class.forName(schema.getFullName()), schema);
                    for (Schema.Field f : udtObject.getSchema().getFields()) {
                        DataType dType = ((UDTValue) paramValue).getType().getFieldType(f.name());
                        Object fieldValue = getValue((UDTValue) paramValue, dType, f.name(), f.schema());
                        udtObject.put(f.pos(), fieldValue);
                    }
                    paramValue = udtObject;
                } catch (ClassNotFoundException e) {
                    throw new RuntimeException("Error occurred while populating data to " + schema.getFullName()
                            + " : " + e.getMessage());
                }
            }
            break;
        case TUPLE:
            paramValue = row.isNull(columnName) ? null : row.getTupleValue(columnName).toString();
            break;
        case CUSTOM:
            paramValue = row.isNull(columnName) ? null : row.getBytes(columnName);
            break;
        default:
            paramValue = row.getString(columnName);
            break;
        }
        return paramValue;
    }

    /**
     * {@inheritDoc}
     *
     * @param key
     * @return
     */
    @Override
    public boolean delete(Object key) throws GoraException {
        ArrayList<String> cassandraKeys = new ArrayList<>();
        ArrayList<Object> cassandraValues = new ArrayList<>();
        try {
            AvroCassandraUtils.processKeys(mapping, key, cassandraKeys, cassandraValues);
            String cqlQuery = CassandraQueryFactory.getDeleteDataQuery(mapping, cassandraKeys);
            SimpleStatement statement = new SimpleStatement(cqlQuery, cassandraValues.toArray());
            if (writeConsistencyLevel != null) {
                statement.setConsistencyLevel(ConsistencyLevel.valueOf(writeConsistencyLevel));
            }
            ResultSet resultSet = client.getSession().execute(statement);
            return resultSet.wasApplied();
        } catch (Exception e) {
            throw new GoraException(e);
        }
    }

    /**
     * {@inheritDoc}
     *
     * @param dataStore
     * @param query
     * @return
     */
    @Override
    public Result execute(DataStore dataStore, Query query) throws GoraException {
        try {
            List<Object> objectArrayList = new ArrayList<>();
            String[] fields = query.getFields();
            if (fields != null) {
                fields = (String[]) ArrayUtils.addAll(fields, mapping.getAllKeys());
            } else {
                fields = mapping.getAllFieldsIncludingKeys();
            }
            CassandraResultSet<K, T> cassandraResult = new CassandraResultSet<>(dataStore, query);
            String cqlQuery = CassandraQueryFactory.getExecuteQuery(mapping, query, objectArrayList, fields);
            ResultSet results;
            SimpleStatement statement;
            if (objectArrayList.size() == 0) {
                statement = new SimpleStatement(cqlQuery);
            } else {
                statement = new SimpleStatement(cqlQuery, objectArrayList.toArray());
            }
            if (readConsistencyLevel != null) {
                statement.setConsistencyLevel(ConsistencyLevel.valueOf(readConsistencyLevel));
            }
            results = client.getSession().execute(statement);
            Iterator<Row> iterator = results.iterator();
            ColumnDefinitions definitions = results.getColumnDefinitions();
            T obj;
            K keyObject;
            CassandraKey cassandraKey = mapping.getCassandraKey();
            while (iterator.hasNext()) {
                AbstractGettableData row = (AbstractGettableData) iterator.next();
                obj = cassandraDataStore.newPersistent();
                keyObject = cassandraDataStore.newKey();
                populateValuesToPersistent(row, definitions, obj, fields);
                if (cassandraKey != null) {
                    populateValuesToPersistent(row, definitions, (PersistentBase) keyObject,
                            cassandraKey.getFieldNames());
                } else {
                    Field key = mapping.getInlinedDefinedPartitionKey();
                    keyObject = (K) getValue(row, definitions.getType(key.getColumnName()), key.getColumnName(),
                            null);
                }
                cassandraResult.addResultElement(keyObject, obj);
            }
            return cassandraResult;
        } catch (Exception e) {
            throw new GoraException(e);
        }
    }

}