package com.facebook.presto;

import com.facebook.presto.execution.ExchangePlanFragmentSource;
import com.facebook.presto.execution.QueryManagerConfig;
import com.facebook.presto.metadata.Metadata;
import com.facebook.presto.operator.Operator;
import com.facebook.presto.operator.OperatorStats;
import com.facebook.presto.operator.SourceHashProviderFactory;
import com.facebook.presto.operator.HackPlanFragmentSourceProvider;
import com.facebook.presto.split.DataStreamProvider;
import com.facebook.presto.sql.analyzer.AnalysisResult;
import com.facebook.presto.sql.analyzer.Analyzer;
import com.facebook.presto.sql.analyzer.Session;
import com.facebook.presto.sql.parser.SqlParser;
import com.facebook.presto.sql.planner.DistributedLogicalPlanner;
import com.facebook.presto.sql.planner.LocalExecutionPlanner;
import com.facebook.presto.sql.planner.LogicalPlanner;
import com.facebook.presto.sql.planner.PlanNodeIdAllocator;
import com.facebook.presto.sql.planner.PlanPrinter;
import com.facebook.presto.sql.planner.SubPlan;
import com.facebook.presto.sql.planner.TableScanPlanFragmentSource;
import com.facebook.presto.sql.planner.plan.PlanNode;
import com.facebook.presto.sql.planner.plan.PlanNodeId;
import com.facebook.presto.sql.planner.plan.TableScanNode;
import com.facebook.presto.sql.tree.Statement;
import com.facebook.presto.tpch.TpchSplit;
import com.facebook.presto.tpch.TpchTableHandle;
import com.facebook.presto.util.MaterializedResult;
import com.google.common.collect.ImmutableMap;
import io.airlift.units.DataSize;
import org.intellij.lang.annotations.Language;

import static com.facebook.presto.util.MaterializedResult.materialize;
import static io.airlift.units.DataSize.Unit.MEGABYTE;
import static org.testng.Assert.assertTrue;

public class TestLocalQueries
    extends AbstractTestQueries
{
    private String catalog;
    private String schema;
    private DataStreamProvider dataStreamProvider;
    private Metadata metadata;

    @Override
    protected void setUpQueryFramework(String catalog, String schema, DataStreamProvider dataStreamProvider, Metadata metadata)
            throws Exception
    {
        this.catalog = catalog;
        this.schema = schema;
        this.dataStreamProvider = dataStreamProvider;
        this.metadata = metadata;
    }

    @Override
    protected MaterializedResult computeActual(@Language("SQL") String sql)
    {
        Statement statement = SqlParser.createStatement(sql);

        Session session = new Session(null, catalog, schema);

        Analyzer analyzer = new Analyzer(session, metadata);

        AnalysisResult analysis = analyzer.analyze(statement);

        PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();
        PlanNode plan = new LogicalPlanner(session, metadata, idAllocator).plan(analysis);
        new PlanPrinter().print(plan, analysis.getTypes());

        SubPlan subplan = new DistributedLogicalPlanner(metadata, idAllocator).createSubplans(plan, analysis.getSymbolAllocator(), true);
        assertTrue(subplan.getChildren().isEmpty(), "Expected subplan to have no children");

        ImmutableMap.Builder<PlanNodeId, TableScanPlanFragmentSource> builder = ImmutableMap.builder();
        for (PlanNode source : subplan.getFragment().getSources()) {
            TableScanNode tableScan = (TableScanNode) source;
            TpchTableHandle handle = (TpchTableHandle) tableScan.getTable();

            builder.put(tableScan.getId(), new TableScanPlanFragmentSource(new TpchSplit(handle)));
        }

        DataSize maxOperatorMemoryUsage = new DataSize(50, MEGABYTE);
        LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(session, metadata,
                new HackPlanFragmentSourceProvider(dataStreamProvider, null, new QueryManagerConfig()),
                analysis.getTypes(),
                null,
                builder.build(),
                ImmutableMap.<PlanNodeId, ExchangePlanFragmentSource>of(),
                new OperatorStats(),
                new SourceHashProviderFactory(maxOperatorMemoryUsage),
                maxOperatorMemoryUsage
        );

        Operator operator = executionPlanner.plan(plan);

        return materialize(operator);
    }
}

@@ -54,1701 +54,275 @@
 sto.
