# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.
import pytest
from datetime import timedelta
import pyarrow as pa
try:
    import pyarrow.parquet as pq
    import pyarrow.parquet.encryption as pe
except ImportError:
    pq = None
    pe = None
else:
    from pyarrow.tests.parquet.encryption import (InMemoryKmsClient,
                                                  MockVersioningKmsClient,
                                                  verify_file_encrypted,
                                                  read_external_keys_to_dict,
                                                  parse_wrapped_key)


PARQUET_NAME = 'encrypted_table.in_mem.parquet'
FOOTER_KEY = b"0123456789112345"
FOOTER_KEY_NAME = "footer_key"
COL_KEY = b"1234567890123450"
COL_KEY_NAME = "col_key"


# Marks all of the tests in this module
# Ignore these with pytest ... -m 'not parquet_encryption'
# Ignore these with pytest ... -m 'not parquet'
pytestmark = [
    pytest.mark.parquet_encryption,
    pytest.mark.parquet
]


@pytest.fixture(scope='module')
def data_table():
    data_table = pa.Table.from_pydict({
        'a': pa.array([1, 2, 3]),
        'b': pa.array(['a', 'b', 'c']),
        'c': pa.array(['x', 'y', 'z'])
    })
    return data_table


@pytest.fixture(scope='module')
def basic_encryption_config():
    basic_encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        })
    return basic_encryption_config


@pytest.fixture(scope='module')
def external_encryption_config():
    external_encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        },
        internal_key_material=False)
    return external_encryption_config


def setup_encryption_environment(custom_kms_conf):
    """
    Sets up and returns the KMS connection configuration and crypto factory
    based on provided KMS configuration parameters.
    """
    kms_connection_config = pe.KmsConnectionConfig(custom_kms_conf=custom_kms_conf)

    def kms_factory(kms_connection_configuration):
        return InMemoryKmsClient(kms_connection_configuration)

    # Create our CryptoFactory
    crypto_factory = pe.CryptoFactory(kms_factory)

    return kms_connection_config, crypto_factory


def write_encrypted_file(path, data_table, footer_key_name, col_key_name,
                         footer_key, col_key, encryption_config):
    """
    Writes an encrypted parquet file based on the provided parameters.
    """
    # Setup the custom KMS configuration with provided keys
    custom_kms_conf = {
        footer_key_name: footer_key.decode("UTF-8"),
        col_key_name: col_key.decode("UTF-8"),
    }

    # Setup encryption environment
    kms_connection_config, crypto_factory = setup_encryption_environment(
        custom_kms_conf)

    # Write the encrypted parquet file
    write_encrypted_parquet(path, data_table, encryption_config,
                            kms_connection_config, crypto_factory)

    return kms_connection_config, crypto_factory


def test_encrypted_parquet_write_read(tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted, and then read it."""
    path = tempdir / PARQUET_NAME

    # Encrypt the footer with the footer key,
    # encrypt column `a` and column `b` with another key,
    # keep `c` plaintext
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        },
        encryption_algorithm="AES_GCM_V1",
        cache_lifetime=timedelta(minutes=5.0),
        data_key_length_bits=256)
    assert encryption_config.uniform_encryption is False

    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
        encryption_config)

    verify_file_encrypted(path)

    # Read with decryption properties
    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))
    result_table = read_encrypted_parquet(
        path, decryption_config, kms_connection_config, crypto_factory)
    assert data_table.equals(result_table)


def test_uniform_encrypted_parquet_write_read(tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted, and then read it."""
    path = tempdir / PARQUET_NAME

    # Encrypt the footer and all columns with the footer key,
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        uniform_encryption=True,
        encryption_algorithm="AES_GCM_V1",
        cache_lifetime=timedelta(minutes=5.0),
        data_key_length_bits=256)
    assert encryption_config.uniform_encryption is True

    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, b"",
        encryption_config)

    verify_file_encrypted(path)

    # Read with decryption properties
    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))
    result_table = read_encrypted_parquet(
        path, decryption_config, kms_connection_config, crypto_factory)
    assert data_table.equals(result_table)


