Skip to content

Commit

Permalink
Add utility to create Session from http request
Browse files Browse the repository at this point in the history
  • Loading branch information
dain committed Sep 29, 2014
1 parent 67823ca commit 6f02cde
Show file tree
Hide file tree
Showing 4 changed files with 109 additions and 115 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;

/**
Expand Down Expand Up @@ -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();
Expand Down
100 changes: 100 additions & 0 deletions presto-main/src/main/java/com/facebook/presto/server/ResourceUtil.java
Original file line number Diff line number Diff line change
@@ -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:https://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());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -143,76 +133,20 @@ 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);
queries.put(query.getQueryId(), query);
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)
Expand Down

0 comments on commit 6f02cde

Please sign in to comment.