Use PreparedQuery for QueryExecutionFactory createQueryExecution

ralpher · Nov 4, 2018 · b77fa0b · b77fa0b
1 parent f8578fb
commit b77fa0b
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 27 deletions.
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/DataDefinitionExecution.java
@@ -14,6 +14,7 @@
 package com.facebook.presto.execution;
 
 import com.facebook.presto.Session;
+import com.facebook.presto.execution.QueryPreparer.PreparedQuery;
 import com.facebook.presto.execution.StateMachine.StateChangeListener;
 import com.facebook.presto.execution.warnings.WarningCollector;
 import com.facebook.presto.memory.VersionedMemoryPoolId;
@@ -38,7 +39,6 @@
 import javax.annotation.Nullable;
 import javax.inject.Inject;
 
-import java.net.URI;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
@@ -307,20 +307,29 @@ public DataDefinitionExecutionFactory(
         public DataDefinitionExecution<?> createQueryExecution(
                 String query,
                 Session session,
-                Statement statement,
-                List<Expression> parameters,
+                PreparedQuery preparedQuery,
                 ResourceGroupId resourceGroup,
                 WarningCollector warningCollector)
         {
-            URI self = locationFactory.createQueryLocation(session.getQueryId());
+            return createDataDefinitionExecution(query, session, resourceGroup, preparedQuery.getStatement(), preparedQuery.getParameters(), warningCollector);
+        }
 
-            DataDefinitionTask<Statement> task = getTask(statement);
+        private <T extends Statement> DataDefinitionExecution<T> createDataDefinitionExecution(
+                String query,
+                Session session,
+                ResourceGroupId resourceGroup,
+                T statement,
+                List<Expression> parameters,
+                WarningCollector warningCollector)
+        {
+            @SuppressWarnings("unchecked")
+            DataDefinitionTask<T> task = (DataDefinitionTask<T>) tasks.get(statement.getClass());
             checkArgument(task != null, "no task for statement: %s", statement.getClass().getSimpleName());
 
             QueryStateMachine stateMachine = QueryStateMachine.begin(
                     query,
                     session,
-                    self,
+                    locationFactory.createQueryLocation(session.getQueryId()),
                     resourceGroup,
                     task.isTransactionControl(),
                     transactionManager,
@@ -331,11 +340,5 @@ public DataDefinitionExecution<?> createQueryExecution(
             stateMachine.setUpdateType(task.getName());
             return new DataDefinitionExecution<>(task, statement, transactionManager, metadata, accessControl, stateMachine, parameters);
         }
-
-        @SuppressWarnings("unchecked")
-        private <T extends Statement> DataDefinitionTask<T> getTask(T statement)
-        {
-            return (DataDefinitionTask<T>) tasks.get(statement.getClass());
-        }
     }
 }
diff --git a/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/QueryExecution.java
@@ -14,15 +14,14 @@
 package com.facebook.presto.execution;
 
 import com.facebook.presto.Session;
+import com.facebook.presto.execution.QueryPreparer.PreparedQuery;
 import com.facebook.presto.execution.QueryTracker.TrackedQuery;
 import com.facebook.presto.execution.StateMachine.StateChangeListener;
 import com.facebook.presto.execution.warnings.WarningCollector;
 import com.facebook.presto.memory.VersionedMemoryPoolId;
 import com.facebook.presto.spi.resourceGroups.ResourceGroupId;
 import com.facebook.presto.spi.type.Type;
 import com.facebook.presto.sql.planner.Plan;
-import com.facebook.presto.sql.tree.Expression;
-import com.facebook.presto.sql.tree.Statement;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -61,7 +60,7 @@ public interface QueryExecution
 
     interface QueryExecutionFactory<T extends QueryExecution>
     {
-        T createQueryExecution(String query, Session session, Statement statement, List<Expression> parameters, ResourceGroupId resourceGroup, WarningCollector warningCollector);
+        T createQueryExecution(String query, Session session, PreparedQuery preparedQuery, ResourceGroupId resourceGroup, WarningCollector warningCollector);
     }
 
     /**

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryExecution.java
@@ -20,6 +20,7 @@
 import com.facebook.presto.connector.ConnectorId;
 import com.facebook.presto.cost.CostCalculator;
 import com.facebook.presto.cost.StatsCalculator;
+import com.facebook.presto.execution.QueryPreparer.PreparedQuery;
 import com.facebook.presto.execution.StateMachine.StateChangeListener;
 import com.facebook.presto.execution.scheduler.ExecutionPolicy;
 import com.facebook.presto.execution.scheduler.NodeScheduler;
@@ -59,8 +60,6 @@
 import com.facebook.presto.sql.planner.SubPlan;
 import com.facebook.presto.sql.planner.optimizations.PlanOptimizer;
 import com.facebook.presto.sql.tree.Explain;
-import com.facebook.presto.sql.tree.Expression;
-import com.facebook.presto.sql.tree.Statement;
 import com.facebook.presto.transaction.TransactionManager;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.util.concurrent.ListenableFuture;
@@ -131,7 +130,7 @@ private SqlQueryExecution(
             Session session,
             URI self,
             ResourceGroupId resourceGroup,
-            Statement statement,
+            PreparedQuery preparedQuery,
             TransactionManager transactionManager,
             Metadata metadata,
             AccessControl accessControl,
@@ -150,7 +149,6 @@ private SqlQueryExecution(
             NodeTaskMap nodeTaskMap,
             QueryExplainer queryExplainer,
             ExecutionPolicy executionPolicy,
-            List<Expression> parameters,
             SplitSchedulerStats schedulerStats,
             StatsCalculator statsCalculator,
             CostCalculator costCalculator,
@@ -193,8 +191,16 @@ private SqlQueryExecution(
                     warningCollector);
 
             // analyze query
-            Analyzer analyzer = new Analyzer(stateMachine.getSession(), metadata, sqlParser, accessControl, Optional.of(queryExplainer), parameters, warningCollector);
-            this.analysis = analyzer.analyze(statement);
+            requireNonNull(preparedQuery, "preparedQuery is null");
+            Analyzer analyzer = new Analyzer(
+                    stateMachine.getSession(),
+                    metadata,
+                    sqlParser,
+                    accessControl,
+                    Optional.of(queryExplainer),
+                    preparedQuery.getParameters(),
+                    warningCollector);
+            this.analysis = analyzer.analyze(preparedQuery.getStatement());
 
             stateMachine.setUpdateType(analysis.getUpdateType());
 
@@ -732,8 +738,7 @@ public static class SqlQueryExecutionFactory
         public QueryExecution createQueryExecution(
                 String query,
                 Session session,
-                Statement statement,
-                List<Expression> parameters,
+                PreparedQuery preparedQuery,
                 ResourceGroupId resourceGroup,
                 WarningCollector warningCollector)
         {
@@ -746,7 +751,7 @@ public QueryExecution createQueryExecution(
                     session,
                     locationFactory.createQueryLocation(session.getQueryId()),
                     resourceGroup,
-                    statement,
+                    preparedQuery,
                     transactionManager,
                     metadata,
                     accessControl,
@@ -765,7 +770,6 @@ public QueryExecution createQueryExecution(
                     nodeTaskMap,
                     queryExplainer,
                     executionPolicy,
-                    parameters,
                     schedulerStats,
                     statsCalculator,
                     costCalculator,

diff --git a/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java b/presto-main/src/main/java/com/facebook/presto/execution/SqlQueryManager.java
@@ -381,8 +381,7 @@ private <C> void createQueryInternal(QueryId queryId, SessionContext sessionCont
             queryExecution = queryExecutionFactory.createQueryExecution(
                     query,
                     session,
-                    preparedQuery.getStatement(),
-                    preparedQuery.getParameters(),
+                    preparedQuery,
                     selectionContext.getResourceGroupId(),
                     warningCollectorFactory.create());