def write_encrypted_parquet(path, table, encryption_config,
                            kms_connection_config, crypto_factory):
    if encryption_config.internal_key_material:
        file_encryption_properties = crypto_factory.file_encryption_properties(
            kms_connection_config, encryption_config)
    else:
        file_encryption_properties = crypto_factory.file_encryption_properties(
            kms_connection_config, encryption_config, path)
    assert file_encryption_properties is not None
    with pq.ParquetWriter(
            path, table.schema,
            encryption_properties=file_encryption_properties) as writer:
        writer.write_table(table)


def read_encrypted_parquet(path, decryption_config,
                           kms_connection_config, crypto_factory,
                           internal_key_material=True):
    if internal_key_material:
        file_decryption_properties = crypto_factory.file_decryption_properties(
            kms_connection_config, decryption_config)
    else:
        file_decryption_properties = crypto_factory.file_decryption_properties(
            kms_connection_config, decryption_config, path)

    assert file_decryption_properties is not None
    meta = pq.read_metadata(
        path, decryption_properties=file_decryption_properties)
    assert meta.num_columns == 3
    schema = pq.read_schema(
        path, decryption_properties=file_decryption_properties)
    assert len(schema.names) == 3

    result = pq.ParquetFile(
        path, decryption_properties=file_decryption_properties)
    return result.read(use_threads=True)


