All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.appform.dropwizard.sharding.dao.RelationalDao Maven / Gradle / Ivy

There is a newer version: 3.0.7-1
Show newest version
/*
 * Copyright 2016 Santanu Sinha 
 *
 * Licensed 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 io.appform.dropwizard.sharding.dao;

import com.google.common.base.Preconditions;
import io.appform.dropwizard.sharding.utils.ShardCalculator;
import io.appform.dropwizard.sharding.utils.Transactions;
import io.dropwizard.hibernate.AbstractDAO;
import lombok.Builder;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import lombok.val;
import org.apache.commons.lang3.reflect.FieldUtils;
import org.hibernate.*;
import org.hibernate.criterion.DetachedCriteria;
import org.hibernate.criterion.Projections;
import org.hibernate.criterion.Restrictions;
import org.hibernate.query.Query;

import javax.persistence.Id;
import java.lang.reflect.Field;
import java.util.Collection;
import java.util.List;
import java.util.Optional;
import java.util.function.BooleanSupplier;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

/**
 * A dao used to work with entities related to a parent shard. The parent may or maynot be physically present.
 * A murmur 128 hash of the string parent key is used to route the save and retrieve calls from the proper shard.
 */
@Slf4j
public class RelationalDao implements ShardedDao {

    private final class RelationalDaoPriv extends AbstractDAO {

        private final SessionFactory sessionFactory;

        /**
         * Creates a new DAO with a given session provider.
         *
         * @param sessionFactory a session provider
         */
        public RelationalDaoPriv(SessionFactory sessionFactory) {
            super(sessionFactory);
            this.sessionFactory = sessionFactory;
        }

        T get(Object lookupKey) {
            return uniqueResult(currentSession()
                                .createCriteria(entityClass)
                                .add(
                                        Restrictions.eq(keyField.getName(), lookupKey)));
        }

        T save(T entity) {
            return persist(entity);
        }

        boolean saveAll(Collection entities) {
            for (T entity : entities) {
                persist(entity);
            }
            return true;
        }

        void update(T oldEntity, T entity) {
            currentSession().evict(oldEntity); //Detach .. otherwise update is a no-op
            currentSession().update(entity);
        }

        List select(SelectParamPriv selectParam) {
            val criteria = selectParam.criteria.getExecutableCriteria(currentSession());
            criteria.setFirstResult(selectParam.start);
            criteria.setMaxResults(selectParam.numRows);
            return list(criteria);
        }

        ScrollableResults scroll(ScrollParamPriv scrollDetails) {
            final Criteria criteria = scrollDetails.getCriteria().getExecutableCriteria(currentSession());
            return criteria.scroll(ScrollMode.FORWARD_ONLY);
        }

        long count(DetachedCriteria criteria) {
            return  (long)criteria.getExecutableCriteria(currentSession())
                            .setProjection(Projections.rowCount())
                            .uniqueResult();
        }

        public int update(final UpdateOperationMeta updateOperationMeta) {
            Query query = currentSession().createNamedQuery(updateOperationMeta.getQueryName());
            updateOperationMeta.getParams().forEach(query::setParameter);
            return query.executeUpdate();
        }

    }

    @Builder
    private static class SelectParamPriv {
        DetachedCriteria criteria;
        int start;
        int numRows;
    }

    @Builder
    private static class ScrollParamPriv {
        @Getter
        private DetachedCriteria criteria;
    }

    private List daos;
    private final Class entityClass;
    @Getter
    private final ShardCalculator shardCalculator;
    private final Field keyField;

    /**
     * Create a relational DAO.
     * @param sessionFactories List of session factories. One for each shard.
     * @param entityClass The class for which the dao will be used.
     * @param shardCalculator
     */
    public RelationalDao(
            List sessionFactories, Class entityClass,
            ShardCalculator shardCalculator ) {
        this.shardCalculator = shardCalculator;
        this.daos = sessionFactories.stream().map(RelationalDaoPriv::new).collect(Collectors.toList());
        this.entityClass = entityClass;

        Field fields[] = FieldUtils.getFieldsWithAnnotation(entityClass, Id.class);
        Preconditions.checkArgument(fields.length != 0, "A field needs to be designated as @Id");
        Preconditions.checkArgument(fields.length == 1, "Only one field can be designated as @Id");
        keyField = fields[0];
        if(!keyField.isAccessible()) {
            try {
                keyField.setAccessible(true);
            } catch (SecurityException e) {
                log.error("Error making key field accessible please use a public method and mark that as @Id", e);
                throw new IllegalArgumentException("Invalid class, DAO cannot be created.", e);
            }
        }
    }


