Using Hibernate DAOs in DropWizard Tasks

I’ve previously blogged about DropWizard, a Java web framework for developing RESTful web services.

One of Dropwizard’s great features is its ability to easily write and deploy runtime administrative tasks for you application. These can run on either under a separate endpoint, like /admin or on an entirely different port. These tasks can be used for a whole range of maintenance operations for your running application, everything from triggering garbage collection, to generating reports, to clearing application caches.

Recently, I had a need to create a maintenance task that required database access. In the application in question, we were using DropWizard’s Hibernate integration to do database work. dropwizard-hibernate comes with a base DAO class, AbstractDAO<T> that provides typesafe database operations for your entity classes. For my database task, I wanted to be able to use my already existing DAO classes to do database queries.

From here out, I’ll dive into example land and show you how to create a simple task that uses a Hibernate DAO to archive users who haven’t logged into the system since before a given date. Here is a snippet from the DAO class in question. It includes a method that takes a timestamp and returns the count of users who have a lastLoggedIn property less than or equal to the provided timestamp.

public class UserDAO extends AbstractDAO {
    ...
    public List getUsersWhoHaveNotLoggedInSinceDate(Long timestamp) {
        Criteria criteria = criteria()
            .add(Restrictions.le("lastLogin", timestamp));
        return list(criteria);
    }
    ...
    public User upsert(User user) {
        return persist(user);
    }
    ...
}

And here is a basic version of the task we want to write, which uses the provided DAO.

public class ArchiveOldUsersTask extends Task {
    private UserDAO userDAO;

    public ArchiveOldUsersTask(UserDAO userDAO) {
        super("archive-old-users");
        this.userDAO = userDAO;
    }

    @Override
    public void execute(ImmutableMultimap parameters, PrintWriter output) throws Exception {
        Long timestamp = Long.valueOf(parameters.get("delete-before").asList.get(0));

        // Archive users based on last login date
        for (User userToArchive : userDAO.getUsersWhoHaveNotLoggedInSinceDate(timestamp)) {
            userToArchive.setActive(false);
            userDAO.upsert(userToArchive);
        }

        output.println("Finished archiving users");
        output.flush();
        output.close();
    }
}

The problem with using DAO classes in a task is that they require some setup of the hibernate session which is handled by DropWizard’s UnitOfWork annotation. UnitOfWork is designed to used with Jersey resource methods and does not work with DropWizard Tasks. This means that the above task will not work as written.

This seemed like a big problem, because it wasn’t at all clear how DropWizard set up up a hibernate session factory for DAO classes to use. Fortunately, the UnitOfWork annotation is actually very simple and easy to read, and I was quickly able to locate the relevant code in UnitOfWorkRequestDispatcher#dispatch.

The code was so short, in fact, that I just ended up duplicating it and using it to wrap the body of my task like so:

public class ArchiveOldUsersTask extends Task {
    private UserDAO userDAO;
    private SessionFactory sessionFactory;

    public ArchiveOldUsersTask(UserDAO userDAO, SessionFactory sessionFactory) {
        super("archive-old-users");
        this.userDAO = userDAO;
    }

    @Override
    public void execute(ImmutableMultimap parameters, PrintWriter output) throws Exception {
        Long timestamp = Long.valueOf(parameters.get("delete-before").asList.get(0));
        Session session = sessionFactory.openSession();
        try {
            ManagedSessionContext.bind(session);
            Transaction transaction = session.beginTransaction();
            try {
                // Archive users based on last login date
                for (User userToArchive : userDAO.getUsersWhoHaveNotLoggedInSinceDate(timestamp)) {
                    userToArchive.setActive(false);
                    userDAO.upsert(userToArchive);
                }
                transaction.commit();
            }
            catch (Exception e) {
                transaction.rollback();
                throw new RuntimeException(e);
            }
        } finally {
            session.close();
            ManagedSessionContext.unbind(sessionFactory);
        }
        output.println("Finished archiving users");
        output.flush();
        output.close();
    }
}

This turned out to be perfectly successful on the first go, and I now had the ability to use all of my Hibernate goodness within Dropwizard Tasks.