def test_encrypted_parquet_write_read_wrong_key(tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted,
    and then read it using wrong keys."""
    path = tempdir / PARQUET_NAME

    # Encrypt the footer with the footer key,
    # encrypt column `a` and column `b` with another key,
    # keep `c` plaintext
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        },
        encryption_algorithm="AES_GCM_V1",
        cache_lifetime=timedelta(minutes=5.0),
        data_key_length_bits=256)

    write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME,
                         FOOTER_KEY, COL_KEY, encryption_config)

    verify_file_encrypted(path)

    wrong_kms_connection_config, wrong_crypto_factory = setup_encryption_environment({
        FOOTER_KEY_NAME: COL_KEY.decode("UTF-8"),  # Intentionally wrong
        COL_KEY_NAME: FOOTER_KEY.decode("UTF-8"),  # Intentionally wrong
    })

    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))
    with pytest.raises(ValueError, match=r"Incorrect master key used"):
        read_encrypted_parquet(
            path, decryption_config, wrong_kms_connection_config,
            wrong_crypto_factory)


def test_encrypted_parquet_read_no_decryption_config(tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted,
    but then try to read it without decryption properties."""
    test_encrypted_parquet_write_read(tempdir, data_table)
    # Read without decryption properties
    with pytest.raises(IOError, match=r"no decryption"):
        pq.ParquetFile(tempdir / PARQUET_NAME).read()


def test_encrypted_parquet_read_metadata_no_decryption_config(
        tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted,
    but then try to read its metadata without decryption properties."""
    test_encrypted_parquet_write_read(tempdir, data_table)
    # Read metadata without decryption properties
    with pytest.raises(IOError, match=r"no decryption"):
        pq.read_metadata(tempdir / PARQUET_NAME)


def test_encrypted_parquet_read_schema_no_decryption_config(
        tempdir, data_table):
    """Write an encrypted parquet, verify it's encrypted,
    but then try to read its schema without decryption properties."""
    test_encrypted_parquet_write_read(tempdir, data_table)
    with pytest.raises(IOError, match=r"no decryption"):
        pq.read_schema(tempdir / PARQUET_NAME)


def test_encrypted_parquet_write_no_col_key(tempdir, data_table):
    """Write an encrypted parquet, but give only footer key,
    without column key."""
    path = tempdir / 'encrypted_table_no_col_key.in_mem.parquet'

    # Encrypt the footer with the footer key
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME)

    with pytest.raises(OSError,
                       match="Either column_keys or uniform_encryption "
                       "must be set"):
        # Write with encryption properties
        write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME,
                             FOOTER_KEY, b"", encryption_config)


def test_encrypted_parquet_write_col_key_and_uniform_encryption(tempdir, data_table):
    """Write an encrypted parquet, but give only footer key,
    without column key."""
    path = tempdir / 'encrypted_table_col_key_and_uniform_encryption.in_mem.parquet'

    # Encrypt the footer with the footer key
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        },
        uniform_encryption=True)

    with pytest.raises(OSError,
                       match=r"Cannot set both column_keys and uniform_encryption"):
        # Write with encryption properties
        write_encrypted_file(path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME,
                             FOOTER_KEY, b"", encryption_config)


def test_encrypted_parquet_write_kms_error(tempdir, data_table,
                                           basic_encryption_config):
    """Write an encrypted parquet, but raise KeyError in KmsClient."""
    path = tempdir / 'encrypted_table_kms_error.in_mem.parquet'
    encryption_config = basic_encryption_config

    # Empty master_keys_map
    kms_connection_config = pe.KmsConnectionConfig()

    def kms_factory(kms_connection_configuration):
        # Empty master keys map will cause KeyError to be raised
        # on wrap/unwrap calls
        return InMemoryKmsClient(kms_connection_configuration)

    crypto_factory = pe.CryptoFactory(kms_factory)
    with pytest.raises(KeyError, match="footer_key"):
        # Write with encryption properties
        write_encrypted_parquet(path, data_table, encryption_config,
                                kms_connection_config, crypto_factory)


def test_encrypted_parquet_write_kms_specific_error(tempdir, data_table,
                                                    basic_encryption_config):
    """Write an encrypted parquet, but raise KeyError in KmsClient."""
    path = tempdir / 'encrypted_table_kms_error.in_mem.parquet'
    encryption_config = basic_encryption_config

    # Empty master_keys_map
    kms_connection_config = pe.KmsConnectionConfig()

    class ThrowingKmsClient(pe.KmsClient):
        """A KmsClient implementation that throws exception in
        wrap/unwrap calls
        """

        def __init__(self, config):
            """Create an InMemoryKmsClient instance."""
            pe.KmsClient.__init__(self)
            self.config = config

        def wrap_key(self, key_bytes, master_key_identifier):
            raise ValueError("Cannot Wrap Key")

        def unwrap_key(self, wrapped_key, master_key_identifier):
            raise ValueError("Cannot Unwrap Key")

    def kms_factory(kms_connection_configuration):
        # Exception thrown in wrap/unwrap calls
        return ThrowingKmsClient(kms_connection_configuration)

    crypto_factory = pe.CryptoFactory(kms_factory)
    with pytest.raises(ValueError, match="Cannot Wrap Key"):
        # Write with encryption properties
        write_encrypted_parquet(path, data_table, encryption_config,
                                kms_connection_config, crypto_factory)


def test_encrypted_parquet_write_kms_factory_error(tempdir, data_table,
                                                   basic_encryption_config):
    """Write an encrypted parquet, but raise ValueError in kms_factory."""
    path = tempdir / 'encrypted_table_kms_factory_error.in_mem.parquet'
    encryption_config = basic_encryption_config

    # Empty master_keys_map
    kms_connection_config = pe.KmsConnectionConfig()

    def kms_factory(kms_connection_configuration):
        raise ValueError('Cannot create KmsClient')

    crypto_factory = pe.CryptoFactory(kms_factory)
    with pytest.raises(ValueError,
                       match="Cannot create KmsClient"):
        # Write with encryption properties
        write_encrypted_parquet(path, data_table, encryption_config,
                                kms_connection_config, crypto_factory)


def test_encrypted_parquet_write_kms_factory_type_error(
        tempdir, data_table, basic_encryption_config):
    """Write an encrypted parquet, but use wrong KMS client type
    that doesn't implement KmsClient."""
    path = tempdir / 'encrypted_table_kms_factory_error.in_mem.parquet'
    encryption_config = basic_encryption_config

    # Empty master_keys_map
    kms_connection_config = pe.KmsConnectionConfig()

    class WrongTypeKmsClient():
        """This is not an implementation of KmsClient.
        """

        def __init__(self, config):
            self.master_keys_map = config.custom_kms_conf

        def wrap_key(self, key_bytes, master_key_identifier):
            return None

        def unwrap_key(self, wrapped_key, master_key_identifier):
            return None

    def kms_factory(kms_connection_configuration):
        return WrongTypeKmsClient(kms_connection_configuration)

    crypto_factory = pe.CryptoFactory(kms_factory)
    with pytest.raises(TypeError):
        # Write with encryption properties
        write_encrypted_parquet(path, data_table, encryption_config,
                                kms_connection_config, crypto_factory)


def test_encrypted_parquet_encryption_configuration():
    def validate_encryption_configuration(encryption_config):
        assert FOOTER_KEY_NAME == encryption_config.footer_key
        assert ["a", "b"] == encryption_config.column_keys[COL_KEY_NAME]
        assert "AES_GCM_CTR_V1" == encryption_config.encryption_algorithm
        assert encryption_config.plaintext_footer
        assert not encryption_config.double_wrapping
        assert timedelta(minutes=10.0) == encryption_config.cache_lifetime
        assert not encryption_config.internal_key_material
        assert 192 == encryption_config.data_key_length_bits

    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={COL_KEY_NAME: ["a", "b"], },
        encryption_algorithm="AES_GCM_CTR_V1",
        plaintext_footer=True,
        double_wrapping=False,
        cache_lifetime=timedelta(minutes=10.0),
        internal_key_material=False,
        data_key_length_bits=192,
    )
    validate_encryption_configuration(encryption_config)

    encryption_config_1 = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME)
    encryption_config_1.column_keys = {COL_KEY_NAME: ["a", "b"], }
    encryption_config_1.encryption_algorithm = "AES_GCM_CTR_V1"
    encryption_config_1.plaintext_footer = True
    encryption_config_1.double_wrapping = False
    encryption_config_1.cache_lifetime = timedelta(minutes=10.0)
    encryption_config_1.internal_key_material = False
    encryption_config_1.data_key_length_bits = 192
    validate_encryption_configuration(encryption_config_1)


