Accept query id and slug in QueuedStatement

Author: prithvip

presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java

@@ -227,6 +227,53 @@ public Response postStatement(
         return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled).build();
     }
 
+    /**
+     * HTTP endpoint for submitting queries to the Presto Coordinator.
+     * Presto performs lazy execution. The submission of a query returns
+     * a placeholder for the result set, but the query gets
+     * scheduled/dispatched only when the client polls for results.
+     * This endpoint accepts a pre-minted queryId and slug, instead of
+     * generating it.
+     * @param statement The statement or sql query string submitted
+     * @param queryId Pre-minted query ID to associate with this query
+     * @param slug Pre-minted slug to protect this query
+     * @param xForwardedProto Forwarded protocol (http or https)
+     * @param servletRequest The http request
+     * @param uriInfo {@link javax.ws.rs.core.UriInfo}
+     * @return {@link javax.ws.rs.core.Response} HTTP response code
+     */
+    @POST
+    @Path("/v1/statement/queued/{queryId}")
+    @Produces(APPLICATION_JSON)
+    public Response postStatement(
+            String statement,
+            @PathParam("queryId") QueryId queryId,
+            @QueryParam("slug") String slug,
+            @DefaultValue("false") @QueryParam("binaryResults") boolean binaryResults,
+            @HeaderParam(X_FORWARDED_PROTO) String xForwardedProto,
+            @HeaderParam(PRESTO_PREFIX_URL) String xPrestoPrefixUrl,
+            @Context HttpServletRequest servletRequest,
+            @Context UriInfo uriInfo)
+    {
+        if (isNullOrEmpty(statement)) {
+            throw badRequest(BAD_REQUEST, "SQL statement is empty");
+        }
+
+        abortIfPrefixUrlInvalid(xPrestoPrefixUrl);
+
+        // TODO: For future cases we may want to start tracing from client. Then continuation of tracing
+        //       will be needed instead of creating a new trace here.
+        SessionContext sessionContext = new HttpRequestSessionContext(
+                servletRequest,
+                sqlParserOptions,
+                tracerProviderManager.getTracerProvider(),
+                Optional.of(sessionPropertyManager));
+        Query query = new Query(statement, sessionContext, dispatchManager, queryResultsProvider, 0, queryId, slug);
+        queries.put(query.getQueryId(), query);
+
+        return withCompressionConfiguration(Response.ok(query.getInitialQueryResults(uriInfo, xForwardedProto, xPrestoPrefixUrl, binaryResults)), compressionEnabled).build();
+    }
+
     /**
      * HTTP endpoint for re-processing a failed query
      * @param queryId Query Identifier of the query to be retried
@@ -455,21 +502,34 @@ private static final class Query
         private final DispatchManager dispatchManager;
         private final LocalQueryProvider queryProvider;
         private final QueryId queryId;
-        private final String slug = "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", "");
+        private final String slug;
         private final AtomicLong lastToken = new AtomicLong();
         private final int retryCount;
 
         @GuardedBy("this")
         private ListenableFuture<?> querySubmissionFuture;
 
         public Query(String query, SessionContext sessionContext, DispatchManager dispatchManager, LocalQueryProvider queryResultsProvider, int retryCount)
+        {
+            this(query, sessionContext, dispatchManager, queryResultsProvider, retryCount, dispatchManager.createQueryId(), createSlug());
+        }
+
+        public Query(
+                String query,
+                SessionContext sessionContext,
+                DispatchManager dispatchManager,
+                LocalQueryProvider queryResultsProvider,
+                int retryCount,
+                QueryId queryId,
+                String slug)
         {
             this.query = requireNonNull(query, "query is null");
             this.sessionContext = requireNonNull(sessionContext, "sessionContext is null");
             this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null");
             this.queryProvider = requireNonNull(queryResultsProvider, "queryExecutor is null");
-            this.queryId = dispatchManager.createQueryId();
             this.retryCount = retryCount;
+            this.queryId = requireNonNull(queryId, "queryId is null");
+            this.slug = requireNonNull(slug, "slug is null");
         }
 
         /**
@@ -640,6 +700,11 @@ private QueryResults createQueryResults(long token, UriInfo uriInfo, String xFor
                     dispatchInfo.getWaitingForPrerequisitesTime());
         }
 
+        private static String createSlug()
+        {
+            return "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", "");
+        }
+
         private URI getNextUri(long token, UriInfo uriInfo, String xForwardedProto, String xPrestoPrefixUrl, DispatchInfo dispatchInfo, boolean binaryResults)
         {
             // if failed, query is complete

presto-main/src/test/java/com/facebook/presto/server/TestServer.java

@@ -26,6 +26,7 @@
 import com.facebook.presto.common.Page;
 import com.facebook.presto.common.block.BlockEncodingManager;
 import com.facebook.presto.common.type.TimeZoneNotSupportedException;
+import com.facebook.presto.execution.QueryIdGenerator;
 import com.facebook.presto.execution.buffer.PagesSerdeFactory;
 import com.facebook.presto.server.testing.TestingPrestoServer;
 import com.facebook.presto.spi.QueryId;
@@ -262,6 +263,40 @@ public void testQuery()
         assertEquals(rows, ImmutableList.of(ImmutableList.of("system")));
     }
 
+    @Test
+    public void testQueryWithPreMintedQueryIdAndSlug()
+    {
+        QueryId queryId = new QueryIdGenerator().createNextQueryId();
+        String slug = "xxx";
+        Request request = preparePost()
+                .setUri(uriFor("/v1/statement/queued/", queryId, slug))
+                .setBodyGenerator(createStaticBodyGenerator("show catalogs", UTF_8))
+                .setHeader(PRESTO_USER, "user")
+                .setHeader(PRESTO_SOURCE, "source")
+                .setHeader(PRESTO_CATALOG, "catalog")
+                .setHeader(PRESTO_SCHEMA, "schema")
+                .build();
+
+        QueryResults queryResults = client.execute(request, createJsonResponseHandler(QUERY_RESULTS_CODEC));
+
+        // verify slug in nextUri is same as requested
+        assertEquals(queryResults.getNextUri().getQuery(), "slug=xxx");
+
+        // verify nextUri points to requested query id
+        assertEquals(queryResults.getNextUri().getPath(), format("/v1/statement/queued/%s/1", queryId));
+
+        while (queryResults.getNextUri() != null) {
+            queryResults = client.execute(prepareGet().setUri(queryResults.getNextUri()).build(), createJsonResponseHandler(QUERY_RESULTS_CODEC));
+        }
+
+        if (queryResults.getError() != null) {
+            fail(queryResults.getError().toString());
+        }
+
+        // verify query id was passed down properly
+        assertEquals(server.getDispatchManager().getQueryInfo(queryId).getQueryId(), queryId);
+    }
+
     @Test
     public void testTransactionSupport()
     {
@@ -327,4 +362,13 @@ public URI uriFor(String path)
     {
         return HttpUriBuilder.uriBuilderFrom(server.getBaseUrl()).replacePath(path).build();
     }
+
+    public URI uriFor(String path, QueryId queryId, String slug)
+    {
+        return HttpUriBuilder.uriBuilderFrom(server.getBaseUrl())
+                .replacePath(path)
+                .appendPath(queryId.getId())
+                .addParameter("slug", slug)
+                .build();
+    }
 }