diff --git a/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java b/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java index 8cb6efea70006..65924f7e25834 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java +++ b/presto-main/src/main/java/com/facebook/presto/server/HttpRequestSessionContext.java @@ -20,6 +20,7 @@ import com.facebook.presto.sql.parser.ParsingException; import com.facebook.presto.sql.parser.ParsingOptions; import com.facebook.presto.sql.parser.SqlParser; +import com.facebook.presto.sql.parser.SqlParserOptions; import com.facebook.presto.transaction.TransactionId; import com.google.common.base.Splitter; import com.google.common.collect.ImmutableMap; @@ -98,7 +99,7 @@ public final class HttpRequestSessionContext private final boolean clientTransactionSupport; private final String clientInfo; - public HttpRequestSessionContext(HttpServletRequest servletRequest) + public HttpRequestSessionContext(HttpServletRequest servletRequest, SqlParserOptions sqlParserOptions) throws WebApplicationException { catalog = trimEmptyToNull(servletRequest.getHeader(PRESTO_CATALOG)); @@ -157,7 +158,7 @@ else if (nameParts.size() == 2) { this.catalogSessionProperties = catalogSessionProperties.entrySet().stream() .collect(toImmutableMap(Entry::getKey, entry -> ImmutableMap.copyOf(entry.getValue()))); - preparedStatements = parsePreparedStatementsHeaders(servletRequest); + preparedStatements = parsePreparedStatementsHeaders(servletRequest, sqlParserOptions); String transactionIdHeader = servletRequest.getHeader(PRESTO_TRANSACTION_ID); clientTransactionSupport = transactionIdHeader != null; @@ -223,7 +224,7 @@ private static void assertRequest(boolean expression, String format, Object... a } } - private static Map parsePreparedStatementsHeaders(HttpServletRequest servletRequest) + private static Map parsePreparedStatementsHeaders(HttpServletRequest servletRequest, SqlParserOptions sqlParserOptions) { ImmutableMap.Builder preparedStatements = ImmutableMap.builder(); for (String header : splitSessionHeader(servletRequest.getHeaders(PRESTO_PREPARED_STATEMENT))) { @@ -241,7 +242,7 @@ private static Map parsePreparedStatementsHeaders(HttpServletReq } // Validate statement - SqlParser sqlParser = new SqlParser(); + SqlParser sqlParser = new SqlParser(sqlParserOptions); try { sqlParser.createStatement(sqlString, new ParsingOptions(AS_DOUBLE /* anything */)); } diff --git a/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java b/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java index 2715d620082cd..cf8d4be544655 100644 --- a/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java +++ b/presto-main/src/main/java/com/facebook/presto/server/protocol/QueuedStatementResource.java @@ -26,6 +26,7 @@ import com.facebook.presto.server.SessionContext; import com.facebook.presto.spi.ErrorCode; import com.facebook.presto.spi.QueryId; +import com.facebook.presto.sql.parser.SqlParserOptions; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.common.collect.Ordering; @@ -110,16 +111,19 @@ public class QueuedStatementResource private final ConcurrentMap queries = new ConcurrentHashMap<>(); private final ScheduledExecutorService queryPurger = newSingleThreadScheduledExecutor(threadsNamed("dispatch-query-purger")); + private final SqlParserOptions sqlParserOptions; + @Inject public QueuedStatementResource( DispatchManager dispatchManager, DispatchExecutor executor, - LocalQueryProvider queryResultsProvider) + LocalQueryProvider queryResultsProvider, + SqlParserOptions sqlParserOptions) { this.dispatchManager = requireNonNull(dispatchManager, "dispatchManager is null"); this.queryResultsProvider = queryResultsProvider; + this.sqlParserOptions = requireNonNull(sqlParserOptions, "sqlParserOptions is null"); - requireNonNull(dispatchManager, "dispatchManager is null"); this.responseExecutor = requireNonNull(executor, "responseExecutor is null").getExecutor(); this.timeoutExecutor = requireNonNull(executor, "timeoutExecutor is null").getScheduledExecutor(); @@ -166,7 +170,7 @@ public Response postStatement( throw badRequest(BAD_REQUEST, "SQL statement is empty"); } - SessionContext sessionContext = new HttpRequestSessionContext(servletRequest); + SessionContext sessionContext = new HttpRequestSessionContext(servletRequest, sqlParserOptions); Query query = new Query(statement, sessionContext, dispatchManager, queryResultsProvider, timeoutExecutor); queries.put(query.getQueryId(), query); diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestHttpRequestSessionContext.java b/presto-main/src/test/java/com/facebook/presto/server/TestHttpRequestSessionContext.java index 58a38cd35da27..719b8cf1cbaa3 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/TestHttpRequestSessionContext.java +++ b/presto-main/src/test/java/com/facebook/presto/server/TestHttpRequestSessionContext.java @@ -15,6 +15,8 @@ import com.facebook.presto.spi.security.Identity; import com.facebook.presto.spi.security.SelectedRole; +import com.facebook.presto.sql.parser.IdentifierSymbol; +import com.facebook.presto.sql.parser.SqlParserOptions; import com.google.common.collect.ImmutableListMultimap; import com.google.common.collect.ImmutableMap; import org.testng.annotations.Test; @@ -24,6 +26,7 @@ import java.io.UnsupportedEncodingException; import java.net.URLEncoder; +import java.util.EnumSet; import java.util.Optional; import static com.facebook.presto.SystemSessionProperties.HASH_PARTITION_COUNT; @@ -67,7 +70,7 @@ public void testSessionContext() .build(), "testRemote"); - HttpRequestSessionContext context = new HttpRequestSessionContext(request); + HttpRequestSessionContext context = new HttpRequestSessionContext(request, new SqlParserOptions()); assertEquals(context.getSource(), "testSource"); assertEquals(context.getCatalog(), "testCatalog"); assertEquals(context.getSchema(), "testSchema"); @@ -99,7 +102,28 @@ public void testPreparedStatementsHeaderDoesNotParse() .put(PRESTO_PREPARED_STATEMENT, "query1=abcdefg") .build(), "testRemote"); - new HttpRequestSessionContext(request); + new HttpRequestSessionContext(request, new SqlParserOptions()); + } + + @Test + public void testPreparedStatementsSpecialCharacters() + { + HttpServletRequest request = new MockHttpServletRequest( + ImmutableListMultimap.builder() + .put(PRESTO_USER, "testUser") + .put(PRESTO_SOURCE, "testSource") + .put(PRESTO_CATALOG, "testCatalog") + .put(PRESTO_SCHEMA, "testSchema") + .put(PRESTO_LANGUAGE, "zh-TW") + .put(PRESTO_TIME_ZONE, "Asia/Taipei") + .put(PRESTO_CLIENT_INFO, "null") + .put(PRESTO_PREPARED_STATEMENT, "query1=select * from tbl:ns") + .build(), + "testRemote"); + SqlParserOptions options = new SqlParserOptions(); + options.allowIdentifierSymbol(EnumSet.allOf(IdentifierSymbol.class)); + + new HttpRequestSessionContext(request, options); } @Test @@ -127,7 +151,7 @@ public void testExtraCredentials() .build(), "testRemote"); - HttpRequestSessionContext context = new HttpRequestSessionContext(request); + HttpRequestSessionContext context = new HttpRequestSessionContext(request, new SqlParserOptions()); assertEquals( context.getIdentity().getExtraCredentials(), ImmutableMap.builder() diff --git a/presto-main/src/test/java/com/facebook/presto/server/TestQuerySessionSupplier.java b/presto-main/src/test/java/com/facebook/presto/server/TestQuerySessionSupplier.java index 1196725f81a58..e80741a9c3fec 100644 --- a/presto-main/src/test/java/com/facebook/presto/server/TestQuerySessionSupplier.java +++ b/presto-main/src/test/java/com/facebook/presto/server/TestQuerySessionSupplier.java @@ -19,6 +19,7 @@ import com.facebook.presto.security.AllowAllAccessControl; import com.facebook.presto.spi.QueryId; import com.facebook.presto.sql.SqlEnvironmentConfig; +import com.facebook.presto.sql.parser.SqlParserOptions; import com.google.common.collect.ImmutableListMultimap; import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; @@ -66,7 +67,7 @@ public class TestQuerySessionSupplier @Test public void testCreateSession() { - HttpRequestSessionContext context = new HttpRequestSessionContext(TEST_REQUEST); + HttpRequestSessionContext context = new HttpRequestSessionContext(TEST_REQUEST, new SqlParserOptions()); QuerySessionSupplier sessionSupplier = new QuerySessionSupplier( createTestTransactionManager(), new AllowAllAccessControl(), @@ -103,7 +104,7 @@ public void testEmptyClientTags() .put(PRESTO_USER, "testUser") .build(), "remoteAddress"); - HttpRequestSessionContext context1 = new HttpRequestSessionContext(request1); + HttpRequestSessionContext context1 = new HttpRequestSessionContext(request1, new SqlParserOptions()); assertEquals(context1.getClientTags(), ImmutableSet.of()); HttpServletRequest request2 = new MockHttpServletRequest( @@ -112,7 +113,7 @@ public void testEmptyClientTags() .put(PRESTO_CLIENT_TAGS, "") .build(), "remoteAddress"); - HttpRequestSessionContext context2 = new HttpRequestSessionContext(request2); + HttpRequestSessionContext context2 = new HttpRequestSessionContext(request2, new SqlParserOptions()); assertEquals(context2.getClientTags(), ImmutableSet.of()); } @@ -125,7 +126,7 @@ public void testInvalidTimeZone() .put(PRESTO_TIME_ZONE, "unknown_timezone") .build(), "testRemote"); - HttpRequestSessionContext context = new HttpRequestSessionContext(request); + HttpRequestSessionContext context = new HttpRequestSessionContext(request, new SqlParserOptions()); QuerySessionSupplier sessionSupplier = new QuerySessionSupplier( createTestTransactionManager(), new AllowAllAccessControl(),