def test_encrypted_parquet_decryption_configuration():
    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=10.0))
    assert timedelta(minutes=10.0) == decryption_config.cache_lifetime

    decryption_config_1 = pe.DecryptionConfiguration()
    decryption_config_1.cache_lifetime = timedelta(minutes=10.0)
    assert timedelta(minutes=10.0) == decryption_config_1.cache_lifetime


def test_encrypted_parquet_kms_configuration():
    def validate_kms_connection_config(kms_connection_config):
        assert "Instance1" == kms_connection_config.kms_instance_id
        assert "URL1" == kms_connection_config.kms_instance_url
        assert "MyToken" == kms_connection_config.key_access_token
        assert ({"key1": "key_material_1", "key2": "key_material_2"} ==
                kms_connection_config.custom_kms_conf)

    kms_connection_config = pe.KmsConnectionConfig(
        kms_instance_id="Instance1",
        kms_instance_url="URL1",
        key_access_token="MyToken",
        custom_kms_conf={
            "key1": "key_material_1",
            "key2": "key_material_2",
        })
    validate_kms_connection_config(kms_connection_config)

    kms_connection_config_1 = pe.KmsConnectionConfig()
    kms_connection_config_1.kms_instance_id = "Instance1"
    kms_connection_config_1.kms_instance_url = "URL1"
    kms_connection_config_1.key_access_token = "MyToken"
    kms_connection_config_1.custom_kms_conf = {
        "key1": "key_material_1",
        "key2": "key_material_2",
    }
    validate_kms_connection_config(kms_connection_config_1)


@pytest.mark.xfail(reason="Plaintext footer - reading plaintext column subset"
                   " reads encrypted columns too")