-execution.ExchangePlanFragmentSource;%0Aimport com.facebook.presto.execution.QueryManagerConfig;%0Aimport com.facebook.presto.metadata.Metadata;%0Aimport com.facebook.presto.operator.Operator;%0Aimport com.facebook.presto.operator.OperatorStats;%0Aimport com.facebook.presto.operator.SourceHashProviderFactory;%0Aimport com.facebook.presto.operator.HackPlanFragmentSourceProvider;%0Aimport com.facebook.presto.split.DataStreamProvider;%0Aimport com.facebook.presto.sql.analyzer.AnalysisResult;%0Aimport com.facebook.presto.sql.analyzer.Analyzer;%0Aimport com.facebook.presto.sql.analyzer.Session;%0Aimport com.facebook.presto.sql.parser.SqlParser;%0Aimport com.facebook.presto.sql.planner.DistributedLogicalPlanner;%0Aimport com.facebook.presto.sql.planner.LocalExecutionPlanner;%0Aimport com.facebook.presto.sql.planner.LogicalPlanner;%0Aimport com.facebook.presto.sql.planner.PlanNodeIdAllocator;%0Aimport com.facebook.presto.sql.planner.PlanPrinter;%0Aimport com.facebook.presto.sql.planner.SubPlan;%0Aimport com.facebook.presto.sql.planner.TableScanPlanFragmentSource;%0Aimport com.facebook.presto.sql.planner.plan.PlanNode;%0Aimport com.facebook.presto.sql.planner.plan.PlanNodeId;%0Aimport com.facebook.presto.sql.planner.plan.TableScanNode;%0Aimport com.facebook.presto.sql.tree.Statement;%0Aimport com.facebook.presto.tpch.TpchSplit;%0Aimport com.facebook.presto.tpch.TpchTableHandle;%0Aimport com.facebook.presto.util.MaterializedResult;%0Aimport com.google.common.collect.ImmutableMap;%0Aimport io.airlift.units.DataSize;%0Aimport org.intellij.lang.annotations.Language;%0A%0Aimport static com.facebook.presto.util.MaterializedResult.materialize;%0Aimport static io.airlift.units.DataSize.Unit.MEGABYTE;%0Aimport static org.testng.Assert.assertTru
+metadata.Metadata;%0Aimport com.facebook.presto.split.DataStreamProvider;%0Aimport com.facebook.presto.sql.analyzer.Session;%0Aimport com.facebook.presto.util.LocalQueryRunner;%0Aimport com.facebook.presto.util.MaterializedResult;%0Aimport org.intellij.lang.annotations.Languag
 e;%0A%0A
@@ -966,341 +966,93 @@
    S
-tatement statement = SqlParser.createStatement(sql);%0A%0A        Session session = new Session(null, catalog, schema);%0A%0A        Analyzer analyzer = new Analyzer(session, metadata);%0A%0A        AnalysisResult analysis = analyzer.analyze(statement);%0A%0A        PlanNodeIdAllocator idAllocator = new PlanNodeIdAllocator();%0A        PlanNode plan
+ession session = new Session(null, catalog, schema);%0A%0A        LocalQueryRunner runner
  = n
@@ -1060,1473 +1060,97 @@
 w Lo
-gi
 cal
-Planner(session, metadata, idAllocator).plan(analysis);%0A        new PlanPrinter().print(plan, analysis.getTypes());%0A%0A        SubPlan subplan = new DistributedLogicalPlanner(metadata, idAllocator).createSubplans(plan, analysis.getSymbolAllocator(), true);%0A        assertTrue(subplan.getChildren().isEmpty(), %22Expected subplan to have no children%22);%0A%0A        ImmutableMap.Builder%3CPlanNodeId, TableScanPlanFragmentSource%3E builder = ImmutableMap.builder();%0A        for (PlanNode source : subplan.getFragment().getSources()) %7B%0A            TableScanNode tableScan = (TableScanNode) source;%0A            TpchTableHandle handle = (TpchTableHandle) tableScan.getTable();%0A%0A            builder.put(tableScan.getId(), new TableScanPlanFragmentSource(new TpchSplit(handle)));%0A        %7D%0A%0A        DataSize maxOperatorMemoryUsage = new DataSize(50, MEGABYTE);%0A        LocalExecutionPlanner executionPlanner = new LocalExecutionPlanner(session, metadata,%0A                new HackPlanFragmentSourceProvider(dataStreamProvider, null, new QueryManagerConfig()),%0A                analysis.getTypes(),%0A                null,%0A                builder.build(),%0A                ImmutableMap.%3CPlanNodeId, ExchangePlanFragmentSource%3Eof(),%0A                new OperatorStats(),%0A                new SourceHashProviderFactory(maxOperatorMemoryUsage),%0A                maxOperatorMemoryUsage%0A        );%0A%0A        Operator operator = executionPlanner.plan(plan);%0A%0A        return materialize(operator
+QueryRunner(dataStreamProvider, metadata, session);%0A%0A        return runner.execute(sql
 );%0A 

Make TestLocalQueries use LocalQueryRunner