Automatic query modification with Spring Data Mongo

Introduction

Today I want to show how one can automatically apply additional restrictions on MongoDB Queries with MongoTemplate or Spring Data Mongo.

Implementation

First, let’s introduce interface QueryModifier.


import org.springframework.data.mongodb.core.query.Query;

public interface QueryModifier {

    /**
     * Modifies source query according to the rules.
     * @param query Source query to modify.
     * @param collectionName name of the collection against which query will be executed.
     * @return Modified query.
     */
    Query modify(Query query, String collectionName);
}

Implementations of this interface will pickup original query and modify it somehow. In order to make it work we need inject this modification after query was generated by Spring Data but before query is sent to the MongoDB instance. The perfect place for this is MongoTemplate`.

So, let’s create subclass of MongoTemplate and override find() and executeQuery() methods.


import com.mongodb.*;
import org.slf4j.*;
import org.springframework.data.authentication.UserCredentials;
import org.springframework.data.mongodb.MongoDbFactory;
import org.springframework.data.mongodb.core.*;

import java.util.*;

public class RestrictingMongoTemplate extends MongoTemplate {

    private List<? extends QueryModifier> queryModifiers = new ArrayList<>();

    /** omitting constructors */
    
    public void setQueryModifiers(List<? extends QueryModifier> queryModifiers) {
        this.queryModifiers = queryModifiers;
    }

    @Override
    public void executeQuery(Query query, String collectionName, DocumentCallbackHandler dch) {
        for (QueryModifier queryModifier : queryModifiers) {
            query = queryModifier.modify(query, collectionName);
        }

        super.executeQuery(query, collectionName, dch);
   }

    public <T> List<T> find(Query query, Class<T> entityClass, String collectionName) {
        int limit = query.getLimit();

        for (QueryModifier queryModifier : queryModifiers) {
            query = queryModifier.modify(query, collectionName);
        }

        query.limit(limit);

        return super.find(query, entityClass, collectionName);
    }

}

Let’s implement simple QueryModifier which will add restriction to return only active documents from Mongo:


import model.Document;
import org.springframework.data.mongodb.core.query.Criteria;
import org.springframework.data.mongodb.core.query.Query;

public class ActiveQueryModifier implements QueryModifier {

    @Override
    public Query modify(Query query, String collectionName) {
        if (Document.COLLECTION_NAME.equals(collectionName)) {
            return query.addCriteria(Criteria.where("active").is(Boolean.TRUE));
        }
        return query;
    }
}

Next step is to provide our custom implementation of MongoTemplate to Spring Data:

Example of Spring Configuration:


@PropertySource("classpath:application.properties")
@EnableMongoRepositories(
        basePackages = "model",
        mongoTemplateRef = "mongoTemplate")
@Configuration("configuration")
public class ServiceConfiguration extends AbstractMongoConfiguration {

    @Value("${mongo.db.name}")
    protected String mongoDatabaseName;

    @Override
    protected String getDatabaseName() {
        return mongoDatabaseName;
    }

    @Override
    public RestrictingMongoTemplate mongoTemplate() throws UnknownHostException {
        RestrictingMongoTemplate mongoTemplate =
                new RestrictingMongoTemplate(mongo(), mongoDatabaseName);

        mongoTemplate.setQueryModifiers(Arrays.asList(
                new ActiveQueryModifier()
        ));

        return mongoTemplate;    
    }

    // creation of Mongo bean omitted.
}

Oleg Atamanenko

thoughts about programming

Automatic query modification with Spring Data Mongo

Introduction

Implementation