def test_encrypted_parquet_write_read_plain_footer_single_wrapping(
        tempdir, data_table):
    """Write an encrypted parquet, with plaintext footer
    and with single wrapping,
    verify it's encrypted, and then read plaintext columns."""
    path = tempdir / PARQUET_NAME

    # Encrypt the footer with the footer key,
    # encrypt column `a` and column `b` with another key,
    # keep `c` plaintext
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={
            COL_KEY_NAME: ["a", "b"],
        },
        plaintext_footer=True,
        double_wrapping=False)

    kms_connection_config = pe.KmsConnectionConfig(
        custom_kms_conf={
            FOOTER_KEY_NAME: FOOTER_KEY.decode("UTF-8"),
            COL_KEY_NAME: COL_KEY.decode("UTF-8"),
        }
    )

    def kms_factory(kms_connection_configuration):
        return InMemoryKmsClient(kms_connection_configuration)

    crypto_factory = pe.CryptoFactory(kms_factory)
    # Write with encryption properties
    write_encrypted_parquet(path, data_table, encryption_config,
                            kms_connection_config, crypto_factory)

    # # Read without decryption properties only the plaintext column
    # result = pq.ParquetFile(path)
    # result_table = result.read(columns='c', use_threads=False)
    # assert table.num_rows == result_table.num_rows


def test_encrypted_parquet_write_read_external(tempdir, data_table,
                                               external_encryption_config):
    """Write an encrypted parquet file with external key material, verify
    it's encrypted, then read both the table and external store.
    """
    path = tempdir / PARQUET_NAME

    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
        external_encryption_config)

    verify_file_encrypted(path)

    decryption_config = pe.DecryptionConfiguration()
    result_table = read_encrypted_parquet(
        path, decryption_config, kms_connection_config, crypto_factory,
        internal_key_material=False)
    store = pa._parquet_encryption.FileSystemKeyMaterialStore.for_file(path)

    assert len(key_ids := store.get_key_id_set()) == (
        len(external_encryption_config.column_keys[COL_KEY_NAME]) + 1)
    assert all([store.get_key_material(k) is not None for k in key_ids])
    assert data_table.equals(result_table)


@pytest.mark.parametrize(
    ("double_wrap_initial", "double_wrap_rotated"), [
        pytest.param(True, True, id="double wrapping"),
        pytest.param(False, True, id="single to double wrapped"),
        pytest.param(True, False, id="double to singe wrapped"),
        pytest.param(False, False, id="single wrapping")])
def test_external_key_material_rotation(
        reusable_tempdir,
        data_table,
        double_wrap_initial,
        double_wrap_rotated):
    """Tests CryptoFactory.rotate_master_keys

    Note: The CryptoFactory.rotate_master_keys() double_wrapping keword arg
    may be either True (the default) or False regardless of whether
    EncryptionConfig.double_wrapping was set to true (also the default) when
    the external key material store was written. This means double wrapping may
    be set one way initially and then applied or removed during rotation.
    """
    path = reusable_tempdir / PARQUET_NAME
    encryption_config = pe.EncryptionConfiguration(
        footer_key=FOOTER_KEY_NAME,
        column_keys={COL_KEY_NAME: ["a", "b"]},
        internal_key_material=False,
        double_wrapping=double_wrap_initial)

    # initial master key version - see MockVersioningKmsClient docstring
    kms_connection_config = pe.KmsConnectionConfig(key_access_token="1")

    def kms_factory(kms_connection_configuration):
        return MockVersioningKmsClient(kms_connection_configuration)
    crypto_factory = pe.CryptoFactory(kms_factory)
    write_encrypted_parquet(
        path,
        data_table,
        encryption_config,
        kms_connection_config,
        crypto_factory)
    before_keys = read_external_keys_to_dict(path)

    # "rotate" kms master key
    kms_connection_config.refresh_key_access_token("2")

    crypto_factory.rotate_master_keys(
        kms_connection_config,
        path,
        double_wrapping=double_wrap_rotated)

    after_keys = read_external_keys_to_dict(path)
    verify_file_encrypted(path)
    table_read_after_rotation = read_encrypted_parquet(
        path,
        pe.DecryptionConfiguration(),
        kms_connection_config,
        crypto_factory,
        internal_key_material=False)
    assert FOOTER_KEY_NAME in before_keys
    assert COL_KEY_NAME in before_keys
    assert FOOTER_KEY_NAME in after_keys
    assert COL_KEY_NAME in after_keys

    def check_rotated_external_keys(master_key_id: str) -> None:
        before_key_mat = before_keys[master_key_id]
        if double_wrap_initial:
            before_key_wrapped = before_key_mat.wrapped_kek
        else:
            before_key_wrapped = before_key_mat.wrapped_dek
        _, before_ver, _ = parse_wrapped_key(before_key_wrapped)

        after_key_mat = after_keys[master_key_id]
        if double_wrap_rotated:
            after_key_wrapped = after_key_mat.wrapped_kek
        else:
            after_key_wrapped = after_key_mat.wrapped_dek
        _, after_ver, _ = parse_wrapped_key(after_key_wrapped)

        # CryptoFactory rewrapped keys if after version is later than before
        assert before_ver < after_ver
    check_rotated_external_keys(FOOTER_KEY_NAME)
    check_rotated_external_keys(COL_KEY_NAME)
    assert data_table.equals(table_read_after_rotation)


