Support Parquet Modular Encryption for Trino

lib/trino-parquet/src/main/java/io/trino/parquet/DataPage.java

@@ -13,6 +13,8 @@
  */
 package io.trino.parquet;
 
+import io.airlift.slice.Slice;
+
 import java.util.OptionalLong;
 
 public abstract sealed class DataPage
@@ -41,4 +43,6 @@ public int getValueCount()
     {
         return valueCount;
     }
+
+    public abstract Slice getSlice();
 }

lib/trino-parquet/src/main/java/io/trino/parquet/EncryptionUtils.java

@@ -0,0 +1,44 @@
+/*
+ * 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.parquet;
+
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.Path;
+import org.apache.parquet.crypto.DecryptionPropertiesFactory;
+import org.apache.parquet.crypto.FileDecryptionProperties;
+import org.apache.parquet.crypto.InternalFileDecryptor;
+
+import java.util.Optional;
+
+import static org.apache.parquet.crypto.DecryptionPropertiesFactory.loadFactory;
+
+public class EncryptionUtils
+{
+    private static final String PARQUET_CRYPTO_FACTORY_CLASS = "parquet.crypto.factory.class";
+    private static final String PARQUET_ENCRYPTION_KMS_CLIENT_CLASS = "parquet.encryption.kms.client.class";
+
+    private EncryptionUtils() {}
+
+    public static Optional<InternalFileDecryptor> createDecryptor(ParquetReaderOptions parquetReaderOptions, String path)
+    {
+        // TODO(wyu): not a fan of recreting configuration with the options, but not sure if I should add this dep to something like trino-hive, hudi, etc.
+        // maybe add it to the trino hadoop-apache dep?
+        Configuration configuration = new Configuration();
+        configuration.set(PARQUET_CRYPTO_FACTORY_CLASS, parquetReaderOptions.getCryptoFactoryClass());
+        configuration.set(PARQUET_ENCRYPTION_KMS_CLIENT_CLASS, parquetReaderOptions.getEncryptionKmsClientClass());
+        DecryptionPropertiesFactory cryptoFactory = loadFactory(configuration);
+        FileDecryptionProperties fileDecryptionProperties = (cryptoFactory == null) ? null : cryptoFactory.getFileDecryptionProperties(configuration, new Path(path));
+        return (fileDecryptionProperties == null) ? Optional.empty() : Optional.of(new InternalFileDecryptor(fileDecryptionProperties));
+    }
+}

lib/trino-parquet/src/main/java/io/trino/parquet/ParquetReaderOptions.java

@@ -35,6 +35,8 @@ public class ParquetReaderOptions
     private final boolean useColumnIndex;
     private final boolean useBloomFilter;
     private final DataSize smallFileThreshold;
+    private final String cryptoFactoryClass;
+    private final String encryptionKmsClientClass;
 
     public ParquetReaderOptions()
     {
@@ -46,6 +48,8 @@ public ParquetReaderOptions()
         useColumnIndex = true;
         useBloomFilter = true;
         smallFileThreshold = DEFAULT_SMALL_FILE_THRESHOLD;
+        cryptoFactoryClass = null;
+        encryptionKmsClientClass = null;
     }
 
     private ParquetReaderOptions(
@@ -56,7 +60,9 @@ private ParquetReaderOptions(
             DataSize maxBufferSize,
             boolean useColumnIndex,
             boolean useBloomFilter,
-            DataSize smallFileThreshold)
+            DataSize smallFileThreshold,
+            String cryptoFactoryClass,
+            String encryptionKmsClientClass)
     {
         this.ignoreStatistics = ignoreStatistics;
         this.maxReadBlockSize = requireNonNull(maxReadBlockSize, "maxReadBlockSize is null");
@@ -67,6 +73,8 @@ private ParquetReaderOptions(
         this.useColumnIndex = useColumnIndex;
         this.useBloomFilter = useBloomFilter;
         this.smallFileThreshold = requireNonNull(smallFileThreshold, "smallFileThreshold is null");
+        this.cryptoFactoryClass = cryptoFactoryClass;
+        this.encryptionKmsClientClass = encryptionKmsClientClass;
     }
 
     public boolean isIgnoreStatistics()
@@ -109,6 +117,16 @@ public DataSize getSmallFileThreshold()
         return smallFileThreshold;
     }
 
+    public String getCryptoFactoryClass()
+    {
+        return cryptoFactoryClass;
+    }
+
+    public String getEncryptionKmsClientClass()
+    {
+        return encryptionKmsClientClass;
+    }
+
     public ParquetReaderOptions withIgnoreStatistics(boolean ignoreStatistics)
     {
         return new ParquetReaderOptions(
@@ -119,7 +137,9 @@ public ParquetReaderOptions withIgnoreStatistics(boolean ignoreStatistics)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withMaxReadBlockSize(DataSize maxReadBlockSize)
@@ -132,7 +152,9 @@ public ParquetReaderOptions withMaxReadBlockSize(DataSize maxReadBlockSize)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withMaxReadBlockRowCount(int maxReadBlockRowCount)
@@ -145,7 +167,9 @@ public ParquetReaderOptions withMaxReadBlockRowCount(int maxReadBlockRowCount)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withMaxMergeDistance(DataSize maxMergeDistance)
@@ -158,7 +182,9 @@ public ParquetReaderOptions withMaxMergeDistance(DataSize maxMergeDistance)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withMaxBufferSize(DataSize maxBufferSize)
@@ -171,7 +197,9 @@ public ParquetReaderOptions withMaxBufferSize(DataSize maxBufferSize)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withUseColumnIndex(boolean useColumnIndex)
@@ -184,7 +212,9 @@ public ParquetReaderOptions withUseColumnIndex(boolean useColumnIndex)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withBloomFilter(boolean useBloomFilter)
@@ -197,7 +227,9 @@ public ParquetReaderOptions withBloomFilter(boolean useBloomFilter)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 
     public ParquetReaderOptions withSmallFileThreshold(DataSize smallFileThreshold)
@@ -210,6 +242,38 @@ public ParquetReaderOptions withSmallFileThreshold(DataSize smallFileThreshold)
                 maxBufferSize,
                 useColumnIndex,
                 useBloomFilter,
-                smallFileThreshold);
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
+    }
+
+    public ParquetReaderOptions withCryptoFactoryClass(String cryptoFactoryClass)
+    {
+        return new ParquetReaderOptions(
+                ignoreStatistics,
+                maxReadBlockSize,
+                maxReadBlockRowCount,
+                maxMergeDistance,
+                maxBufferSize,
+                useColumnIndex,
+                useBloomFilter,
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
+    }
+
+    public ParquetReaderOptions withEncryptionKmsClientClass(String encryptionKmsClientClass)
+    {
+        return new ParquetReaderOptions(
+                ignoreStatistics,
+                maxReadBlockSize,
+                maxReadBlockRowCount,
+                maxMergeDistance,
+                maxBufferSize,
+                useColumnIndex,
+                useBloomFilter,
+                smallFileThreshold,
+                cryptoFactoryClass,
+                encryptionKmsClientClass);
     }
 }