Compress large prepared statements headers

client/trino-client/src/main/java/io/trino/client/StatementClientV1.java

@@ -64,7 +64,7 @@ class StatementClientV1
     private static final MediaType MEDIA_TYPE_TEXT = MediaType.parse("text/plain; charset=utf-8");
     private static final JsonCodec<QueryResults> QUERY_RESULTS_CODEC = jsonCodec(QueryResults.class);
 
-    private static final Splitter SESSION_HEADER_SPLITTER = Splitter.on('=').limit(2).trimResults();
+    private static final Splitter COLLECTION_HEADER_SPLITTER = Splitter.on('=').limit(2).trimResults();
     private static final String USER_AGENT_VALUE = StatementClientV1.class.getSimpleName() +
             "/" +
             firstNonNull(StatementClientV1.class.getPackage().getImplementationVersion(), "unknown");
@@ -395,7 +395,7 @@ private void processResponse(Headers headers, QueryResults results)
         setPath.set(headers.get(TRINO_HEADERS.responseSetPath()));
 
         for (String setSession : headers.values(TRINO_HEADERS.responseSetSession())) {
-            List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(setSession);
+            List<String> keyValue = COLLECTION_HEADER_SPLITTER.splitToList(setSession);
             if (keyValue.size() != 2) {
                 continue;
             }
@@ -404,15 +404,15 @@ private void processResponse(Headers headers, QueryResults results)
         resetSessionProperties.addAll(headers.values(TRINO_HEADERS.responseClearSession()));
 
         for (String setRole : headers.values(TRINO_HEADERS.responseSetRole())) {
-            List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(setRole);
+            List<String> keyValue = COLLECTION_HEADER_SPLITTER.splitToList(setRole);
             if (keyValue.size() != 2) {
                 continue;
             }
             setRoles.put(keyValue.get(0), ClientSelectedRole.valueOf(urlDecode(keyValue.get(1))));
         }
 
         for (String entry : headers.values(TRINO_HEADERS.responseAddedPrepare())) {
-            List<String> keyValue = SESSION_HEADER_SPLITTER.splitToList(entry);
+            List<String> keyValue = COLLECTION_HEADER_SPLITTER.splitToList(entry);
             if (keyValue.size() != 2) {
                 continue;
             }

client/trino-jdbc/src/test/java/io/trino/jdbc/TestJdbcPreparedStatement.java

@@ -14,6 +14,8 @@
 package io.trino.jdbc;
 
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.common.collect.Iterables;
 import io.airlift.log.Logging;
 import io.trino.client.ClientTypeSignature;
 import io.trino.client.ClientTypeSignatureParameter;
@@ -49,6 +51,7 @@
 import java.time.ZoneOffset;
 import java.util.Calendar;
 import java.util.TimeZone;
+import java.util.stream.LongStream;
 
 import static com.google.common.base.Strings.repeat;
 import static com.google.common.base.Verify.verify;
@@ -72,14 +75,21 @@
 
 public class TestJdbcPreparedStatement
 {
+    private static final int HEADER_SIZE_LIMIT = 16 * 1024;
+
     private TestingTrinoServer server;
 
     @BeforeClass
     public void setup()
             throws Exception
     {
         Logging.initialize();
-        server = TestingTrinoServer.create();
+        server = TestingTrinoServer.builder()
+                .setProperties(ImmutableMap.<String, String>builder()
+                        .put("http-server.max-request-header-size", format("%sB", HEADER_SIZE_LIMIT))
+                        .put("http-server.max-response-header-size", format("%sB", HEADER_SIZE_LIMIT))
+                        .buildOrThrow())
+                .build();
         server.installPlugin(new BlackHolePlugin());
         server.installPlugin(new MemoryPlugin());
         server.createCatalog("blackhole", "blackhole");
@@ -208,9 +218,9 @@ public void testGetParameterMetaData()
 
             try (PreparedStatement statement = connection.prepareStatement(
                     "SELECT ? FROM test_get_parameterMetaData WHERE c_boolean = ? AND c_decimal = ? " +
-                        "AND c_decimal_2 = ? AND c_varchar = ? AND c_varchar_2 = ? AND c_row = ? " +
-                        "AND c_array = ? AND c_map = ? AND c_tinyint = ? AND c_integer = ? AND c_bigint = ? " +
-                        "AND c_smallint = ? AND c_real = ? AND c_double = ?")) {
+                            "AND c_decimal_2 = ? AND c_varchar = ? AND c_varchar_2 = ? AND c_row = ? " +
+                            "AND c_array = ? AND c_map = ? AND c_tinyint = ? AND c_integer = ? AND c_bigint = ? " +
+                            "AND c_smallint = ? AND c_real = ? AND c_double = ?")) {
                 ParameterMetaData parameterMetaData = statement.getParameterMetaData();
                 assertEquals(parameterMetaData.getParameterCount(), 15);
 
@@ -387,6 +397,23 @@ public void testDeallocate()
         }
     }
 
+    @Test
+    public void testLargePreparedStatement()
+            throws Exception
+    {
+        int elements = HEADER_SIZE_LIMIT + 1;
+        try (Connection connection = createConnection();
+                PreparedStatement statement = connection.prepareStatement("VALUES ?" + repeat(", ?", elements - 1))) {
+            for (int i = 0; i < elements; i++) {
+                statement.setLong(i + 1, i);
+            }
+            try (ResultSet resultSet = statement.executeQuery()) {
+                assertThat(readRows(resultSet).stream().map(Iterables::getOnlyElement))
+                        .containsExactlyInAnyOrder(LongStream.range(0, elements).boxed().toArray());
+            }
+        }
+    }
+
     @Test
     public void testExecuteUpdate()
             throws Exception

core/trino-main/src/main/java/io/trino/server/HttpRequestSessionContextFactory.java

@@ -24,6 +24,7 @@
 import io.trino.client.ProtocolHeaders;
 import io.trino.metadata.Metadata;
 import io.trino.security.AccessControl;
+import io.trino.server.protocol.PreparedStatementEncoder;
 import io.trino.spi.security.AccessDeniedException;
 import io.trino.spi.security.GroupProvider;
 import io.trino.spi.security.Identity;
@@ -72,13 +73,15 @@ public class HttpRequestSessionContextFactory
     private static final Splitter DOT_SPLITTER = Splitter.on('.');
     public static final String AUTHENTICATED_IDENTITY = "trino.authenticated-identity";
 
+    private final PreparedStatementEncoder preparedStatementEncoder;
     private final Metadata metadata;
     private final GroupProvider groupProvider;
     private final AccessControl accessControl;
 
     @Inject
-    public HttpRequestSessionContextFactory(Metadata metadata, GroupProvider groupProvider, AccessControl accessControl)
+    public HttpRequestSessionContextFactory(PreparedStatementEncoder preparedStatementEncoder, Metadata metadata, GroupProvider groupProvider, AccessControl accessControl)
     {
+        this.preparedStatementEncoder = requireNonNull(preparedStatementEncoder, "preparedStatementEncoder is null");
         this.metadata = requireNonNull(metadata, "metadata is null");
         this.groupProvider = requireNonNull(groupProvider, "groupProvider is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
@@ -378,17 +381,18 @@ private static void assertRequest(boolean expression, String format, Object... a
         }
     }
 
-    private static Map<String, String> parsePreparedStatementsHeaders(ProtocolHeaders protocolHeaders, MultivaluedMap<String, String> headers)
+    private Map<String, String> parsePreparedStatementsHeaders(ProtocolHeaders protocolHeaders, MultivaluedMap<String, String> headers)
     {
         ImmutableMap.Builder<String, String> preparedStatements = ImmutableMap.builder();
-        parseProperty(headers, protocolHeaders.requestPreparedStatement()).forEach((key, sqlString) -> {
+        parseProperty(headers, protocolHeaders.requestPreparedStatement()).forEach((key, value) -> {
             String statementName;
             try {
                 statementName = urlDecode(key);
             }
             catch (IllegalArgumentException e) {
                 throw badRequest(format("Invalid %s header: %s", protocolHeaders.requestPreparedStatement(), e.getMessage()));
             }
+            String sqlString = preparedStatementEncoder.decodePreparedStatementFromHeader(value);
 
             // Validate statement
             SqlParser sqlParser = new SqlParser();

core/trino-main/src/main/java/io/trino/server/ProtocolConfig.java

@@ -16,13 +16,16 @@
 import io.airlift.configuration.Config;
 import io.airlift.configuration.ConfigDescription;
 
+import javax.annotation.Nonnull;
 import javax.validation.constraints.Pattern;
 
 import java.util.Optional;
 
 public class ProtocolConfig
 {
     private String alternateHeaderName;
+    private int preparedStatementCompressionThreshold = 2 * 1024;
+    private int preparedStatementCompressionMinimalGain = 512;
 
     @Deprecated
     public Optional<@Pattern(regexp = "[A-Za-z]+") String> getAlternateHeaderName()
@@ -38,4 +41,32 @@ public ProtocolConfig setAlternateHeaderName(String alternateHeaderName)
         this.alternateHeaderName = alternateHeaderName;
         return this;
     }
+
+    @Nonnull
+    public int getPreparedStatementCompressionThreshold()
+    {
+        return preparedStatementCompressionThreshold;
+    }
+
+    @Config("protocol.v1.prepared-statement-compression.length-threshold")
+    @ConfigDescription("Prepared statements longer than the configured value will be compressed")
+    public ProtocolConfig setPreparedStatementCompressionThreshold(int preparedStatementCompressionThreshold)
+    {
+        this.preparedStatementCompressionThreshold = preparedStatementCompressionThreshold;
+        return this;
+    }
+
+    @Nonnull
+    public int getPreparedStatementCompressionMinimalGain()
+    {
+        return preparedStatementCompressionMinimalGain;
+    }
+
+    @Config("protocol.v1.prepared-statement-compression.min-gain")
+    @ConfigDescription("Prepared statement compression will not be applied if size gain is less than the configured value")
+    public ProtocolConfig setPreparedStatementCompressionMinimalGain(int preparedStatementCompressionMinimalGain)
+    {
+        this.preparedStatementCompressionMinimalGain = preparedStatementCompressionMinimalGain;
+        return this;
+    }
 }

core/trino-main/src/main/java/io/trino/server/ServerMainModule.java

@@ -101,6 +101,7 @@
 import io.trino.server.PluginManager.PluginsProvider;
 import io.trino.server.SliceSerialization.SliceDeserializer;
 import io.trino.server.SliceSerialization.SliceSerializer;
+import io.trino.server.protocol.PreparedStatementEncoder;
 import io.trino.server.remotetask.HttpLocationFactory;
 import io.trino.spi.PageIndexerFactory;
 import io.trino.spi.PageSorter;
@@ -207,6 +208,7 @@ protected void setup(Binder binder)
             httpServerConfig.setAdminEnabled(false);
         });
 
+        binder.bind(PreparedStatementEncoder.class).in(Scopes.SINGLETON);
         binder.bind(HttpRequestSessionContextFactory.class).in(Scopes.SINGLETON);
         install(new InternalCommunicationModule());
 

core/trino-main/src/main/java/io/trino/server/protocol/ExecutingStatementResource.java

@@ -88,6 +88,7 @@ public class ExecutingStatementResource
 
     private final ConcurrentMap<QueryId, Query> queries = new ConcurrentHashMap<>();
     private final ScheduledExecutorService queryPurger = newSingleThreadScheduledExecutor(threadsNamed("execution-query-purger"));
+    private final PreparedStatementEncoder preparedStatementEncoder;
     private final boolean compressionEnabled;
 
     @Inject
@@ -98,6 +99,7 @@ public ExecutingStatementResource(
             QueryInfoUrlFactory queryInfoUrlTemplate,
             @ForStatementResource BoundedExecutor responseExecutor,
             @ForStatementResource ScheduledExecutorService timeoutExecutor,
+            PreparedStatementEncoder preparedStatementEncoder,
             ServerConfig serverConfig)
     {
         this.queryManager = requireNonNull(queryManager, "queryManager is null");
@@ -106,6 +108,7 @@ public ExecutingStatementResource(
         this.queryInfoUrlFactory = requireNonNull(queryInfoUrlTemplate, "queryInfoUrlTemplate is null");
         this.responseExecutor = requireNonNull(responseExecutor, "responseExecutor is null");
         this.timeoutExecutor = requireNonNull(timeoutExecutor, "timeoutExecutor is null");
+        this.preparedStatementEncoder = requireNonNull(preparedStatementEncoder, "preparedStatementEncoder is null");
         this.compressionEnabled = requireNonNull(serverConfig, "serverConfig is null").isQueryResultsCompressionEnabled();
 
         queryPurger.scheduleWithFixedDelay(
@@ -210,12 +213,12 @@ private void asyncQueryResults(
         }
         ListenableFuture<QueryResults> queryResultsFuture = query.waitForResults(token, uriInfo, wait, targetResultSize);
 
-        ListenableFuture<Response> response = Futures.transform(queryResultsFuture, queryResults -> toResponse(query, queryResults, compressionEnabled), directExecutor());
+        ListenableFuture<Response> response = Futures.transform(queryResultsFuture, queryResults -> toResponse(query, queryResults), directExecutor());
 
         bindAsyncResponse(asyncResponse, response, responseExecutor);
     }
 
-    private static Response toResponse(Query query, QueryResults queryResults, boolean compressionEnabled)
+    private Response toResponse(Query query, QueryResults queryResults)
     {
         ResponseBuilder response = Response.ok(queryResults);
 
@@ -239,7 +242,7 @@ private static Response toResponse(Query query, QueryResults queryResults, boole
         // add added prepare statements
         for (Entry<String, String> entry : query.getAddedPreparedStatements().entrySet()) {
             String encodedKey = urlEncode(entry.getKey());
-            String encodedValue = urlEncode(entry.getValue());
+            String encodedValue = urlEncode(preparedStatementEncoder.encodePreparedStatementForHeader(entry.getValue()));
             response.header(protocolHeaders.responseAddedPrepare(), encodedKey + '=' + encodedValue);
         }
 

core/trino-main/src/main/java/io/trino/server/protocol/PreparedStatementEncoder.java

@@ -0,0 +1,74 @@
+/*
+ * 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.trino.server.protocol;
+
+import io.airlift.compress.zstd.ZstdCompressor;
+import io.airlift.compress.zstd.ZstdDecompressor;
+import io.trino.server.ProtocolConfig;
+
+import javax.inject.Inject;
+
+import static com.google.common.io.BaseEncoding.base64Url;
+import static java.lang.Math.toIntExact;
+import static java.nio.charset.StandardCharsets.UTF_8;
+import static java.util.Objects.requireNonNull;
+
+public class PreparedStatementEncoder
+{
+    // No valid SQL statement starts with $
+    private static final String PREFIX = "$zstd:";
+
+    private final int compressionThreshold;
+    private final int compressionMinGain;
+
+    @Inject
+    public PreparedStatementEncoder(ProtocolConfig protocolConfig)
+    {
+        requireNonNull(protocolConfig, "protocolConfig is null");
+        this.compressionThreshold = protocolConfig.getPreparedStatementCompressionThreshold();
+        this.compressionMinGain = protocolConfig.getPreparedStatementCompressionMinimalGain();
+    }
+
+    public String encodePreparedStatementForHeader(String preparedStatement)
+    {
+        if (preparedStatement.length() < compressionThreshold) {
+            return preparedStatement;
+        }
+
+        ZstdCompressor compressor = new ZstdCompressor();
+        byte[] inputBytes = preparedStatement.getBytes(UTF_8);
+        byte[] compressed = new byte[compressor.maxCompressedLength(inputBytes.length)];
+        int outputSize = compressor.compress(inputBytes, 0, inputBytes.length, compressed, 0, compressed.length);
+        String encoded = base64Url().encode(compressed, 0, outputSize);
+
+        if (encoded.length() + PREFIX.length() + compressionMinGain > preparedStatement.length()) {
+            return preparedStatement;
+        }
+        return PREFIX + encoded;
+    }
+
+    public String decodePreparedStatementFromHeader(String headerValue)
+    {
+        if (!headerValue.startsWith(PREFIX)) {
+            return headerValue;
+        }
+
+        String encoded = headerValue.substring(PREFIX.length());
+        byte[] compressed = base64Url().decode(encoded);
+
+        byte[] preparedStatement = new byte[toIntExact(ZstdDecompressor.getDecompressedSize(compressed, 0, compressed.length))];
+        new ZstdDecompressor().decompress(compressed, 0, compressed.length, preparedStatement, 0, preparedStatement.length);
+        return new String(preparedStatement, UTF_8);
+    }
+}

core/trino-main/src/test/java/io/trino/server/TestHttpRequestSessionContextFactory.java

@@ -19,6 +19,7 @@
 import io.airlift.jaxrs.testing.GuavaMultivaluedMap;
 import io.trino.client.ProtocolHeaders;
 import io.trino.security.AllowAllAccessControl;
+import io.trino.server.protocol.PreparedStatementEncoder;
 import io.trino.spi.security.Identity;
 import io.trino.spi.security.SelectedRole;
 import org.testng.annotations.Test;
@@ -40,7 +41,11 @@
 
 public class TestHttpRequestSessionContextFactory
 {
-    private static final HttpRequestSessionContextFactory SESSION_CONTEXT_FACTORY = new HttpRequestSessionContextFactory(createTestMetadataManager(), ImmutableSet::of, new AllowAllAccessControl());
+    private static final HttpRequestSessionContextFactory SESSION_CONTEXT_FACTORY = new HttpRequestSessionContextFactory(
+            new PreparedStatementEncoder(new ProtocolConfig()),
+            createTestMetadataManager(),
+            ImmutableSet::of,
+            new AllowAllAccessControl());
 
     @Test
     public void testSessionContext()

core/trino-main/src/test/java/io/trino/server/TestProtocolConfig.java

@@ -28,18 +28,24 @@ public class TestProtocolConfig
     public void testDefaults()
     {
         assertRecordedDefaults(recordDefaults(ProtocolConfig.class)
-                .setAlternateHeaderName(null));
+                .setAlternateHeaderName(null)
+                .setPreparedStatementCompressionThreshold(2 * 1024)
+                .setPreparedStatementCompressionMinimalGain(512));
     }
 
     @Test
     public void testExplicitPropertyMappings()
     {
         Map<String, String> properties = new ImmutableMap.Builder<String, String>()
                 .put("protocol.v1.alternate-header-name", "taco")
+                .put("protocol.v1.prepared-statement-compression.length-threshold", "412")
+                .put("protocol.v1.prepared-statement-compression.min-gain", "0")
                 .buildOrThrow();
 
         ProtocolConfig expected = new ProtocolConfig()
-                .setAlternateHeaderName("taco");
+                .setAlternateHeaderName("taco")
+                .setPreparedStatementCompressionThreshold(412)
+                .setPreparedStatementCompressionMinimalGain(0);
 
         assertFullMapping(properties, expected);
     }