From 6f02cde48b569f9520f16757bdb5d176a65bdb2b Mon Sep 17 00:00:00 2001 From: Dain Sundstrom Date: Sun, 28 Sep 2014 12:29:10 -0700 Subject: [PATCH] Add utility to create Session from http request --- .../presto/server/ExecuteResource.java | 2 +- .../facebook/presto/server/QueryResource.java | 48 +-------- .../facebook/presto/server/ResourceUtil.java | 100 ++++++++++++++++++ .../presto/server/StatementResource.java | 74 +------------ 4 files changed, 109 insertions(+), 115 deletions(-) create mode 100644 presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java diff --git a/presto-main/src/main/java/com/facebook/presto/server/ExecuteResource.java b/presto-main/src/main/java/com/facebook/presto/server/ExecuteResource.java index bc5b7f4b1453..fffaea0279fc 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/ExecuteResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/ExecuteResource.java @@ -47,7 +47,7 @@ import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE; import static com.facebook.presto.client.PrestoHeaders.PRESTO_TIME_ZONE; import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER; -import static com.facebook.presto.server.StatementResource.assertRequest; +import static com.facebook.presto.server.ResourceUtil.assertRequest; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Preconditions.checkState; import static com.google.common.base.Strings.isNullOrEmpty; diff --git a/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java b/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java index a70ba707d09c..731e708c4760 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/QueryResource.java @@ -24,7 +24,6 @@ import javax.servlet.http.HttpServletRequest; import javax.ws.rs.DELETE; import javax.ws.rs.GET; -import javax.ws.rs.HeaderParam; import javax.ws.rs.POST; import javax.ws.rs.Path; import javax.ws.rs.PathParam; @@ -37,21 +36,12 @@ import java.net.URI; import java.util.List; -import java.util.Locale; import java.util.NoSuchElementException; -import java.util.TimeZone; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_CATALOG; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_LANGUAGE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_SCHEMA; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_TIME_ZONE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER; -import static com.facebook.presto.server.StatementResource.assertRequest; -import static com.facebook.presto.server.StatementResource.getTimeZoneKey; +import static com.facebook.presto.server.ResourceUtil.assertRequest; +import static com.facebook.presto.server.ResourceUtil.createSessionForRequest; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Strings.isNullOrEmpty; -import static com.google.common.net.HttpHeaders.USER_AGENT; import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom; /** @@ -102,42 +92,12 @@ public Response getQueryInfo(@PathParam("queryId") QueryId queryId) @Produces(MediaType.APPLICATION_JSON) public Response createQuery( String statement, - @HeaderParam(PRESTO_USER) String user, - @HeaderParam(PRESTO_SOURCE) String source, - @HeaderParam(PRESTO_CATALOG) String catalog, - @HeaderParam(PRESTO_SCHEMA) String schema, - @HeaderParam(PRESTO_TIME_ZONE) String timeZoneId, - @HeaderParam(PRESTO_LANGUAGE) String language, - @HeaderParam(USER_AGENT) String userAgent, - @Context HttpServletRequest requestContext, + @Context HttpServletRequest servletRequest, @Context UriInfo uriInfo) { assertRequest(!isNullOrEmpty(statement), "SQL statement is empty"); - assertRequest(!isNullOrEmpty(user), "User (%s) is empty", PRESTO_USER); - assertRequest(!isNullOrEmpty(catalog), "Catalog (%s) is empty", PRESTO_CATALOG); - assertRequest(!isNullOrEmpty(schema), "Schema (%s) is empty", PRESTO_SCHEMA); - if (timeZoneId == null) { - timeZoneId = TimeZone.getDefault().getID(); - } - - Locale locale = Locale.getDefault(); - if (language != null) { - locale = Locale.forLanguageTag(language); - } - - String remoteUserAddress = requestContext.getRemoteAddr(); - - Session session = Session.builder() - .setUser(user) - .setSource(source) - .setCatalog(catalog) - .setSchema(schema) - .setTimeZoneKey(getTimeZoneKey(timeZoneId)) - .setLocale(locale) - .setRemoteUserAddress(remoteUserAddress) - .setUserAgent(userAgent) - .build(); + Session session = createSessionForRequest(servletRequest); QueryInfo queryInfo = queryManager.createQuery(session, statement); URI pagesUri = uriBuilderFrom(uriInfo.getRequestUri()).appendPath(queryInfo.getQueryId().toString()).build(); diff --git a/presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java b/presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java new file mode 100644 index 000000000000..b8e72f63aecb --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java @@ -0,0 +1,100 @@ +/* + * 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 com.facebook.presto.server; + +import com.facebook.presto.Session; +import com.facebook.presto.Session.SessionBuilder; +import com.facebook.presto.spi.type.TimeZoneKey; +import com.facebook.presto.spi.type.TimeZoneNotSupportedException; + +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.WebApplicationException; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import javax.ws.rs.core.Response.Status; + +import java.util.Locale; + +import static com.facebook.presto.client.PrestoHeaders.PRESTO_CATALOG; +import static com.facebook.presto.client.PrestoHeaders.PRESTO_LANGUAGE; +import static com.facebook.presto.client.PrestoHeaders.PRESTO_SCHEMA; +import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE; +import static com.facebook.presto.client.PrestoHeaders.PRESTO_TIME_ZONE; +import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER; +import static com.google.common.base.Strings.isNullOrEmpty; +import static com.google.common.net.HttpHeaders.USER_AGENT; +import static java.lang.String.format; + +final class ResourceUtil +{ + private ResourceUtil() + { + } + + public static Session createSessionForRequest(HttpServletRequest servletRequest) + { + SessionBuilder sessionBuilder = Session.builder() + .setUser(getRequiredHeader(servletRequest, PRESTO_USER, "User")) + .setSource(servletRequest.getHeader(PRESTO_SOURCE)) + .setCatalog(getRequiredHeader(servletRequest, PRESTO_CATALOG, "Catalog")) + .setSchema(getRequiredHeader(servletRequest, PRESTO_SCHEMA, "Schema")) + .setRemoteUserAddress(servletRequest.getRemoteAddr()) + .setUserAgent(servletRequest.getHeader(USER_AGENT)); + + String timeZoneId = servletRequest.getHeader(PRESTO_TIME_ZONE); + if (timeZoneId != null) { + sessionBuilder.setTimeZoneKey(getTimeZoneKey(timeZoneId)); + } + + String language = servletRequest.getHeader(PRESTO_LANGUAGE); + if (language != null) { + sessionBuilder.setLocale(Locale.forLanguageTag(language)); + } + + return sessionBuilder.build(); + } + + private static String getRequiredHeader(HttpServletRequest servletRequest, String name, String description) + { + String user = servletRequest.getHeader(name); + assertRequest(!isNullOrEmpty(user), description + " (%s) is empty", PRESTO_USER); + return user; + } + + public static void assertRequest(boolean expression, String format, Object... args) + { + if (!expression) { + throw badRequest(format(format, args)); + } + } + + private static TimeZoneKey getTimeZoneKey(String timeZoneId) + { + try { + return TimeZoneKey.getTimeZoneKey(timeZoneId); + } + catch (TimeZoneNotSupportedException e) { + throw badRequest(e.getMessage()); + } + } + + private static WebApplicationException badRequest(String message) + { + throw new WebApplicationException(Response + .status(Status.BAD_REQUEST) + .type(MediaType.TEXT_PLAIN) + .entity(message) + .build()); + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/server/StatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/StatementResource.java index 3d222881ff84..23b4460065c7 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/StatementResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/StatementResource.java @@ -34,8 +34,6 @@ import com.facebook.presto.spi.ConnectorSession; import com.facebook.presto.spi.Page; import com.facebook.presto.spi.block.Block; -import com.facebook.presto.spi.type.TimeZoneKey; -import com.facebook.presto.spi.type.TimeZoneNotSupportedException; import com.facebook.presto.spi.type.Type; import com.facebook.presto.util.IterableTransformer; import com.google.common.base.Preconditions; @@ -59,7 +57,6 @@ import javax.servlet.http.HttpServletRequest; import javax.ws.rs.DELETE; import javax.ws.rs.GET; -import javax.ws.rs.HeaderParam; import javax.ws.rs.POST; import javax.ws.rs.Path; import javax.ws.rs.PathParam; @@ -79,30 +76,23 @@ import java.util.HashSet; import java.util.Iterator; import java.util.List; -import java.util.Locale; import java.util.Set; -import java.util.TimeZone; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.ConcurrentMap; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicLong; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_CATALOG; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_LANGUAGE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_SCHEMA; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_SOURCE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_TIME_ZONE; -import static com.facebook.presto.client.PrestoHeaders.PRESTO_USER; import static com.facebook.presto.execution.QueryInfo.queryIdGetter; import static com.facebook.presto.execution.StageInfo.getAllStages; import static com.facebook.presto.execution.StageInfo.stageStateGetter; +import static com.facebook.presto.server.ResourceUtil.assertRequest; +import static com.facebook.presto.server.ResourceUtil.createSessionForRequest; import static com.facebook.presto.util.Failures.toFailure; import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkNotNull; import static com.google.common.base.Strings.isNullOrEmpty; import static com.google.common.collect.Iterables.transform; -import static com.google.common.net.HttpHeaders.USER_AGENT; import static io.airlift.concurrent.Threads.threadsNamed; import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom; import static io.airlift.units.DataSize.Unit.MEGABYTE; @@ -143,43 +133,13 @@ public void stop() @Produces(MediaType.APPLICATION_JSON) public Response createQuery( String statement, - @HeaderParam(PRESTO_USER) String user, - @HeaderParam(PRESTO_SOURCE) String source, - @HeaderParam(PRESTO_CATALOG) String catalog, - @HeaderParam(PRESTO_SCHEMA) String schema, - @HeaderParam(PRESTO_TIME_ZONE) String timeZoneId, - @HeaderParam(PRESTO_LANGUAGE) String language, - @HeaderParam(USER_AGENT) String userAgent, - @Context HttpServletRequest requestContext, + @Context HttpServletRequest servletRequest, @Context UriInfo uriInfo) throws InterruptedException { assertRequest(!isNullOrEmpty(statement), "SQL statement is empty"); - assertRequest(!isNullOrEmpty(user), "User (%s) is empty", PRESTO_USER); - assertRequest(!isNullOrEmpty(catalog), "Catalog (%s) is empty", PRESTO_CATALOG); - assertRequest(!isNullOrEmpty(schema), "Schema (%s) is empty", PRESTO_SCHEMA); - if (timeZoneId == null) { - timeZoneId = TimeZone.getDefault().getID(); - } - - Locale locale = Locale.getDefault(); - if (language != null) { - locale = Locale.forLanguageTag(language); - } - - String remoteUserAddress = requestContext.getRemoteAddr(); - - Session session = Session.builder() - .setUser(user) - .setSource(source) - .setCatalog(catalog) - .setSchema(schema) - .setTimeZoneKey(getTimeZoneKey(timeZoneId)) - .setLocale(locale) - .setRemoteUserAddress(remoteUserAddress) - .setUserAgent(userAgent) - .build(); + Session session = createSessionForRequest(servletRequest); ExchangeClient exchangeClient = exchangeClientSupplier.get(); Query query = new Query(session, statement, queryManager, exchangeClient); @@ -187,32 +147,6 @@ public Response createQuery( return Response.ok(query.getNextResults(uriInfo, new Duration(1, TimeUnit.MILLISECONDS))).build(); } - static void assertRequest(boolean expression, String format, Object... args) - { - if (!expression) { - throw badRequest(format(format, args)); - } - } - - static TimeZoneKey getTimeZoneKey(String timeZoneId) - { - try { - return TimeZoneKey.getTimeZoneKey(timeZoneId); - } - catch (TimeZoneNotSupportedException e) { - throw badRequest(e.getMessage()); - } - } - - private static WebApplicationException badRequest(String message) - { - throw new WebApplicationException(Response - .status(Status.BAD_REQUEST) - .type(MediaType.TEXT_PLAIN) - .entity(message) - .build()); - } - @GET @Path("{queryId}/{token}") @Produces(MediaType.APPLICATION_JSON)