From 34bd7ba69aff3952820e2343b57adef024da8dc2 Mon Sep 17 00:00:00 2001 From: Dong Shi Date: Mon, 9 Nov 2020 10:38:44 -0800 Subject: [PATCH] Pass in SqlParserOptions to HttpRequestSessionContext for Prepared Statements. Parsing for prepared statements has differences from the regular statement handler in that configured parsing options are not performed consistently. This change passes the SqlParsingOptions through to the prepared statement handling code to make that logic consistent. --- .../server/HttpRequestSessionContext.java | 9 +++--- .../protocol/QueuedStatementResource.java | 10 +++++-- .../server/TestHttpRequestSessionContext.java | 30 +++++++++++++++++-- .../server/TestQuerySessionSupplier.java | 9 +++--- 4 files changed, 44 insertions(+), 14 deletions(-) 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(),