com.yammer.dropwizard.tasks.TaskServlet Maven / Gradle / Ivy

Go to download
package com.yammer.dropwizard.tasks;

import com.google.common.collect.ImmutableMap;
import com.google.common.collect.ImmutableMultimap;
import com.yammer.dropwizard.logging.Log;

import javax.servlet.ServletException;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import javax.ws.rs.core.MediaType;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.Enumeration;

/**
 * A servlet which provides access to administrative {@link Task}s. It only responds to {@code POST}
 * requests, since most {@link Task}s aren't side-effect free, and passes along the query string
 * parameters of the request to the task as a multimap.
 *
 * @see Task
 */
public class TaskServlet extends HttpServlet {
    private static final long serialVersionUID = 7404713218661358124L;
    private static final Log LOG = Log.forClass(TaskServlet.class);
    private final ImmutableMap tasks;

    /**
     * Creates a new {@link TaskServlet} given the provided set of {@link Task} instances.
     *
     * @param tasks a series of tasks which the servlet will provide access to
     */
    public TaskServlet(Iterable tasks) {
        final ImmutableMap.Builder builder = ImmutableMap.builder();
        for (Task task : tasks) {
            builder.put("/tasks/" + task.getName(), task);
        }
        this.tasks = builder.build();
    }

    @Override
    protected void doPost(HttpServletRequest req,
                          HttpServletResponse resp) throws ServletException, IOException {
        final Task task = tasks.get(req.getRequestURI());
        if (task != null) {
            try {
                resp.setContentType(MediaType.TEXT_PLAIN);
                final PrintWriter output = resp.getWriter();
                try {
                    task.execute(getParams(req), output);
                } finally {
                    output.close();
                }
            } catch (Exception e) {
                LOG.error(e, "Error running {}", task.getName());
                resp.sendError(HttpServletResponse.SC_INTERNAL_SERVER_ERROR);
            }
        } else {
            resp.sendError(HttpServletResponse.SC_NOT_FOUND);
        }
    }

    private static ImmutableMultimap getParams(HttpServletRequest req) {
        final ImmutableMultimap.Builder results = ImmutableMultimap.builder();
        final Enumeration names = req.getParameterNames();
        while (names.hasMoreElements()) {
            final String name = names.nextElement();
            final String[] values = req.getParameterValues(name);
            results.putAll(name, values);
        }
        return results.build();
    }
}