diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java index eaef784ac19c..25d32af3b09f 100644 --- a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java +++ b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchManager.java @@ -30,6 +30,7 @@ import io.prestosql.server.SessionContext; import io.prestosql.server.SessionPropertyDefaults; import io.prestosql.server.SessionSupplier; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.PrestoException; import io.prestosql.spi.QueryId; import io.prestosql.spi.resourcegroups.SelectionContext; @@ -131,7 +132,7 @@ public QueryId createQueryId() return queryIdGenerator.createNextQueryId(); } - public ListenableFuture createQuery(QueryId queryId, String slug, SessionContext sessionContext, String query) + public ListenableFuture createQuery(QueryId queryId, Slug slug, SessionContext sessionContext, String query) { requireNonNull(queryId, "queryId is null"); requireNonNull(sessionContext, "sessionFactory is null"); @@ -155,7 +156,7 @@ public ListenableFuture createQuery(QueryId queryId, String slug, SessionCont * Creates and registers a dispatch query with the query tracker. This method will never fail to register a query with the query * tracker. If an error occurs while creating a dispatch query, a failed dispatch will be created and registered. */ - private void createQueryInternal(QueryId queryId, String slug, SessionContext sessionContext, String query, ResourceGroupManager resourceGroupManager) + private void createQueryInternal(QueryId queryId, Slug slug, SessionContext sessionContext, String query, ResourceGroupManager resourceGroupManager) { Session session = null; PreparedQuery preparedQuery = null; diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchQueryFactory.java b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchQueryFactory.java index 3706da2b0b32..cda42eb08c32 100644 --- a/presto-main/src/main/java/io/prestosql/dispatcher/DispatchQueryFactory.java +++ b/presto-main/src/main/java/io/prestosql/dispatcher/DispatchQueryFactory.java @@ -15,6 +15,7 @@ import io.prestosql.Session; import io.prestosql.execution.QueryPreparer.PreparedQuery; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.resourcegroups.ResourceGroupId; public interface DispatchQueryFactory @@ -23,6 +24,6 @@ DispatchQuery createDispatchQuery( Session session, String query, PreparedQuery preparedQuery, - String slug, + Slug slug, ResourceGroupId resourceGroup); } diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java b/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java index 6173c939b7fa..5ebb0b2c76ec 100644 --- a/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java +++ b/presto-main/src/main/java/io/prestosql/dispatcher/LocalDispatchQueryFactory.java @@ -28,6 +28,7 @@ import io.prestosql.execution.warnings.WarningCollectorFactory; import io.prestosql.metadata.Metadata; import io.prestosql.security.AccessControl; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.PrestoException; import io.prestosql.spi.resourcegroups.ResourceGroupId; import io.prestosql.sql.tree.Statement; @@ -90,7 +91,7 @@ public DispatchQuery createDispatchQuery( Session session, String query, PreparedQuery preparedQuery, - String slug, + Slug slug, ResourceGroupId resourceGroup) { WarningCollector warningCollector = warningCollectorFactory.create(); diff --git a/presto-main/src/main/java/io/prestosql/dispatcher/QueuedStatementResource.java b/presto-main/src/main/java/io/prestosql/dispatcher/QueuedStatementResource.java index c29726cc661e..c8ae122939d0 100644 --- a/presto-main/src/main/java/io/prestosql/dispatcher/QueuedStatementResource.java +++ b/presto-main/src/main/java/io/prestosql/dispatcher/QueuedStatementResource.java @@ -28,6 +28,7 @@ import io.prestosql.execution.QueryState; import io.prestosql.server.HttpRequestSessionContext; import io.prestosql.server.SessionContext; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.ErrorCode; import io.prestosql.spi.QueryId; @@ -70,10 +71,10 @@ import static io.airlift.http.server.AsyncResponseHandler.bindAsyncResponse; import static io.prestosql.execution.QueryState.FAILED; import static io.prestosql.execution.QueryState.QUEUED; +import static io.prestosql.server.protocol.Slug.Context.EXECUTING_QUERY; +import static io.prestosql.server.protocol.Slug.Context.QUEUED_QUERY; import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; -import static java.util.Locale.ENGLISH; import static java.util.Objects.requireNonNull; -import static java.util.UUID.randomUUID; import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; import static java.util.concurrent.TimeUnit.MILLISECONDS; import static java.util.concurrent.TimeUnit.SECONDS; @@ -175,7 +176,7 @@ public void getStatus( @Context UriInfo uriInfo, @Suspended AsyncResponse asyncResponse) { - Query query = getQuery(queryId, slug); + Query query = getQuery(queryId, slug, token); // wait for query to be dispatched, up to the wait timeout ListenableFuture futureStateChange = addTimeout( @@ -206,15 +207,15 @@ public Response cancelQuery( @PathParam("slug") String slug, @PathParam("token") long token) { - getQuery(queryId, slug) + getQuery(queryId, slug, token) .cancel(); return Response.noContent().build(); } - private Query getQuery(QueryId queryId, String slug) + private Query getQuery(QueryId queryId, String slug, long token) { Query query = queries.get(queryId); - if (query == null || !query.getSlug().equals(slug)) { + if (query == null || !query.getSlug().isValid(QUEUED_QUERY, slug, token)) { throw badRequest(NOT_FOUND, "Query not found"); } return query; @@ -229,13 +230,13 @@ private static URI getQueryHtmlUri(QueryId queryId, UriInfo uriInfo, String xFor .build(); } - private static URI getQueuedUri(QueryId queryId, String slug, long token, UriInfo uriInfo, String xForwardedProto) + private static URI getQueuedUri(QueryId queryId, Slug slug, long token, UriInfo uriInfo, String xForwardedProto) { return uriInfo.getBaseUriBuilder() .scheme(getScheme(xForwardedProto, uriInfo)) .replacePath("/v1/statement/queued/") .path(queryId.toString()) - .path(slug) + .path(slug.makeSlug(QUEUED_QUERY, token)) .path(String.valueOf(token)) .replaceQuery("") .build(); @@ -290,7 +291,7 @@ private static final class Query private final SessionContext sessionContext; private final DispatchManager dispatchManager; private final QueryId queryId; - private final String slug = "x" + randomUUID().toString().toLowerCase(ENGLISH).replace("-", ""); + private final Slug slug = Slug.createNew(); private final AtomicLong lastToken = new AtomicLong(); @GuardedBy("this") @@ -309,7 +310,7 @@ public QueryId getQueryId() return queryId; } - public String getSlug() + public Slug getSlug() { return slug; } @@ -412,7 +413,7 @@ private URI getRedirectUri(CoordinatorLocation coordinatorLocation, UriInfo uriI return uriBuilderFrom(coordinatorUri) .appendPath("/v1/statement/executing") .appendPath(queryId.toString()) - .appendPath(slug) + .appendPath(slug.makeSlug(EXECUTING_QUERY, 0)) .appendPath("0") .build(); } diff --git a/presto-main/src/main/java/io/prestosql/execution/DataDefinitionExecution.java b/presto-main/src/main/java/io/prestosql/execution/DataDefinitionExecution.java index bd3dbaa3beb6..f98d7979d96e 100644 --- a/presto-main/src/main/java/io/prestosql/execution/DataDefinitionExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/DataDefinitionExecution.java @@ -26,6 +26,7 @@ import io.prestosql.metadata.Metadata; import io.prestosql.security.AccessControl; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.QueryId; import io.prestosql.sql.planner.Plan; import io.prestosql.sql.tree.Expression; @@ -53,7 +54,7 @@ public class DataDefinitionExecution { private final DataDefinitionTask task; private final T statement; - private final String slug; + private final Slug slug; private final TransactionManager transactionManager; private final Metadata metadata; private final AccessControl accessControl; @@ -63,7 +64,7 @@ public class DataDefinitionExecution private DataDefinitionExecution( DataDefinitionTask task, T statement, - String slug, + Slug slug, TransactionManager transactionManager, Metadata metadata, AccessControl accessControl, @@ -81,7 +82,7 @@ private DataDefinitionExecution( } @Override - public String getSlug() + public Slug getSlug() { return slug; } @@ -300,7 +301,7 @@ public DataDefinitionExecutionFactory( public DataDefinitionExecution createQueryExecution( PreparedQuery preparedQuery, QueryStateMachine stateMachine, - String slug, + Slug slug, WarningCollector warningCollector) { return createDataDefinitionExecution(preparedQuery.getStatement(), preparedQuery.getParameters(), stateMachine, slug); @@ -310,7 +311,7 @@ private DataDefinitionExecution createDataDefinitionExe T statement, List parameters, QueryStateMachine stateMachine, - String slug) + Slug slug) { @SuppressWarnings("unchecked") DataDefinitionTask task = (DataDefinitionTask) tasks.get(statement.getClass()); diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java b/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java index 8a76c5275d31..89fc54cb321f 100644 --- a/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/QueryExecution.java @@ -24,6 +24,7 @@ import io.prestosql.execution.warnings.WarningCollector; import io.prestosql.memory.VersionedMemoryPoolId; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.type.Type; import io.prestosql.sql.planner.Plan; @@ -51,7 +52,7 @@ public interface QueryExecution QueryInfo getQueryInfo(); - String getSlug(); + Slug getSlug(); Duration getTotalCpuTime(); @@ -80,7 +81,7 @@ public interface QueryExecution interface QueryExecutionFactory { - T createQueryExecution(PreparedQuery preparedQuery, QueryStateMachine stateMachine, String slug, WarningCollector warningCollector); + T createQueryExecution(PreparedQuery preparedQuery, QueryStateMachine stateMachine, Slug slug, WarningCollector warningCollector); } /** diff --git a/presto-main/src/main/java/io/prestosql/execution/QueryManager.java b/presto-main/src/main/java/io/prestosql/execution/QueryManager.java index 8bfc167abc47..2d30ca7aa44d 100644 --- a/presto-main/src/main/java/io/prestosql/execution/QueryManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/QueryManager.java @@ -17,6 +17,7 @@ import io.prestosql.Session; import io.prestosql.execution.StateMachine.StateChangeListener; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.QueryId; import java.util.List; @@ -73,7 +74,7 @@ QueryInfo getFullQueryInfo(QueryId queryId) /** * @throws NoSuchElementException if query does not exist */ - boolean isQuerySlugValid(QueryId queryId, String slug); + Slug getQuerySlug(QueryId queryId); /** * @throws NoSuchElementException if query does not exist diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java index 9a57d20ab203..bb486b89d2fd 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryExecution.java @@ -41,6 +41,7 @@ import io.prestosql.operator.ForScheduler; import io.prestosql.security.AccessControl; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.PrestoException; import io.prestosql.spi.QueryId; import io.prestosql.spi.connector.ConnectorTableHandle; @@ -99,7 +100,7 @@ public class SqlQueryExecution private static final OutputBufferId OUTPUT_BUFFER_ID = new OutputBufferId(0); private final QueryStateMachine stateMachine; - private final String slug; + private final Slug slug; private final Metadata metadata; private final SqlParser sqlParser; private final SplitManager splitManager; @@ -126,7 +127,7 @@ public class SqlQueryExecution private SqlQueryExecution( PreparedQuery preparedQuery, QueryStateMachine stateMachine, - String slug, + Slug slug, Metadata metadata, AccessControl accessControl, SqlParser sqlParser, @@ -225,7 +226,7 @@ private Analysis analyze( } @Override - public String getSlug() + public Slug getSlug() { return slug; } @@ -703,7 +704,7 @@ public static class SqlQueryExecutionFactory public QueryExecution createQueryExecution( PreparedQuery preparedQuery, QueryStateMachine stateMachine, - String slug, + Slug slug, WarningCollector warningCollector) { String executionPolicyName = SystemSessionProperties.getExecutionPolicy(stateMachine.getSession()); diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java b/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java index e8e124cbb2b2..f28fa2855da7 100644 --- a/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java +++ b/presto-main/src/main/java/io/prestosql/execution/SqlQueryManager.java @@ -25,6 +25,7 @@ import io.prestosql.execution.StateMachine.StateChangeListener; import io.prestosql.memory.ClusterMemoryManager; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.PrestoException; import io.prestosql.spi.QueryId; import io.prestosql.sql.planner.Plan; @@ -187,9 +188,9 @@ public Session getQuerySession(QueryId queryId) } @Override - public boolean isQuerySlugValid(QueryId queryId, String slug) + public Slug getQuerySlug(QueryId queryId) { - return queryTracker.getQuery(queryId).getSlug().equals(slug); + return queryTracker.getQuery(queryId).getSlug(); } public Plan getQueryPlan(QueryId queryId) diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/ExecutingStatementResource.java b/presto-main/src/main/java/io/prestosql/server/protocol/ExecutingStatementResource.java index ef21a6249b8d..5ff050e6108e 100644 --- a/presto-main/src/main/java/io/prestosql/server/protocol/ExecutingStatementResource.java +++ b/presto-main/src/main/java/io/prestosql/server/protocol/ExecutingStatementResource.java @@ -74,6 +74,7 @@ import static io.prestosql.client.PrestoHeaders.PRESTO_SET_SESSION; import static io.prestosql.client.PrestoHeaders.PRESTO_STARTED_TRANSACTION_ID; import static io.prestosql.memory.context.AggregatedMemoryContext.newSimpleAggregatedMemoryContext; +import static io.prestosql.server.protocol.Slug.Context.EXECUTING_QUERY; import static java.util.Objects.requireNonNull; import static java.util.concurrent.Executors.newSingleThreadScheduledExecutor; import static java.util.concurrent.TimeUnit.MILLISECONDS; @@ -156,7 +157,7 @@ public void getQueryResults( @Context UriInfo uriInfo, @Suspended AsyncResponse asyncResponse) { - Query query = getQuery(queryId, slug); + Query query = getQuery(queryId, slug, token); if (isNullOrEmpty(proto)) { proto = uriInfo.getRequestUri().getScheme(); } @@ -164,11 +165,11 @@ public void getQueryResults( asyncQueryResults(query, token, maxWait, targetResultSize, uriInfo, proto, asyncResponse); } - protected Query getQuery(QueryId queryId, String slug) + protected Query getQuery(QueryId queryId, String slug, long token) { Query query = queries.get(queryId); if (query != null) { - if (!query.isSlugValid(slug)) { + if (!query.isSlugValid(slug, token)) { throw badRequest(NOT_FOUND, "Query not found"); } return query; @@ -176,11 +177,13 @@ protected Query getQuery(QueryId queryId, String slug) // this is the first time the query has been accessed on this coordinator Session session; + Slug querySlug; try { - if (!queryManager.isQuerySlugValid(queryId, slug)) { + session = queryManager.getQuerySession(queryId); + querySlug = queryManager.getQuerySlug(queryId); + if (!querySlug.isValid(EXECUTING_QUERY, slug, token)) { throw badRequest(NOT_FOUND, "Query not found"); } - session = queryManager.getQuerySession(queryId); } catch (NoSuchElementException e) { throw badRequest(NOT_FOUND, "Query not found"); @@ -190,7 +193,7 @@ protected Query getQuery(QueryId queryId, String slug) ExchangeClient exchangeClient = exchangeClientSupplier.get(new SimpleLocalMemoryContext(newSimpleAggregatedMemoryContext(), ExecutingStatementResource.class.getSimpleName())); return Query.create( session, - slug, + querySlug, queryManager, exchangeClient, responseExecutor, @@ -277,7 +280,7 @@ public Response cancelQuery( { Query query = queries.get(queryId); if (query != null) { - if (!query.isSlugValid(slug)) { + if (!query.isSlugValid(slug, token)) { throw badRequest(NOT_FOUND, "Query not found"); } query.cancel(); @@ -286,7 +289,7 @@ public Response cancelQuery( // cancel the query execution directly instead of creating the statement client try { - if (!queryManager.isQuerySlugValid(queryId, slug)) { + if (!queryManager.getQuerySlug(queryId).isValid(EXECUTING_QUERY, slug, token)) { throw badRequest(NOT_FOUND, "Query not found"); } queryManager.cancelQuery(queryId); @@ -298,13 +301,14 @@ public Response cancelQuery( } @DELETE - @Path("/v1/statement/partialCancel/{queryId}/{stage}/{slug}") + @Path("/v1/statement/partialCancel/{queryId}/{stage}/{slug}/{token}") public void partialCancel( @PathParam("queryId") QueryId queryId, @PathParam("stage") int stage, - @PathParam("slug") String slug) + @PathParam("slug") String slug, + @PathParam("token") long token) { - Query query = getQuery(queryId, slug); + Query query = getQuery(queryId, slug, token); query.partialCancel(stage); } diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/Query.java b/presto-main/src/main/java/io/prestosql/server/protocol/Query.java index 718236434719..98e1a378931d 100644 --- a/presto-main/src/main/java/io/prestosql/server/protocol/Query.java +++ b/presto-main/src/main/java/io/prestosql/server/protocol/Query.java @@ -89,6 +89,7 @@ import static io.airlift.concurrent.MoreFutures.addTimeout; import static io.prestosql.SystemSessionProperties.isExchangeCompressionEnabled; import static io.prestosql.execution.QueryState.FAILED; +import static io.prestosql.server.protocol.Slug.Context.EXECUTING_QUERY; import static io.prestosql.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR; import static io.prestosql.util.Failures.toFailure; import static io.prestosql.util.MoreLists.mappedCopy; @@ -103,7 +104,7 @@ class Query private final QueryManager queryManager; private final QueryId queryId; private final Session session; - private final String slug; + private final Slug slug; @GuardedBy("this") private final ExchangeClient exchangeClient; @@ -163,7 +164,7 @@ class Query public static Query create( Session session, - String slug, + Slug slug, QueryManager queryManager, ExchangeClient exchangeClient, Executor dataProcessorExecutor, @@ -186,7 +187,7 @@ public static Query create( private Query( Session session, - String slug, + Slug slug, QueryManager queryManager, ExchangeClient exchangeClient, Executor resultsProcessorExecutor, @@ -240,9 +241,9 @@ public QueryId getQueryId() return queryId; } - public boolean isSlugValid(String slug) + public boolean isSlugValid(String slug, long token) { - return this.slug.equals(slug); + return this.slug.isValid(EXECUTING_QUERY, slug, token); } public QueryInfo getQueryInfo() @@ -455,11 +456,13 @@ private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, Str } URI nextResultsUri = null; + URI partialCancelUri = null; if (nextToken.isPresent()) { nextResultsUri = createNextResultsUri(scheme, uriInfo, nextToken.getAsLong()); + partialCancelUri = findCancelableLeafStage(queryInfo) + .map(stage -> this.createPartialCancelUri(stage, scheme, uriInfo, nextToken.getAsLong())) + .orElse(null); } - Optional partialCancelUri = findCancelableLeafStage(queryInfo) - .map(stage -> this.createPartialCancelUri(stage, scheme, uriInfo)); // update catalog, schema, and path setCatalog = queryInfo.getSetCatalog(); @@ -485,7 +488,7 @@ private synchronized QueryResults getNextResult(long token, UriInfo uriInfo, Str QueryResults queryResults = new QueryResults( queryId.toString(), queryHtmlUri, - partialCancelUri.orElse(null), + partialCancelUri, nextResultsUri, columns, data, @@ -551,20 +554,21 @@ private synchronized URI createNextResultsUri(String scheme, UriInfo uriInfo, lo .scheme(scheme) .replacePath("/v1/statement/executing") .path(queryId.toString()) - .path(slug) + .path(slug.makeSlug(EXECUTING_QUERY, nextToken)) .path(String.valueOf(nextToken)) .replaceQuery("") .build(); } - private URI createPartialCancelUri(int stage, String scheme, UriInfo uriInfo) + private URI createPartialCancelUri(int stage, String scheme, UriInfo uriInfo, long nextToken) { return uriInfo.getBaseUriBuilder() .scheme(scheme) .replacePath("/v1/statement/partialCancel") .path(queryId.toString()) .path(String.valueOf(stage)) - .path(slug) + .path(slug.makeSlug(EXECUTING_QUERY, nextToken)) + .path(String.valueOf(nextToken)) .replaceQuery("") .build(); } diff --git a/presto-main/src/main/java/io/prestosql/server/protocol/Slug.java b/presto-main/src/main/java/io/prestosql/server/protocol/Slug.java new file mode 100644 index 000000000000..c5a70c0ff6f1 --- /dev/null +++ b/presto-main/src/main/java/io/prestosql/server/protocol/Slug.java @@ -0,0 +1,60 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.prestosql.server.protocol; + +import com.google.common.hash.HashFunction; +import com.google.common.hash.Hashing; + +import java.security.SecureRandom; + +import static java.util.Objects.requireNonNull; + +public final class Slug +{ + public enum Context { + QUEUED_QUERY, + EXECUTING_QUERY, + } + + private static final SecureRandom RANDOM = new SecureRandom(); + + public static Slug createNew() + { + byte[] randomBytes = new byte[16]; + RANDOM.nextBytes(randomBytes); + return new Slug(randomBytes); + } + + private final HashFunction hmac; + + private Slug(byte[] slugKey) + { + this.hmac = Hashing.hmacSha1(requireNonNull(slugKey, "slugKey is null")); + } + + public String makeSlug(Context context, long token) + { + // "y" is an arbitrary prefix distinguishing this slug version. Added for troubleshooting purposes. + return "y" + hmac.newHasher() + .putInt(requireNonNull(context, "context is null").ordinal()) + .putLong(token) + .hash() + .toString(); + } + + public boolean isValid(Context context, String slug, long token) + { + return makeSlug(context, token).equals(slug); + } +} diff --git a/presto-tests/src/test/java/io/prestosql/execution/TestQueryRunnerUtil.java b/presto-tests/src/test/java/io/prestosql/execution/TestQueryRunnerUtil.java index e877e6bf90f2..f64b2dc11284 100644 --- a/presto-tests/src/test/java/io/prestosql/execution/TestQueryRunnerUtil.java +++ b/presto-tests/src/test/java/io/prestosql/execution/TestQueryRunnerUtil.java @@ -19,6 +19,7 @@ import io.prestosql.dispatcher.DispatchManager; import io.prestosql.plugin.tpch.TpchPlugin; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.QueryId; import io.prestosql.tests.DistributedQueryRunner; @@ -37,7 +38,7 @@ private TestQueryRunnerUtil() {} public static QueryId createQuery(DistributedQueryRunner queryRunner, Session session, String sql) { DispatchManager dispatchManager = queryRunner.getCoordinator().getDispatchManager(); - getFutureValue(dispatchManager.createQuery(session.getQueryId(), "slug", new TestingSessionContext(session), sql)); + getFutureValue(dispatchManager.createQuery(session.getQueryId(), Slug.createNew(), new TestingSessionContext(session), sql)); return session.getQueryId(); } diff --git a/presto-tests/src/test/java/io/prestosql/tests/TestMetadataManager.java b/presto-tests/src/test/java/io/prestosql/tests/TestMetadataManager.java index 92a1aac0b2af..4b41c06ebcd1 100644 --- a/presto-tests/src/test/java/io/prestosql/tests/TestMetadataManager.java +++ b/presto-tests/src/test/java/io/prestosql/tests/TestMetadataManager.java @@ -20,6 +20,7 @@ import io.prestosql.execution.TestingSessionContext; import io.prestosql.metadata.MetadataManager; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.Plugin; import io.prestosql.spi.QueryId; import io.prestosql.spi.connector.ConnectorFactory; @@ -114,7 +115,7 @@ public void testMetadataIsClearedAfterQueryCanceled() QueryId queryId = dispatchManager.createQueryId(); dispatchManager.createQuery( queryId, - "slug", + Slug.createNew(), new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get(); diff --git a/presto-tests/src/test/java/io/prestosql/tests/TestQueryManager.java b/presto-tests/src/test/java/io/prestosql/tests/TestQueryManager.java index 1dff54dffb61..64b36ad2bf7a 100644 --- a/presto-tests/src/test/java/io/prestosql/tests/TestQueryManager.java +++ b/presto-tests/src/test/java/io/prestosql/tests/TestQueryManager.java @@ -19,6 +19,7 @@ import io.prestosql.execution.QueryState; import io.prestosql.execution.TestingSessionContext; import io.prestosql.server.BasicQueryInfo; +import io.prestosql.server.protocol.Slug; import io.prestosql.spi.PrestoException; import io.prestosql.spi.QueryId; import io.prestosql.tests.tpch.TpchQueryRunnerBuilder; @@ -64,7 +65,7 @@ public void testFailQuery() QueryId queryId = dispatchManager.createQueryId(); dispatchManager.createQuery( queryId, - "slug", + Slug.createNew(), new TestingSessionContext(TEST_SESSION), "SELECT * FROM lineitem") .get();