Compress large prepared statements in HTTP headers

This allows working with larger prepared statements before hitting
header size limits (on either server, or intermediate proxy).

Since the client should not make assumption on any specific format of
prepared statements in the headers, the compression is applied
unconditionally.

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");
@@ -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

@@ -100,6 +100,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;
@@ -206,6 +207,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);
     }

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

@@ -25,6 +25,7 @@
 import io.trino.metadata.Metadata;
 import io.trino.metadata.SessionPropertyManager;
 import io.trino.security.AllowAllAccessControl;
+import io.trino.server.protocol.PreparedStatementEncoder;
 import io.trino.spi.QueryId;
 import io.trino.spi.TrinoException;
 import io.trino.sql.SqlEnvironmentConfig;
@@ -67,7 +68,11 @@ public class TestQuerySessionSupplier
             .put(TRINO_HEADERS.requestSession(), JOIN_DISTRIBUTION_TYPE + "=partitioned," + HASH_PARTITION_COUNT + " = 43")
             .put(TRINO_HEADERS.requestPreparedStatement(), "query1=select * from foo,query2=select * from bar")
             .build());
-    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 testCreateSession()

core/trino-main/src/test/java/io/trino/server/security/TestResourceSecurity.java

@@ -32,6 +32,8 @@
 import io.trino.security.AccessControl;
 import io.trino.security.AccessControlManager;
 import io.trino.server.HttpRequestSessionContextFactory;
+import io.trino.server.ProtocolConfig;
+import io.trino.server.protocol.PreparedStatementEncoder;
 import io.trino.server.security.oauth2.OAuth2Client;
 import io.trino.server.testing.TestingTrinoServer;
 import io.trino.spi.security.AccessDeniedException;
@@ -949,7 +951,11 @@ public static class TestResource
         @Inject
         public TestResource(AccessControl accessControl)
         {
-            this.sessionContextFactory = new HttpRequestSessionContextFactory(createTestMetadataManager(), user -> ImmutableSet.of(), accessControl);
+            this.sessionContextFactory = new HttpRequestSessionContextFactory(
+                    new PreparedStatementEncoder(new ProtocolConfig()),
+                    createTestMetadataManager(),
+                    user -> ImmutableSet.of(),
+                    accessControl);
         }
 
         @ResourceSecurity(AUTHENTICATED_USER)