(
read_keys,
read_columns,
to_table_success,
dataset_success=True,
)
| 123 | |
| 124 | # define the actual test |
| 125 | def assert_decrypts( |
| 126 | read_keys, |
| 127 | read_columns, |
| 128 | to_table_success, |
| 129 | dataset_success=True, |
| 130 | ): |
| 131 | # use all keys for writing |
| 132 | write_keys = keys |
| 133 | encryption_config = create_encryption_config(FOOTER_KEY_NAME, column_keys) |
| 134 | decryption_config = create_decryption_config() |
| 135 | encrypt_kms_connection_config = create_kms_connection_config(write_keys) |
| 136 | decrypt_kms_connection_config = create_kms_connection_config(read_keys) |
| 137 | |
| 138 | crypto_factory = pe.CryptoFactory(kms_factory) |
| 139 | parquet_encryption_cfg = ds.ParquetEncryptionConfig( |
| 140 | crypto_factory, encrypt_kms_connection_config, encryption_config |
| 141 | ) |
| 142 | parquet_decryption_cfg = ds.ParquetDecryptionConfig( |
| 143 | crypto_factory, decrypt_kms_connection_config, decryption_config |
| 144 | ) |
| 145 | |
| 146 | # create write_options with dataset encryption config |
| 147 | pformat = pa.dataset.ParquetFileFormat() |
| 148 | write_options = pformat.make_write_options( |
| 149 | encryption_config=parquet_encryption_cfg |
| 150 | ) |
| 151 | |
| 152 | mockfs = fs._MockFileSystem() |
| 153 | mockfs.create_dir("/") |
| 154 | |
| 155 | ds.write_dataset( |
| 156 | data=table, |
| 157 | base_dir="sample_dataset", |
| 158 | format=pformat, |
| 159 | file_options=write_options, |
| 160 | filesystem=mockfs, |
| 161 | ) |
| 162 | |
| 163 | # read without decryption config -> errors if dataset was properly encrypted |
| 164 | pformat = pa.dataset.ParquetFileFormat() |
| 165 | with pytest.raises(IOError, match=r"no decryption"): |
| 166 | ds.dataset("sample_dataset", format=pformat, filesystem=mockfs) |
| 167 | |
| 168 | # set decryption config for parquet fragment scan options |
| 169 | pq_scan_opts = ds.ParquetFragmentScanOptions( |
| 170 | decryption_config=parquet_decryption_cfg |
| 171 | ) |
| 172 | pformat = pa.dataset.ParquetFileFormat( |
| 173 | default_fragment_scan_options=pq_scan_opts |
| 174 | ) |
| 175 | with cond_raises(dataset_success, ValueError, match="Unknown master key"): |
| 176 | dataset = ds.dataset("sample_dataset", format=pformat, filesystem=mockfs) |
| 177 | with cond_raises(to_table_success, ValueError, match="Unknown master key"): |
| 178 | assert table.select(read_columns).equals(dataset.to_table(read_columns)) |
| 179 | |
| 180 | # set decryption properties for parquet fragment scan options |
| 181 | decryption_properties = crypto_factory.file_decryption_properties( |
| 182 | decrypt_kms_connection_config, decryption_config) |
no test coverage detected