ExecutionClient.java

package com.airbnb.airpal.core.execution;

import com.airbnb.airpal.api.ExecutionRequest;
import com.airbnb.airpal.api.Job;
import com.airbnb.airpal.api.JobState;
import com.airbnb.airpal.api.event.JobFinishedEvent;
import com.airbnb.airpal.api.output.HiveTablePersistentOutput;
import com.airbnb.airpal.api.output.PersistentJobOutputFactory;
import com.airbnb.airpal.api.output.builders.OutputBuilderFactory;
import com.airbnb.airpal.api.output.persistors.PersistorFactory;
import com.airbnb.airpal.core.AirpalUser;
import com.airbnb.airpal.core.store.history.JobHistoryStore;
import com.airbnb.airpal.core.store.jobs.ActiveJobsStore;
import com.airbnb.airpal.core.store.usage.UsageStore;
import com.airbnb.airpal.presto.QueryInfoClient;
import com.airbnb.airpal.presto.Table;
import com.airbnb.airpal.presto.metadata.ColumnCache;
import com.airbnb.airpal.presto.metadata.SchemaCache;
import com.facebook.presto.client.Column;
import com.facebook.presto.client.QueryError;
import com.google.common.eventbus.EventBus;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.inject.Inject;
import lombok.Getter;
import org.joda.time.DateTime;
import org.joda.time.Duration;

import javax.annotation.Nullable;
import javax.validation.constraints.NotNull;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executors;

import static com.airbnb.airpal.presto.QueryRunner.QueryRunnerFactory;

public class ExecutionClient
{
    private final ListeningExecutorService executor = MoreExecutors.listeningDecorator(
            Executors.newCachedThreadPool(SchemaCache.daemonThreadsNamed("execution-client-%d")));

    @Getter
    private final EventBus eventBus;
    @Getter
    private final JobHistoryStore historyStore;
    private final PersistentJobOutputFactory persistentJobOutputFactory;
    @Getter
    private final UsageStore usageStore;
    @Getter
    private final SchemaCache schemaCache;
    private final ColumnCache columnCache;
    private final QueryInfoClient queryInfoClient;
    private final QueryRunnerFactory queryRunnerFactory;
    private final ActiveJobsStore activeJobsStore;
    private final OutputBuilderFactory outputBuilderFactory;
    private final PersistorFactory persistorFactory;
    private final Map<UUID, Execution> executionMap = new ConcurrentHashMap<>();

    @Inject
    public ExecutionClient(QueryRunnerFactory queryRunnerFactory,
            EventBus eventBus,
            JobHistoryStore historyStore,
            PersistentJobOutputFactory persistentJobOutputFactory,
            UsageStore usageStore,
            SchemaCache schemaCache,
            ColumnCache columnCache,
            QueryInfoClient queryInfoClient,
            ActiveJobsStore activeJobsStore,
            OutputBuilderFactory outputBuilderFactory,
            PersistorFactory persistorFactory)
    {
        this.queryRunnerFactory = queryRunnerFactory;
        this.eventBus = eventBus;
        this.historyStore = historyStore;
        this.persistentJobOutputFactory = persistentJobOutputFactory;
        this.usageStore = usageStore;
        this.schemaCache = schemaCache;
        this.columnCache = columnCache;
        this.queryInfoClient = queryInfoClient;
        this.activeJobsStore = activeJobsStore;
        this.outputBuilderFactory = outputBuilderFactory;
        this.persistorFactory = persistorFactory;
    }

    public UUID runQuery(final ExecutionRequest request,
            final AirpalUser user,
            final String schema,
            final Duration timeout)
    {
        return runQuery(request.getQuery(), request.getTmpTable(), user, schema, timeout);
    }

    public UUID runQuery(final String query,
            final String tmpTable,
            final AirpalUser user,
            final String schema,
            final Duration timeout)
    {
        final UUID uuid = UUID.randomUUID();
        final Job job = new Job(user.getUserName(),
                                query,
                                uuid,
                                persistentJobOutputFactory.create(tmpTable, uuid),
                                null,
                                JobState.QUEUED,
                                Collections.<Column>emptyList(),
                                null,
                                null
        );

        final Execution execution = new Execution(job,
                eventBus,
                queryRunnerFactory.create(user.getUserName(), schema),
                queryInfoClient,
                new QueryExecutionAuthorizer(user, "hive", user.getDefaultSchema()),
                timeout,
                columnCache,
                outputBuilderFactory,
                persistorFactory);

        executionMap.put(uuid, execution);
        activeJobsStore.jobStarted(job);

        ListenableFuture<Job> result = executor.submit(execution);
        Futures.addCallback(result, new FutureCallback<Job>()
        {
            @Override
            public void onSuccess(@Nullable Job result)
            {
                if (result != null) {
                    result.setState(JobState.FINISHED);
                }
                jobFinished(result);
            }

            @Override
            public void onFailure(@NotNull Throwable t)
            {
                if (t instanceof ExecutionFailureException) {
                    ExecutionFailureException e = (ExecutionFailureException) t;
                    Job j = e.getJob();
                    j.setState(JobState.FAILED);
                    if (j.getError() == null) {
                        j.setError(new QueryError(e.getMessage(), null, -1, null, null, null, null));
                    }

                    jobFinished(j);
                }
            }
        });

        return uuid;
    }

    protected void jobFinished(Job job)
    {
        job.setQueryFinished(new DateTime());
        activeJobsStore.jobFinished(job);
        historyStore.addRun(job);

        for (Table t : job.getTablesUsed()) {
            usageStore.markUsage(t);
        }

        if (job.getOutput() instanceof HiveTablePersistentOutput && job.getOutput().getLocation() != null) {
            String[] parts = job.getOutput().getLocation().toString().split("\\.");
            if (parts.length == 2) {
                Map<String, List<String>> cache = schemaCache.getSchemaMap("hive");
                List<String> tables = cache.get(parts[0]);
                tables.add(parts[1]);
            }
        }

        eventBus.post(new JobFinishedEvent(job));
        executionMap.remove(job.getUuid());
    }

    public boolean cancelQuery(
            AirpalUser user,
            UUID uuid)
    {
        Execution execution = executionMap.get(uuid);

        if ((execution != null) && (execution.getJob().getUser().equals(user.getUserName()))) {
            execution.cancel();
            return true;
        } else {
            return false;
        }
    }

    public static class ExecutionFailureException extends RuntimeException
    {
        @Getter
        private final Job job;

        public ExecutionFailureException(Job job, String message, Throwable cause)
        {
            super(message, cause);
            this.job = job;
        }
    }
}