diff --git a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java index 82f4577b83..022e082ce8 100644 --- a/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java +++ b/parquet-hadoop/src/main/java/org/apache/parquet/hadoop/ParquetFileWriter.java @@ -519,36 +519,44 @@ private ParquetFileWriter( return; } - if (null == encryptionProperties) { - encryptionProperties = encryptor.getEncryptionProperties(); - } + try { + if (null == encryptionProperties) { + encryptionProperties = encryptor.getEncryptionProperties(); + } - // Verify that every encrypted column is in file schema - Map columnEncryptionProperties = - encryptionProperties.getEncryptedColumns(); - if (null != columnEncryptionProperties) { // if null, every column in file schema will be encrypted with footer - // key - for (Map.Entry entry : columnEncryptionProperties.entrySet()) { - String[] path = entry.getKey().toArray(); - if (!schema.containsPath(path)) { - StringBuilder columnList = new StringBuilder(); - columnList.append("["); - for (String[] columnPath : schema.getPaths()) { - columnList - .append(ColumnPath.get(columnPath).toDotString()) - .append("], ["); + // Verify that every encrypted column is in file schema + Map columnEncryptionProperties = + encryptionProperties.getEncryptedColumns(); + if (null != columnEncryptionProperties) { // if null, every column in file schema will be encrypted with + // footer + // key + for (Map.Entry entry : columnEncryptionProperties.entrySet()) { + String[] path = entry.getKey().toArray(); + if (!schema.containsPath(path)) { + StringBuilder columnList = new StringBuilder(); + columnList.append("["); + for (String[] columnPath : schema.getPaths()) { + columnList + .append(ColumnPath.get(columnPath).toDotString()) + .append("], ["); + } + throw new ParquetCryptoRuntimeException("Encrypted column [" + + entry.getKey().toDotString() + "] not in file schema column list: " + + columnList.substring(0, columnList.length() - 3)); } - throw new ParquetCryptoRuntimeException( - "Encrypted column [" + entry.getKey().toDotString() + "] not in file schema column list: " - + columnList.substring(0, columnList.length() - 3)); } } - } - if (null == encryptor) { - this.fileEncryptor = new InternalFileEncryptor(encryptionProperties); - } else { - this.fileEncryptor = encryptor; + if (null == encryptor) { + this.fileEncryptor = new InternalFileEncryptor(encryptionProperties); + } else { + this.fileEncryptor = encryptor; + } + } catch (Exception e) { + // If encryption setup throws in the constructor, the output stream opened above should be + // closed. Otherwise, there's no way to close it outside since the object is never returned. + out.close(); + throw e; } } diff --git a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileWriter.java b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileWriter.java index ca03ef4db8..f8f4c72591 100644 --- a/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileWriter.java +++ b/parquet-hadoop/src/test/java/org/apache/parquet/hadoop/TestParquetFileWriter.java @@ -33,6 +33,7 @@ import static org.apache.parquet.schema.Type.Repetition.OPTIONAL; import static org.apache.parquet.schema.Type.Repetition.REPEATED; import static org.apache.parquet.schema.Type.Repetition.REQUIRED; +import static org.assertj.core.api.Assertions.assertThatThrownBy; import static org.junit.Assert.assertArrayEquals; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; @@ -51,6 +52,7 @@ import java.util.List; import java.util.Map; import java.util.concurrent.Callable; +import java.util.concurrent.atomic.AtomicBoolean; import org.apache.hadoop.conf.Configuration; import org.apache.hadoop.fs.FSDataInputStream; import org.apache.hadoop.fs.FileStatus; @@ -79,6 +81,9 @@ import org.apache.parquet.column.statistics.LongStatistics; import org.apache.parquet.column.values.bloomfilter.BlockSplitBloomFilter; import org.apache.parquet.column.values.bloomfilter.BloomFilter; +import org.apache.parquet.crypto.ColumnEncryptionProperties; +import org.apache.parquet.crypto.FileEncryptionProperties; +import org.apache.parquet.crypto.ParquetCryptoRuntimeException; import org.apache.parquet.example.data.Group; import org.apache.parquet.example.data.simple.SimpleGroup; import org.apache.parquet.format.Statistics; @@ -87,6 +92,7 @@ import org.apache.parquet.hadoop.example.GroupWriteSupport; import org.apache.parquet.hadoop.metadata.BlockMetaData; import org.apache.parquet.hadoop.metadata.ColumnChunkMetaData; +import org.apache.parquet.hadoop.metadata.ColumnPath; import org.apache.parquet.hadoop.metadata.CompressionCodecName; import org.apache.parquet.hadoop.metadata.ConcatenatingKeyValueMetadataMergeStrategy; import org.apache.parquet.hadoop.metadata.FileMetaData; @@ -101,7 +107,9 @@ import org.apache.parquet.internal.column.columnindex.BoundaryOrder; import org.apache.parquet.internal.column.columnindex.ColumnIndex; import org.apache.parquet.internal.column.columnindex.OffsetIndex; +import org.apache.parquet.io.OutputFile; import org.apache.parquet.io.ParquetEncodingException; +import org.apache.parquet.io.PositionOutputStream; import org.apache.parquet.io.api.Binary; import org.apache.parquet.schema.MessageType; import org.apache.parquet.schema.MessageTypeParser; @@ -1498,6 +1506,91 @@ public void testMergeMetadataWithNoConflictingKeyValues() { assertEquals("d", mergedValues.get("c")); } + @Test + public void testConstructorClosesStreamWhenEncryptedColumnMissing() throws Exception { + ColumnEncryptionProperties missingColumn = ColumnEncryptionProperties.builder("not_in_schema") + .withKey("0123456789012345".getBytes(StandardCharsets.UTF_8)) + .build(); + Map encryptedColumns = new HashMap<>(); + encryptedColumns.put(missingColumn.getPath(), missingColumn); + FileEncryptionProperties encryptionProperties = FileEncryptionProperties.builder( + "0123456789012345".getBytes(StandardCharsets.UTF_8)) + .withEncryptedColumns(encryptedColumns) + .build(); + + RecordingOutputFile file = new RecordingOutputFile(); + assertThatThrownBy(() -> new ParquetFileWriter( + file, + SCHEMA, + CREATE, + DEFAULT_BLOCK_SIZE, + MAX_PADDING_SIZE_DEFAULT, + ParquetProperties.DEFAULT_COLUMN_INDEX_TRUNCATE_LENGTH, + ParquetProperties.DEFAULT_STATISTICS_TRUNCATE_LENGTH, + ParquetProperties.DEFAULT_PAGE_WRITE_CHECKSUM_ENABLED, + encryptionProperties)) + .isInstanceOf(ParquetCryptoRuntimeException.class) + .hasMessage("Encrypted column [not_in_schema] not in file schema column list: [a.b], [c.d]"); + + assertTrue( + "The output stream opened by the constructor must be closed when validation fails", + file.isStreamClosed()); + } + + /** + * An {@link OutputFile} whose {@link PositionOutputStream} records whether it was closed, used to + * assert that a failed {@link ParquetFileWriter} construction does not leak the open stream. + */ + private static class RecordingOutputFile implements OutputFile { + + private final AtomicBoolean streamClosed = new AtomicBoolean(false); + + boolean isStreamClosed() { + return streamClosed.get(); + } + + private PositionOutputStream newRecordingStream() { + return new PositionOutputStream() { + private long pos = 0; + + @Override + public long getPos() { + return pos; + } + + @Override + public void write(int b) { + pos++; + } + + @Override + public void close() { + streamClosed.set(true); + } + }; + } + + @Override + public PositionOutputStream create(long blockSizeHint) { + return newRecordingStream(); + } + + @Override + public PositionOutputStream createOrOverwrite(long blockSizeHint) { + return newRecordingStream(); + } + + @Override + public boolean supportsBlockSize() { + return false; + } + + @Override + public long defaultBlockSize() { + return 0; + } + } + private org.apache.parquet.column.statistics.Statistics statsC1(Binary... values) { org.apache.parquet.column.statistics.Statistics stats = org.apache.parquet.column.statistics.Statistics.createStats(C1.getPrimitiveType());