From a85e79a9355bf1657d53922e9c0dfd90207ffa5f Mon Sep 17 00:00:00 2001
From: Star Poon <ypoon@lycorp.co.jp>
Date: Wed, 5 Feb 2025 17:08:55 +0900
Subject: [PATCH] Update aircompressor to v3

---
 gateway-ha/pom.xml                                 | 14 ++++++++++++--
 .../gateway/ha/router/TrinoQueryProperties.java    |  7 ++++---
 2 files changed, 16 insertions(+), 5 deletions(-)

diff --git a/gateway-ha/pom.xml b/gateway-ha/pom.xml
index 09a22109d..93c921fcf 100644
--- a/gateway-ha/pom.xml
+++ b/gateway-ha/pom.xml
@@ -109,8 +109,8 @@
 
         <dependency>
             <groupId>io.airlift</groupId>
-            <artifactId>aircompressor</artifactId>
-            <version>0.27</version>
+            <artifactId>aircompressor-v3</artifactId>
+            <version>3.2</version>
         </dependency>
 
         <dependency>
@@ -497,6 +497,16 @@
                 </configuration>
             </plugin>
 
+            <plugin>
+                <groupId>org.basepom.maven</groupId>
+                <artifactId>duplicate-finder-maven-plugin</artifactId>
+                <configuration>
+                    <ignoredResourcePatterns combine.children="append">
+                        <ignoredResourcePattern>aircompressor.*</ignoredResourcePattern>
+                    </ignoredResourcePatterns>
+                </configuration>
+            </plugin>
+
             <plugin>
                 <groupId>org.apache.maven.plugins</groupId>
                 <artifactId>maven-resources-plugin</artifactId>
diff --git a/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java b/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java
index da119cbd7..52260dbe8 100644
--- a/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java
+++ b/gateway-ha/src/main/java/io/trino/gateway/ha/router/TrinoQueryProperties.java
@@ -23,7 +23,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
-import io.airlift.compress.zstd.ZstdDecompressor;
+import io.airlift.compress.v3.zstd.ZstdDecompressor;
 import io.airlift.json.JsonCodec;
 import io.airlift.log.Logger;
 import io.trino.sql.parser.ParsingException;
@@ -106,6 +106,7 @@ public class TrinoQueryProperties
     private boolean isNewQuerySubmission;
     private Optional<String> errorMessage = Optional.empty();
     private Optional<String> queryId = Optional.empty();
+    private final ZstdDecompressor decompressor = ZstdDecompressor.create();
 
     public static final String TRINO_CATALOG_HEADER_NAME = "X-Trino-Catalog";
     public static final String TRINO_SCHEMA_HEADER_NAME = "X-Trino-Schema";
@@ -269,8 +270,8 @@ private String decodePreparedStatementFromHeader(String 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);
+        byte[] preparedStatement = new byte[toIntExact(decompressor.getDecompressedSize(compressed, 0, compressed.length))];
+        decompressor.decompress(compressed, 0, compressed.length, preparedStatement, 0, preparedStatement.length);
         return new String(preparedStatement, UTF_8);
     }