trinodb · martint · Jun 5, 2019 · Dec 1, 2018
diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java b/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java
@@ -96,6 +96,7 @@ public DispatchQuery createDispatchQuery(
         WarningCollector warningCollector = warningCollectorFactory.create();
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 query,
+                preparedQuery.getPrepareSql(),
                 session,
                 locationFactory.createQueryLocation(session.getQueryId()),
                 resourceGroup,

diff --git a/presto-main/src/main/java/io/prestosql/event/QueryMonitor.java b/presto-main/src/main/java/io/prestosql/event/QueryMonitor.java
@@ -125,6 +125,7 @@ public void queryCreatedEvent(BasicQueryInfo queryInfo)
                                 queryInfo.getQueryId().toString(),
                                 queryInfo.getSession().getTransactionId().map(TransactionId::toString),
                                 queryInfo.getQuery(),
+                                queryInfo.getPreparedQuery(),
                                 QUEUED.toString(),
                                 queryInfo.getSelf(),
                                 Optional.empty(),
@@ -138,6 +139,7 @@ public void queryImmediateFailureEvent(BasicQueryInfo queryInfo, ExecutionFailur
                         queryInfo.getQueryId().toString(),
                         queryInfo.getSession().getTransactionId().map(TransactionId::toString),
                         queryInfo.getQuery(),
+                        queryInfo.getPreparedQuery(),
                         queryInfo.getState().toString(),
                         queryInfo.getSelf(),
                         Optional.empty(),
@@ -205,6 +207,7 @@ private QueryMetadata createQueryMetadata(QueryInfo queryInfo)
                 queryInfo.getQueryId().toString(),
                 queryInfo.getSession().getTransactionId().map(TransactionId::toString),
                 queryInfo.getQuery(),
+                queryInfo.getPreparedQuery(),
                 queryInfo.getState().toString(),
                 queryInfo.getSelf(),
                 createTextQueryPlan(queryInfo),

diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryInfo.java b/presto-main/src/main/java/io/prestosql/execution/QueryInfo.java
@@ -52,6 +52,7 @@ public class QueryInfo
     private final URI self;
     private final List<String> fieldNames;
     private final String query;
+    private final Optional<String> preparedQuery;
     private final QueryStats queryStats;
     private final Optional<String> setCatalog;
     private final Optional<String> setSchema;
@@ -84,6 +85,7 @@ public QueryInfo(
             @JsonProperty("self") URI self,
             @JsonProperty("fieldNames") List<String> fieldNames,
             @JsonProperty("query") String query,
+            @JsonProperty("preparedQuery") Optional<String> preparedQuery,
             @JsonProperty("queryStats") QueryStats queryStats,
             @JsonProperty("setCatalog") Optional<String> setCatalog,
             @JsonProperty("setSchema") Optional<String> setSchema,
@@ -120,6 +122,7 @@ public QueryInfo(
         requireNonNull(deallocatedPreparedStatements, "deallocatedPreparedStatements is null");
         requireNonNull(startedTransactionId, "startedTransactionId is null");
         requireNonNull(query, "query is null");
+        requireNonNull(preparedQuery, "preparedQuery is null");
         requireNonNull(outputStage, "outputStage is null");
         requireNonNull(inputs, "inputs is null");
         requireNonNull(output, "output is null");
@@ -134,6 +137,7 @@ public QueryInfo(
         this.self = self;
         this.fieldNames = ImmutableList.copyOf(fieldNames);
         this.query = query;
+        this.preparedQuery = preparedQuery;
         this.queryStats = queryStats;
         this.setCatalog = setCatalog;
         this.setSchema = setSchema;
@@ -205,6 +209,12 @@ public String getQuery()
         return query;
     }
 
+    @JsonProperty
+    public Optional<String> getPreparedQuery()
+    {
+        return preparedQuery;
+    }
+
     @JsonProperty
     public QueryStats getQueryStats()
     {

diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryPreparer.java b/presto-main/src/main/java/io/prestosql/execution/QueryPreparer.java
@@ -59,7 +59,13 @@ public PreparedQuery prepareQuery(Session session, String query)
     public PreparedQuery prepareQuery(Session session, Statement wrappedStatement)
             throws ParsingException, PrestoException, SemanticException
     {
-        Statement statement = unwrapExecuteStatement(wrappedStatement, sqlParser, session);
+        Statement statement = wrappedStatement;
+        Optional<String> prepareSql = Optional.empty();
+        if (statement instanceof Execute) {
+            prepareSql = Optional.of(session.getPreparedStatementFromExecute((Execute) statement));
+            statement = sqlParser.createStatement(prepareSql.get(), createParsingOptions(session));
+        }
+
         if (statement instanceof Explain && ((Explain) statement).isAnalyze()) {
             Statement innerStatement = ((Explain) statement).getStatement();
             Optional<QueryType> innerQueryType = StatementUtils.getQueryType(innerStatement.getClass());
@@ -72,17 +78,7 @@ public PreparedQuery prepareQuery(Session session, Statement wrappedStatement)
             parameters = ((Execute) wrappedStatement).getParameters();
         }
         validateParameters(statement, parameters);
-        return new PreparedQuery(statement, parameters);
-    }
-
-    private static Statement unwrapExecuteStatement(Statement statement, SqlParser sqlParser, Session session)
-    {
-        if (!(statement instanceof Execute)) {
-            return statement;
-        }
-
-        String sql = session.getPreparedStatementFromExecute((Execute) statement);
-        return sqlParser.createStatement(sql, createParsingOptions(session));
+        return new PreparedQuery(statement, parameters, prepareSql);
     }
 
     private static void validateParameters(Statement node, List<Expression> parameterValues)
@@ -100,11 +96,13 @@ public static class PreparedQuery
     {
         private final Statement statement;
         private final List<Expression> parameters;
+        private final Optional<String> prepareSql;
 
-        public PreparedQuery(Statement statement, List<Expression> parameters)
+        public PreparedQuery(Statement statement, List<Expression> parameters, Optional<String> prepareSql)
         {
             this.statement = requireNonNull(statement, "statement is null");
             this.parameters = ImmutableList.copyOf(requireNonNull(parameters, "parameters is null"));
+            this.prepareSql = requireNonNull(prepareSql, "prepareSql is null");
         }
 
         public Statement getStatement()
@@ -116,5 +114,10 @@ public List<Expression> getParameters()
         {
             return parameters;
         }
+
+        public Optional<String> getPrepareSql()
+        {
+            return prepareSql;
+        }
     }
 }
diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryStateMachine.java b/presto-main/src/main/java/io/prestosql/execution/QueryStateMachine.java
@@ -96,6 +96,7 @@ public class QueryStateMachine
 
     private final QueryId queryId;
     private final String query;
+    private final Optional<String> preparedQuery;
     private final Session session;
     private final URI self;
     private final ResourceGroupId resourceGroup;
@@ -150,6 +151,7 @@ public class QueryStateMachine
 
     private QueryStateMachine(
             String query,
+            Optional<String> preparedQuery,
             Session session,
             URI self,
             ResourceGroupId resourceGroup,
@@ -160,6 +162,7 @@ private QueryStateMachine(
             WarningCollector warningCollector)
     {
         this.query = requireNonNull(query, "query is null");
+        this.preparedQuery = requireNonNull(preparedQuery, "preparedQuery is null");
         this.session = requireNonNull(session, "session is null");
         this.queryId = session.getQueryId();
         this.self = requireNonNull(self, "self is null");
@@ -179,6 +182,7 @@ private QueryStateMachine(
      */
     public static QueryStateMachine begin(
             String query,
+            Optional<String> preparedQuery,
             Session session,
             URI self,
             ResourceGroupId resourceGroup,
@@ -191,6 +195,7 @@ public static QueryStateMachine begin(
     {
         return beginWithTicker(
                 query,
+                preparedQuery,
                 session,
                 self,
                 resourceGroup,
@@ -205,6 +210,7 @@ public static QueryStateMachine begin(
 
     static QueryStateMachine beginWithTicker(
             String query,
+            Optional<String> preparedQuery,
             Session session,
             URI self,
             ResourceGroupId resourceGroup,
@@ -225,6 +231,7 @@ static QueryStateMachine beginWithTicker(
 
         QueryStateMachine queryStateMachine = new QueryStateMachine(
                 query,
+                preparedQuery,
                 session,
                 self,
                 resourceGroup,
@@ -356,6 +363,7 @@ public BasicQueryInfo getBasicQueryInfo(Optional<BasicStageStats> rootStage)
                 stageStats.isScheduled(),
                 self,
                 query,
+                preparedQuery,
                 queryStats,
                 errorCode == null ? null : errorCode.getType(),
                 errorCode);
@@ -391,6 +399,7 @@ QueryInfo getQueryInfo(Optional<StageInfo> rootStage)
                 self,
                 outputManager.getQueryOutputInfo().map(QueryOutputInfo::getColumnNames).orElse(ImmutableList.of()),
                 query,
+                preparedQuery,
                 getQueryStats(rootStage),
                 Optional.ofNullable(setCatalog.get()),
                 Optional.ofNullable(setSchema.get()),
@@ -985,6 +994,7 @@ public void pruneQueryInfo()
                 queryInfo.getSelf(),
                 queryInfo.getFieldNames(),
                 queryInfo.getQuery(),
+                queryInfo.getPreparedQuery(),
                 pruneQueryStats(queryInfo.getQueryStats()),
                 queryInfo.getSetCatalog(),
                 queryInfo.getSetSchema(),

diff --git a/presto-main/src/main/java/io/prestosql/server/BasicQueryInfo.java b/presto-main/src/main/java/io/prestosql/server/BasicQueryInfo.java
@@ -52,6 +52,7 @@ public class BasicQueryInfo
     private final boolean scheduled;
     private final URI self;
     private final String query;
+    private final Optional<String> preparedQuery;
     private final BasicQueryStats queryStats;
     private final ErrorType errorType;
     private final ErrorCode errorCode;
@@ -66,6 +67,7 @@ public BasicQueryInfo(
             @JsonProperty("scheduled") boolean scheduled,
             @JsonProperty("self") URI self,
             @JsonProperty("query") String query,
+            @JsonProperty("preparedQuery") Optional<String> preparedQuery,
             @JsonProperty("queryStats") BasicQueryStats queryStats,
             @JsonProperty("errorType") ErrorType errorType,
             @JsonProperty("errorCode") ErrorCode errorCode)
@@ -80,6 +82,7 @@ public BasicQueryInfo(
         this.scheduled = scheduled;
         this.self = requireNonNull(self, "self is null");
         this.query = requireNonNull(query, "query is null");
+        this.preparedQuery = requireNonNull(preparedQuery, "preparedQuery is null");
         this.queryStats = requireNonNull(queryStats, "queryStats is null");
     }
 
@@ -93,6 +96,7 @@ public BasicQueryInfo(QueryInfo queryInfo)
                 queryInfo.isScheduled(),
                 queryInfo.getSelf(),
                 queryInfo.getQuery(),
+                queryInfo.getPreparedQuery(),
                 new BasicQueryStats(queryInfo.getQueryStats()),
                 queryInfo.getErrorType(),
                 queryInfo.getErrorCode());
@@ -109,6 +113,7 @@ public static BasicQueryInfo immediateFailureQueryInfo(Session session, String q
                 false,
                 self,
                 query,
+                Optional.empty(),
                 immediateFailureQueryStats(),
                 errorCode == null ? null : errorCode.getType(),
                 errorCode);
@@ -162,6 +167,12 @@ public String getQuery()
         return query;
     }
 
+    @JsonProperty
+    public Optional<String> getPreparedQuery()
+    {
+        return preparedQuery;
+    }
+
     @JsonProperty
     public BasicQueryStats getQueryStats()
     {

diff --git a/presto-main/src/main/java/io/prestosql/server/QueryResource.java b/presto-main/src/main/java/io/prestosql/server/QueryResource.java
@@ -230,6 +230,7 @@ private static QueryInfo toFullQueryInfo(DispatchQuery query)
                 info.getSelf(),
                 ImmutableList.of(),
                 info.getQuery(),
+                info.getPreparedQuery(),
                 queryStats,
                 Optional.empty(),
                 Optional.empty(),

diff --git a/presto-main/src/test/java/io/prestosql/execution/MockManagedQueryExecution.java b/presto-main/src/test/java/io/prestosql/execution/MockManagedQueryExecution.java
@@ -106,6 +106,7 @@ public BasicQueryInfo getBasicQueryInfo()
                 !state.isDone(),
                 URI.create("http://test"),
                 "SELECT 1",
+                Optional.empty(),
                 new BasicQueryStats(
                         new DateTime(1),
                         new DateTime(2),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestCommitTask.java b/presto-main/src/test/java/io/prestosql/execution/TestCommitTask.java
@@ -125,6 +125,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session,
     {
         return QueryStateMachine.begin(
                 query,
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestDeallocateTask.java b/presto-main/src/test/java/io/prestosql/execution/TestDeallocateTask.java
@@ -29,6 +29,7 @@
 import org.testng.annotations.Test;
 
 import java.net.URI;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 
@@ -83,6 +84,7 @@ private Set<String> executeDeallocate(String statementName, String sqlString, Se
         AccessControl accessControl = new AccessControlManager(transactionManager);
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 sqlString,
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestPrepareTask.java b/presto-main/src/test/java/io/prestosql/execution/TestPrepareTask.java
@@ -34,6 +34,7 @@
 
 import java.net.URI;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 
 import static io.airlift.concurrent.Threads.daemonThreadsNamed;
@@ -104,6 +105,7 @@ private Map<String, String> executePrepare(String statementName, Statement state
         TransactionManager transactionManager = createTestTransactionManager();
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 sqlString,
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestQueryStateMachine.java b/presto-main/src/test/java/io/prestosql/execution/TestQueryStateMachine.java
@@ -494,6 +494,7 @@ private QueryStateMachine createQueryStateMachineWithTicker(Ticker ticker)
         AccessControl accessControl = new AccessControlManager(transactionManager);
         QueryStateMachine stateMachine = QueryStateMachine.beginWithTicker(
                 QUERY,
+                Optional.empty(),
                 TEST_SESSION,
                 LOCATION,
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestResetSessionTask.java b/presto-main/src/test/java/io/prestosql/execution/TestResetSessionTask.java
@@ -38,6 +38,7 @@
 import org.testng.annotations.Test;
 
 import java.net.URI;
+import java.util.Optional;
 import java.util.Set;
 import java.util.concurrent.ExecutorService;
 
@@ -107,6 +108,7 @@ public void test()
 
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 "reset foo",
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestRollbackTask.java b/presto-main/src/test/java/io/prestosql/execution/TestRollbackTask.java
@@ -117,6 +117,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session,
     {
         return QueryStateMachine.begin(
                 query,
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSetPathTask.java b/presto-main/src/test/java/io/prestosql/execution/TestSetPathTask.java
@@ -107,6 +107,7 @@ private QueryStateMachine createQueryStateMachine(String query)
     {
         return QueryStateMachine.begin(
                 query,
+                Optional.empty(),
                 TEST_SESSION,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSetRoleTask.java b/presto-main/src/test/java/io/prestosql/execution/TestSetRoleTask.java
@@ -109,6 +109,7 @@ private void assertSetRole(String statement, Map<String, SelectedRole> expected)
         SetRole setRole = (SetRole) parser.createStatement(statement);
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 statement,
+                Optional.empty(),
                 testSessionBuilder()
                         .setCatalog(CATALOG_NAME)
                         .build(),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestSetSessionTask.java b/presto-main/src/test/java/io/prestosql/execution/TestSetSessionTask.java
@@ -47,6 +47,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 
 import static io.airlift.concurrent.MoreFutures.getFutureValue;
@@ -182,6 +183,7 @@ private void testSetSessionWithParameters(String property, Expression expression
         QualifiedName qualifiedPropName = QualifiedName.of(CATALOG_NAME, property);
         QueryStateMachine stateMachine = QueryStateMachine.begin(
                 format("set %s = 'old_value'", qualifiedPropName),
+                Optional.empty(),
                 TEST_SESSION,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),

diff --git a/presto-main/src/test/java/io/prestosql/execution/TestStartTransactionTask.java b/presto-main/src/test/java/io/prestosql/execution/TestStartTransactionTask.java
@@ -260,6 +260,7 @@ private QueryStateMachine createQueryStateMachine(String query, Session session,
     {
         return QueryStateMachine.begin(
                 query,
+                Optional.empty(),
                 session,
                 URI.create("fake://uri"),
                 new ResourceGroupId("test"),