AirpalApplicationBase.java

package com.airbnb.airpal;

import com.airbnb.airpal.core.AirpalUserFactory;
import com.airbnb.airpal.core.health.PrestoHealthCheck;
import com.airbnb.airpal.resources.ExecuteResource;
import com.airbnb.airpal.resources.FilesResource;
import com.airbnb.airpal.resources.HealthResource;
import com.airbnb.airpal.resources.PingResource;
import com.airbnb.airpal.resources.QueriesResource;
import com.airbnb.airpal.resources.QueryResource;
import com.airbnb.airpal.resources.ResultsPreviewResource;
import com.airbnb.airpal.resources.S3FilesResource;
import com.airbnb.airpal.resources.SessionResource;
import com.airbnb.airpal.resources.TablesResource;
import com.airbnb.airpal.resources.UserResource;
import com.airbnb.airpal.resources.UsersResource;
import com.airbnb.airpal.resources.sse.SSEEventSourceServlet;
import com.google.inject.AbstractModule;
import com.google.inject.Guice;
import com.google.inject.Injector;
import com.google.inject.Stage;
import io.dropwizard.Application;
import io.dropwizard.Bundle;
import io.dropwizard.ConfiguredBundle;
import io.dropwizard.assets.AssetsBundle;
import io.dropwizard.db.DataSourceFactory;
import io.dropwizard.flyway.FlywayBundle;
import io.dropwizard.flyway.FlywayFactory;
import io.dropwizard.jetty.BiDiGzipHandler;
import io.dropwizard.setup.Bootstrap;
import io.dropwizard.setup.Environment;
import io.dropwizard.views.ViewBundle;
import org.eclipse.jetty.server.Handler;

import javax.servlet.ServletRegistration;

import java.util.Arrays;

import static org.glassfish.jersey.message.MessageProperties.IO_BUFFER_SIZE;

public abstract class AirpalApplicationBase<T extends AirpalConfiguration>
        extends Application<T>
{
    private static final String SERVER_SENT_EVENTS = "text/event-stream";

    protected Injector injector;

    @Override
    public void initialize(Bootstrap<T> bootstrap)
    {
        for (ConfiguredBundle<T> configuredBundle : getConfiguredBundles()) {
            bootstrap.addBundle(configuredBundle);
        }
        for (Bundle bundle : getBundles()) {
            bootstrap.addBundle(bundle);
        }
    }

    public abstract Iterable<AbstractModule> getModules(T config, Environment environment);

    public Iterable<ConfiguredBundle<T>> getConfiguredBundles()
    {
        return Arrays.asList(new ViewBundle());
    }

    public Iterable<Bundle> getBundles()
    {
        return Arrays.asList(
                new AssetsBundle("/assets", "/app", "index.html"),
                new FlywayBundle<T>()
                {
                    @Override
                    public DataSourceFactory getDataSourceFactory(T configuration)
                    {
                        return configuration.getDataSourceFactory();
                    }

                    @Override
                    public FlywayFactory getFlywayFactory(T configuration)
                    {
                        return configuration.getFlywayFactory();
                    }
                });
    }

    @Override
    public void run(T config, Environment environment)
            throws Exception
    {
        this.injector = Guice.createInjector(Stage.PRODUCTION, getModules(config, environment));

        System.setProperty(IO_BUFFER_SIZE, String.valueOf(config.getBufferSize().toBytes()));

        environment.healthChecks().register("presto", injector.getInstance(PrestoHealthCheck.class));

        environment.jersey().register(injector.getInstance(ExecuteResource.class));
        environment.jersey().register(injector.getInstance(QueryResource.class));
        environment.jersey().register(injector.getInstance(QueriesResource.class));
        environment.jersey().register(injector.getInstance(UserResource.class));
        environment.jersey().register(injector.getInstance(UsersResource.class));
        environment.jersey().register(injector.getInstance(TablesResource.class));
        environment.jersey().register(injector.getInstance(HealthResource.class));
        environment.jersey().register(injector.getInstance(PingResource.class));
        environment.jersey().register(injector.getInstance(SessionResource.class));
        environment.jersey().register(injector.getInstance(FilesResource.class));
        environment.jersey().register(injector.getInstance(ResultsPreviewResource.class));
        environment.jersey().register(injector.getInstance(S3FilesResource.class));

        environment.jersey().register(injector.getInstance(AirpalUserFactory.class));

        // Setup SSE (Server Sent Events)
        ServletRegistration.Dynamic sseServlet = environment.servlets()
                .addServlet("updates", injector.getInstance(SSEEventSourceServlet.class));
        sseServlet.setAsyncSupported(true);
        sseServlet.addMapping("/api/updates/subscribe");

        // Disable GZIP content encoding for SSE endpoints
        environment.lifecycle().addServerLifecycleListener(server -> {
            for (Handler handler : server.getChildHandlersByClass(BiDiGzipHandler.class)) {
                ((BiDiGzipHandler) handler).addExcludedMimeTypes(SERVER_SENT_EVENTS);
            }
        });
    }
}