def test_encrypted_parquet_loop(tempdir, data_table, basic_encryption_config):
    """Write an encrypted parquet, verify it's encrypted,
    and then read it multithreaded in a loop."""
    path = tempdir / PARQUET_NAME

    # Encrypt the footer with the footer key,
    # encrypt column `a` and column `b` with another key,
    # keep `c` plaintext, defined in basic_encryption_config
    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
        basic_encryption_config)

    verify_file_encrypted(path)

    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))

    for i in range(50):
        # Read with decryption properties
        file_decryption_properties = crypto_factory.file_decryption_properties(
            kms_connection_config, decryption_config)
        assert file_decryption_properties is not None

        result = pq.ParquetFile(
            path, decryption_properties=file_decryption_properties)
        result_table = result.read(use_threads=True)
        assert data_table.equals(result_table)


def test_read_with_deleted_crypto_factory(tempdir, data_table, basic_encryption_config):
    """
    Test that decryption properties can be used if the crypto factory is no longer alive
    """
    path = tempdir / PARQUET_NAME
    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
        basic_encryption_config)
    verify_file_encrypted(path)

    # Create decryption properties and delete the crypto factory that created
    # the properties afterwards.
    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))
    file_decryption_properties = crypto_factory.file_decryption_properties(
        kms_connection_config, decryption_config)
    del crypto_factory

    result = pq.ParquetFile(
        path, decryption_properties=file_decryption_properties)
    result_table = result.read(use_threads=True)
    assert data_table.equals(result_table)


def test_encrypted_parquet_read_table(tempdir, data_table, basic_encryption_config):
    """Write an encrypted parquet then read it back using read_table."""
    path = tempdir / PARQUET_NAME

    # Write the encrypted parquet file using the utility function
    kms_connection_config, crypto_factory = write_encrypted_file(
        path, data_table, FOOTER_KEY_NAME, COL_KEY_NAME, FOOTER_KEY, COL_KEY,
        basic_encryption_config)

    decryption_config = pe.DecryptionConfiguration(
        cache_lifetime=timedelta(minutes=5.0))
    file_decryption_properties = crypto_factory.file_decryption_properties(
        kms_connection_config, decryption_config)

    # Read the encrypted parquet file using read_table
    result_table = pq.read_table(path, decryption_properties=file_decryption_properties)

    # Assert that the read table matches the original data
    assert data_table.equals(result_table)

    # Read the encrypted parquet folder using read_table
    result_table = pq.read_table(
        tempdir, decryption_properties=file_decryption_properties)
    assert data_table.equals(result_table)