package com.hubspot.singularity.guice;
import static com.google.common.base.Preconditions.checkNotNull;
import io.dropwizard.jersey.jackson.JacksonMessageBodyProvider;
import io.dropwizard.setup.Environment;
import java.io.IOException;
import java.io.OutputStream;
import java.lang.annotation.Annotation;
import java.lang.reflect.Type;
import javax.inject.Inject;
import javax.ws.rs.Produces;
import javax.ws.rs.core.Context;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.MultivaluedMap;
import javax.ws.rs.core.UriInfo;
import javax.ws.rs.ext.Provider;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SharedMetricRegistries;
import com.codahale.metrics.Timer;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.hubspot.jackson.jaxrs.PropertyFilter;
import com.hubspot.jackson.jaxrs.PropertyFiltering;
@Provider
@Produces(MediaType.APPLICATION_JSON)
public class GuicePropertyFilteringMessageBodyWriter extends JacksonMessageBodyProvider {
private static final Logger LOG = LoggerFactory.getLogger(GuicePropertyFilteringMessageBodyWriter.class);
@Context
private volatile UriInfo uriInfo;
private final Environment environment;
private final ObjectMapper objectMapper;
@Inject
public GuicePropertyFilteringMessageBodyWriter(final Environment environment, final ObjectMapper objectMapper) {
super(objectMapper, environment.getValidator());
this.environment = checkNotNull(environment, "environment is null");
this.objectMapper = checkNotNull(objectMapper, "objectMapper is null");
}
@Override
public boolean isWriteable(final Class<?> type, final Type genericType, final Annotation[] annotations, final MediaType mediaType) {
return hasMatchingMediaType(mediaType) &&
filteringEnabled(annotations) &&
super.isWriteable(type, genericType, annotations, mediaType);
}
@Override
public long getSize(final Object object, final Class<?> type, final Type genericType, final Annotation[] annotations, final MediaType mediaType) {
return -1;
}
@Override
public void writeTo(final Object o, final Class<?> type, final Type genericType, final Annotation[] annotations, final MediaType mediaType,
final MultivaluedMap<String, Object> httpHeaders, final OutputStream os) throws IOException {
final PropertyFiltering annotation = findPropertyFiltering(annotations);
final PropertyFilter propertyFilter = new PropertyFilter(uriInfo.getQueryParameters().get(annotation.using()));
if (!propertyFilter.hasFilters()) {
super.writeTo(o, type, genericType, annotations, mediaType, httpHeaders, os);
return;
}
final Timer timer = getTimer();
final Timer.Context context = timer.time();
try {
final JsonNode tree = objectMapper.valueToTree(o);
propertyFilter.filter(tree);
super.writeTo(tree, tree.getClass(), tree.getClass(), annotations, mediaType, httpHeaders, os);
} finally {
context.stop();
}
}
private Timer getTimer() {
return getMetricRegistry().timer(MetricRegistry.name(GuicePropertyFilteringMessageBodyWriter.class, "filter"));
}
private MetricRegistry getMetricRegistry() {
MetricRegistry registry = environment.metrics();
if (registry == null) {
LOG.warn("No environment metrics found!");
registry = SharedMetricRegistries.getOrCreate("com.hubspot");
}
return registry;
}
private static boolean filteringEnabled(final Annotation... annotations) {
return findPropertyFiltering(annotations) != null;
}
private static PropertyFiltering findPropertyFiltering(final Annotation... annotations) {
for (final Annotation annotation : annotations) {
if (annotation.annotationType() == PropertyFiltering.class) {
return (PropertyFiltering) annotation;
}
}
return null;
}
}