    public Optional get(String parentKey, Object key) throws Exception {
        return Optional.ofNullable(get(parentKey, key, t-> t));
    }

    public U get(String parentKey, Object key, Function function) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, true, dao::get, key, function);
    }

    public Optional save(String parentKey, T entity) throws Exception {
        return Optional.ofNullable(save(parentKey, entity, t -> t));
    }

    public  U save(String parentKey, T entity, Function handler) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, false, dao::save, entity, handler);
    }

    public boolean saveAll(String parentKey, Collection entities) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, false, dao::saveAll, entities);
    }

     void save(LookupDao.LockedContext context, T entity) {
        RelationalDaoPriv dao = daos.get(context.getShardId());
        Transactions.execute(context.getSessionFactory(), false, dao::save, entity, t->t, false);
    }

     void save(LookupDao.LockedContext context, T entity, Function handler) {
        RelationalDaoPriv dao = daos.get(context.getShardId());
        Transactions.execute(context.getSessionFactory(), false, dao::save, entity, handler, false);
    }

     boolean update(LookupDao.LockedContext context, Object id, Function updater) {
        RelationalDaoPriv dao = daos.get(context.getShardId());
        return update(context.getSessionFactory(), dao, id, updater, false);
    }

     boolean update(LookupDao.LockedContext context,
                       DetachedCriteria criteria,
                       Function updater,
                       BooleanSupplier updateNext) {
        final RelationalDaoPriv dao = daos.get(context.getShardId());

        try {
            final ScrollParamPriv scrollParam = ScrollParamPriv.builder()
                    .criteria(criteria)
                    .build();

            return Transactions.execute(context.getSessionFactory(), true, dao::scroll, scrollParam, scrollableResults -> {
                boolean updateNextObject = true;
                try {
                    while(scrollableResults.next() && updateNextObject) {
                        final T entity = (T) scrollableResults.get(0);
                        if (null == entity) {
                            return false;
                        }
                        final T newEntity = updater.apply(entity);
                        if(null == newEntity) {
                            return false;
                        }
                        dao.update(entity, newEntity);
                        updateNextObject = updateNext.getAsBoolean();
                    }
                }
                finally {
                    scrollableResults.close();
                }
                return true;
            }, false);
        } catch (Exception e) {
            throw new RuntimeException("Error updating entity with scroll: " + criteria, e);
        }
    }

     List select(LookupDao.ReadOnlyContext context, DetachedCriteria criteria, int first, int numResults) throws Exception {
        final RelationalDaoPriv dao = daos.get(context.getShardId());
        SelectParamPriv selectParam = SelectParamPriv.builder()
                .criteria(criteria)
                .start(first)
                .numRows(numResults)
                .build();
        return Transactions.execute(context.getSessionFactory(), true, dao::select, selectParam, t -> t, false);
    }

    public boolean update(String parentKey, Object id, Function updater) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        return update(dao.sessionFactory, dao, id, updater, true);
    }

    public  U runInSession(String id, Function handler) {
        int shardId = shardCalculator.shardId(id);
        RelationalDaoPriv dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, handler);
    }

    private boolean update(SessionFactory daoSessionFactory, RelationalDaoPriv dao, Object id, Function updater, boolean completeTransaction){
        try {
            return Transactions.execute(daoSessionFactory, true, dao::get, id, (T entity) -> {
                if(null == entity) {
                    return false;
                }
                T newEntity = updater.apply(entity);
                if(null == newEntity) {
                    return false;
                }
                dao.update(entity, newEntity);
                return true;
            }, completeTransaction);
        } catch (Exception e) {
            throw new RuntimeException("Error updating entity: " + id, e);
        }
    }

    public boolean update(String parentKey, DetachedCriteria criteria, Function updater) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        try {
            SelectParamPriv selectParam = SelectParamPriv.builder()
                                                .criteria(criteria)
                                                .start(0)
                                                .numRows(1)
                                                .build();
            return Transactions., SelectParamPriv, Boolean>execute(dao.sessionFactory, true, dao::select, selectParam, (List entityList) -> {
                if(entityList == null || entityList.isEmpty()) {
                    return false;
                }
                T oldEntity = entityList.get(0);
                if(null == oldEntity) {
                    return false;
                }
                T newEntity = updater.apply(oldEntity);
                if(null == newEntity) {
                    return false;
                }
                dao.update(oldEntity, newEntity);
                return true;
            });
        } catch (Exception e) {
            throw new RuntimeException("Error updating entity with criteria: " + criteria, e);
        }
    }


    public int updateUsingQuery(String parentKey, UpdateOperationMeta updateOperationMeta) {
        int shardId = shardCalculator.shardId(parentKey);
        val dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, false, dao::update, updateOperationMeta);
    }

    public  int updateUsingQuery(LookupDao.LockedContext lockedContext, UpdateOperationMeta updateOperationMeta) {
        val dao = daos.get(lockedContext.getShardId());
        return Transactions.execute(lockedContext.getSessionFactory(), false, dao::update, updateOperationMeta, false);
    }


     boolean createOrUpdate(LookupDao.LockedContext context,
                               DetachedCriteria criteria,
                               Function updater,
                               Supplier entityGenerator) {
        final RelationalDaoPriv dao = daos.get(context.getShardId());

        try {
            final SelectParamPriv selectParam = SelectParamPriv.builder()
                    .criteria(criteria)
                    .start(0)
                    .numRows(1)
                    .build();

            return Transactions., SelectParamPriv, Boolean>execute(context.getSessionFactory(), true, dao::select, selectParam, (List entityList) -> {
                if(entityList == null || entityList.isEmpty()) {
                    Preconditions.checkNotNull(entityGenerator, "Entity generator can't be null");
                    final T newEntity = entityGenerator.get();
                    Preconditions.checkNotNull(newEntity, "Generated entity can't be null");
                    dao.save(newEntity);
                    return true;
                }

                final T oldEntity = entityList.get(0);
                if(null == oldEntity) {
                    return false;
                }
                final T newEntity = updater.apply(oldEntity);
                if(null == newEntity) {
                    return false;
                }
                dao.update(oldEntity, newEntity);
                return true;
            }, false);
        } catch (Exception e) {
            throw new RuntimeException("Error updating entity with criteria: " + criteria, e);
        }
    }

    public boolean updateAll(String parentKey, int start, int numRows, DetachedCriteria criteria, Function updater) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        try {
            SelectParamPriv selectParam = SelectParamPriv.builder()
                    .criteria(criteria)
                    .start(start)
                    .numRows(numRows)
                    .build();
            return Transactions., SelectParamPriv, Boolean>execute(dao.sessionFactory, true, dao::select, selectParam, entityList -> {
                if (entityList == null || entityList.isEmpty()) {
                    return false;
                }
                for (T oldEntity : entityList) {
                    if (null == oldEntity) {
                        return false;
                    }
                    T newEntity = updater.apply(oldEntity);
                    if (null == newEntity) {
                        return false;
                    }
                    dao.update(oldEntity, newEntity);
                }
                return true;
            });
        } catch (Exception e) {
            throw new RuntimeException("Error updating entity with criteria: " + criteria, e);
        }
    }

    public List select(String parentKey, DetachedCriteria criteria, int first, int numResults) throws Exception {
        return select(parentKey, criteria, first, numResults, t-> t);
    }

    public U select(String parentKey, DetachedCriteria criteria, int first, int numResults, Function, U> handler) throws Exception {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        SelectParamPriv selectParam = SelectParamPriv.builder()
                .criteria(criteria)
                .start(first)
                .numRows(numResults)
                .build();
        return Transactions.execute(dao.sessionFactory, true, dao::select, selectParam, handler);
    }

    public long count(String parentKey, DetachedCriteria criteria) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        return Transactions.execute(dao.sessionFactory, true, dao::count, criteria);
    }


    public boolean exists(String parentKey, Object key) {
        int shardId = shardCalculator.shardId(parentKey);
        RelationalDaoPriv dao = daos.get(shardId);
        Optional result = Transactions.executeAndResolve(dao.sessionFactory, true, dao::get, key);
        return result.isPresent();
    }

    /**
     * Queries using the specified criteria across all shards and returns the counts of rows satisfying the criteria.
     * Note: This method runs the query serially and it's usage is not recommended.
     * @param criteria The select criteria
     * @return List of counts in each shard
     */
    public List countScatterGather(DetachedCriteria criteria) {
        return daos.stream().map(dao -> {
            try {
                return Transactions.execute(dao.sessionFactory, true, dao::count, criteria);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }).collect(Collectors.toList());
    }

    public List scatterGather(DetachedCriteria criteria, int start, int numRows) {
        return daos.stream().map(dao -> {
            try {
                SelectParamPriv selectParam = SelectParamPriv.builder()
                        .criteria(criteria)
                        .start(start)
                        .numRows(numRows)
                        .build();
                return Transactions.execute(dao.sessionFactory, true, dao::select, selectParam);
            } catch (Exception e) {
                throw new RuntimeException(e);
            }
        }).flatMap(Collection::stream).collect(Collectors.toList());
    }

    protected Field getKeyField() {
        return this.keyField;
    }
}