Auto generated ID with Morphia and MongoDB

Following Fredriks post on unit testing Morphia and MongoDB, here’s a short how-to (with sources for the lazy) for using automatically generated integer ID with Morphia.

The problem is this:

• In SQL database we’re used to integer fields that can be safely incremented server-side, that can be used as entity ID. Not so in MongoDB.

MongoDB only accepts one ID type for a document, and that is the ObjectId. While this is no doubt fine for MongoDB, as Java developers it stings a bit. So here’s what we need to do:

• Creating a separate document with an ID counter per entity we want to use.

So let’s dive in shall we. First, let’s create an entity:

@Entity
public class User {

    @Id
    private Long id;

    public User(Long id) {
        this.id = id;
    }

    public User() { }

    [...]
}

Beware: If you’re new to Morphia, watch out, these annotations are not JPA as you would expect, they are Morphia annotations.

Now, this won’t work as MongoDB will complain about the ID type. But before we solve it, let’s create a simple DAO shall we?

public class UserDao {

    private Datastore store;

    public UserDao(Datastore store) {
        this.store = store;
    }

    public User get(Long id) {
        return store.get(User.class, id);
    }

    public User save(User u) {
        store.save(u);
        return u;
    }
}

This still won’t work as we haven’t solved the ID problem. In order to do that, we’ll first create a new document type, which we’ll use to store ID counters per entity we want to use (because you do want more than one entity don’t you?).

@Entity
public class EntityId {

    @Id
    private String className; 

    // this is the actual ID counter, will
    // be incremented atomically
    private Long counter = 1L;

    public EntityId() { }

    public EntityId(String className) {
        this.className = className;
    }

    [...]
}

The EntityId will store a counter for each entity collection we use. So the “User” type will have one EntityId document in the database, and the counter will be atomically updated when we create new users. So when the entity is saved, if it doesn’t have an ID, create one (in the DAO):

public User save(User u) {
    if(u.getId() == null) {
        Long id = generateId(u);
        u.setId(id);
    }
    store.save(u);
    return u;
}

// --- PROTECTED METHODS --- //

protected Long generateId(User entity) {
    // lookup the collection name for the entity
    String collName = store.getCollection(getClass()).getName();
    // find any existing counters for the type
    Query<EntityId> q = store.find(EntityId.class, "_id", collName);
    // create an update operation which increments the counter
    UpdateOperations<EntityId> update = store.createUpdateOperations(EntityId.class).inc("counter");
    // execute on server, if not found null is return,
    // else the counter is incremented atomically
    EntityId counter = store.findAndModify(q, update);
    if (counter == null) {
        // so just create one
        counter = new EntityId(collName);
        store.save(counter);
    }
    // return new id
    return counter.getCounter();
}

And that is, as they say, it. If you’ve read my friend Fredriks excellent post on how to unit test Morphia just go ahead and test it. Or if you’re lazy, download my sources, which contains a ready-made unit test and are mavenized, and spin it yourself.

Next time I’ll extend this little experiment to custom data type converters as well. Stay tuned!