DATAMONGO-1142

gistfile1.java

public class InheritanceAwareMongoRepository<T, ID extends Serializable>
        extends SimpleMongoRepository<T, ID> {

    private final MongoOperations mongoOperations;
    private final MongoEntityInformation<T, ID> entityInformation;

    public InheritanceAwareMongoRepository(
            MongoEntityInformation<T, ID> metadata,
            MongoOperations mongoOperations) {
        super(metadata, mongoOperations);
        this.mongoOperations = mongoOperations;
        this.entityInformation = metadata;
    }

    @Override
    public T findOne(ID id) {
        Assert.notNull(id, "The given id must not be null!");
        return mongoOperations.findOne(getIdQuery(id),
                entityInformation.getJavaType(),
                entityInformation.getCollectionName());
    }

    private Query getIdQuery(Object id) {
        return new Query(getIdCriteria(id)).addCriteria(getClassCriteria());
    }

    private Criteria getIdCriteria(Object id) {
        return where(entityInformation.getIdAttribute()).is(id);
    }

    private Criteria getClassCriteria() {
        return where("_class").is(entityInformation.getJavaType().getTypeName());
    }

    @Override
    public boolean exists(ID id) {
        Assert.notNull(id, "The given id must not be null!");
        return mongoOperations.exists(getIdQuery(id),
                entityInformation.getJavaType(),
                entityInformation.getCollectionName());
    }

    @Override
    public long count() {
        return mongoOperations.getCollection(
                entityInformation.getCollectionName()).count(
                        getClassCriteria().getCriteriaObject());
    }

    @Override
    public void delete(ID id) {
        Assert.notNull(id, "The given id must not be null!");
        mongoOperations.remove(getIdQuery(id),
                entityInformation.getJavaType(),
                entityInformation.getCollectionName());
    }

    @Override
    public void delete(T entity) {
        Assert.notNull(entity, "The given entity must not be null!");
        delete(entityInformation.getId(entity));
    }

    @Override
    public void delete(Iterable<? extends T> entities) {
        Assert.notNull(entities, "The given Iterable of entities not be null!");
        for (T entity : entities) {
            delete(entity);
        }
    }

    @Override
    public void deleteAll() {
        mongoOperations.remove(new Query(getClassCriteria()),
                entityInformation.getCollectionName());
    }

    @Override
    public List<T> findAll() {
        return findAll(new Query());
    }

    @Override
    public Iterable<T> findAll(Iterable<ID> ids) {

        Set<ID> parameters = new HashSet<>(tryDetermineRealSizeOrReturn(ids, 10));
        for (ID id : ids) {
            parameters.add(id);
        }

        return findAll(new Query(new Criteria(
                entityInformation.getIdAttribute()).in(parameters)));
    }

    @Override
    public Page<T> findAll(final Pageable pageable) {

        Long count = count();
        List<T> list = findAll(new Query().with(pageable));

        return new PageImpl<>(list, pageable, count);
    }

    @Override
    public List<T> findAll(Sort sort) {
        return findAll(new Query().with(sort));
    }

    private List<T> findAll(Query query) {

        if (query == null) {
            return Collections.emptyList();
        }

        query.addCriteria(getClassCriteria());

        return mongoOperations.find(query, entityInformation.getJavaType(),
                entityInformation.getCollectionName());
    }

    private static int tryDetermineRealSizeOrReturn(
            Iterable<?> iterable, int defaultSize) {
        return iterable == null ? 0 : (iterable instanceof Collection)
                ? ((Collection<?>) iterable).size() : defaultSize;
    }

}