universal-driver/python/tests/e2e/put_get/test_put_get_auto_compress.py at 9a12f44c5b0ced800bd979c76773f9cdcb70ec1f · snowflakedb/universal-driver · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
import tempfile

from pathlib import Path

from tests.compatibility import NEW_DRIVER_ONLY, OLD_DRIVER_ONLY
from tests.e2e.put_get.put_get_helper import (
    create_temporary_stage_and_upload_file,
    get_file_from_stage,
)
from tests.utils import shared_test_data_dir


def test_should_compress_the_file_before_uploading_to_stage_when_auto_compress_set_to_true(
    connection,
):
    uncompressed_file_path = shared_test_data_dir() / "compression" / "test_data.csv"
    compressed_file_path = shared_test_data_dir() / "compression" / "test_data.csv.gz"
    uncompressed_filename = "test_data.csv"
    compressed_filename = "test_data.csv.gz"
    # Given Snowflake client is logged in
    with connection.cursor() as cursor:
        # When File is uploaded to stage with AUTO_COMPRESS set to true
        stage_name, _ = create_temporary_stage_and_upload_file(
            cursor,
            "TEST_PUT_GET_AUTO_COMPRESS_TRUE",
            uncompressed_file_path,
            auto_compress=True,
            overwrite=True,
        )

        with tempfile.TemporaryDirectory() as temp_dir:
            # Then Only compressed file should be downloaded
            download_dir = Path(temp_dir)

            get_result = get_file_from_stage(cursor, stage_name, uncompressed_filename, download_dir)

            assert get_result[2] == "DOWNLOADED"

            expected_file_path = download_dir / compressed_filename
            assert expected_file_path.exists()

            not_expected_file_path = download_dir / uncompressed_filename
            assert not not_expected_file_path.exists()

            # And Have correct content
            downloaded_content = expected_file_path.read_bytes()
            reference_content = compressed_file_path.read_bytes()

            if OLD_DRIVER_ONLY("BD#1"):
                assert downloaded_content != reference_content

            elif NEW_DRIVER_ONLY("BD#1"):
                assert downloaded_content == reference_content


def test_should_not_compress_the_file_before_uploading_to_stage_when_auto_compress_set_to_false(
    connection,
):
    uncompressed_file_path = shared_test_data_dir() / "compression" / "test_data.csv"
    uncompressed_filename = "test_data.csv"
    compressed_filename = "test_data.csv.gz"

    # Given Snowflake client is logged in
    with connection.cursor() as cursor:
        # When File is uploaded to stage with AUTO_COMPRESS set to false
        stage_name, _ = create_temporary_stage_and_upload_file(
            cursor,
            "TEST_PUT_GET_AUTO_COMPRESS_FALSE",
            uncompressed_file_path,
            auto_compress=False,
            overwrite=True,
        )

        with tempfile.TemporaryDirectory() as temp_dir:
            # Then Only uncompressed file should be downloaded
            download_dir = Path(temp_dir)
            get_result = get_file_from_stage(cursor, stage_name, uncompressed_filename, download_dir)

            assert get_result[2] == "DOWNLOADED"

            expected_file_path = download_dir / uncompressed_filename
            assert expected_file_path.exists()

            not_expected_file_path = download_dir / compressed_filename
            assert not not_expected_file_path.exists()

            # And Have correct content
            downloaded_content = expected_file_path.read_bytes()
            reference_content = uncompressed_file_path.read_bytes()
            assert downloaded_content == reference_content