def test_upload_one_file_to_s3_wsaeconnaborted():
"""Tests Upload one file to S3 with retry on ERRORNO_WSAECONNABORTED.
Notes:
The last attempted max_currency should be (initial_parallel/max_retry).
"""
upload_file = MagicMock(
side_effect=OpenSSL.SSL.SysCallError(
ERRORNO_WSAECONNABORTED, 'mock err. connection aborted'))
s3object = MagicMock(metadata=defaultdict(str), upload_file=upload_file)
client = Mock()
client.Object.return_value = s3object
initial_parallel = 100
client_meta_dict = {
'stage_info': {
'location': 'sfc-customer-stage/rwyi-testacco/users/9220/',
'locationType': 'S3',
},
'cloud_client': client,
}
client_meta = SFResourceMeta(**client_meta_dict)
upload_meta = {
'name': 'data1.txt.gz',
'stage_location_type': 'S3',
'no_sleeping_time': True,
'parallel': initial_parallel,
'put_callback': None,
'put_callback_output_stream': None,
'client_meta': client_meta,
SHA256_DIGEST: '123456789abcdef',
'dst_file_name': 'data1.txt.gz',
'src_file_name': path.join(THIS_DIR, '../data', 'put_get_1.txt'),
'overwrite': True,
}
upload_meta['real_src_file_name'] = upload_meta['src_file_name']
upload_meta['upload_size'] = os.stat(upload_meta['src_file_name']).st_size
meta = SnowflakeFileMeta(**upload_meta)
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert meta.last_max_concurrency is not None
assert meta.last_max_concurrency == initial_parallel / DEFAULT_MAX_RETRY
# min parallel == 1
upload_file.reset_mock()
initial_parallel = 4
meta.parallel = initial_parallel
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert meta.last_max_concurrency is not None
assert meta.last_max_concurrency == 1
def test_upload_one_file_to_s3_wsaeconnaborted():
"""Tests Upload one file to S3 with retry on ERRORNO_WSAECONNABORTED.
Notes:
The last attempted max_currency should be (initial_parallel/max_retry).
"""
upload_file = MagicMock(side_effect=OpenSSL.SSL.SysCallError(
ERRORNO_WSAECONNABORTED, "mock err. connection aborted"))
s3object = MagicMock(metadata=defaultdict(str), upload_file=upload_file)
client = Mock()
client.Object.return_value = s3object
initial_parallel = 100
client_meta_dict = {
"stage_info": {
"location": "sfc-customer-stage/rwyi-testacco/users/9220/",
"locationType": "S3",
},
"cloud_client": client,
}
client_meta = SFResourceMeta(**client_meta_dict)
upload_meta = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"parallel": initial_parallel,
"put_callback": None,
"put_callback_output_stream": None,
"client_meta": client_meta,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
upload_meta["real_src_file_name"] = upload_meta["src_file_name"]
upload_meta["upload_size"] = os.stat(upload_meta["src_file_name"]).st_size
meta = SnowflakeFileMeta(**upload_meta)
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert meta.last_max_concurrency is not None
assert meta.last_max_concurrency == initial_parallel / DEFAULT_MAX_RETRY
# min parallel == 1
upload_file.reset_mock()
initial_parallel = 4
meta.parallel = initial_parallel
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert meta.last_max_concurrency is not None
assert meta.last_max_concurrency == 1
def test_download_unknown_error(caplog):
"""Tests whether an unknown error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
mock_resource = MagicMock()
mock_resource.download_file.side_effect = botocore.exceptions.ClientError(
{'Error': {'Code': 'unknown', 'Message': 'Just testing'}}, 'Testing')
client_meta = {
'cloud_client': mock_resource,
'stage_info': {'location': 'loc'},
}
meta = {'name': 'f',
'src_file_name': 'f',
'stage_location_type': 'S3',
'client_meta': SFResourceMeta(**client_meta),
'sha256_digest': 'asd',
'src_file_size': 99,
'get_callback_output_stream': None,
'show_progress_bar': False,
'get_callback': None}
meta = SnowflakeFileMeta(**meta)
with mock.patch('snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object', return_value=mock_resource):
with pytest.raises(botocore.exceptions.ClientError,
match=r'An error occurred \(unknown\) when calling the Testing operation: Just testing'):
SnowflakeS3Util._native_download_file(meta, 'f', 4)
assert ('snowflake.connector.s3_util',
logging.DEBUG,
'Failed to download a file: f, err: An error occurred (unknown) when '
'calling the Testing operation: Just testing') in caplog.record_tuples
def test_upload_unknown_error(caplog):
"""Tests whether unknown errors are handled as expected when uploading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
mock_resource, mock_object = MagicMock(), MagicMock()
mock_resource.Object.return_value = mock_object
mock_object.bucket_name = 'bucket'
mock_object.key = 'key'
mock_object.upload_file.side_effect = botocore.exceptions.ClientError(
{'Error': {'Code': 'unknown', 'Message': 'Just testing'}}, 'Testing')
client_meta = {
'cloud_client': mock_resource,
'stage_info': {'location': 'loc'},
}
meta = {'name': 'f',
'src_file_name': 'f',
'stage_location_type': 'S3',
'client_meta': SFResourceMeta(**client_meta),
'sha256_digest': 'asd',
'dst_file_name': 'f',
'put_callback': None}
meta = SnowflakeFileMeta(**meta)
with mock.patch('snowflake.connector.s3_util.SnowflakeS3Util.extract_bucket_name_and_path'):
with pytest.raises(botocore.exceptions.ClientError,
match=r'An error occurred \(unknown\) when calling the Testing operation: Just testing'):
SnowflakeS3Util.upload_file('f', meta, {}, 4, 67108864)
def test_upload_failed_error(caplog):
"""Tests whether token expiry error is handled as expected when uploading."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
mock_resource, mock_object = MagicMock(), MagicMock()
mock_resource.Object.return_value = mock_object
mock_object.upload_file.side_effect = S3UploadFailedError("ExpiredToken")
client_meta = {
"cloud_client": mock_resource,
"stage_info": {
"location": "loc"
},
}
meta = {
"name": "f",
"src_file_name": "f",
"stage_location_type": "S3",
"client_meta": SFResourceMeta(**client_meta),
"sha256_digest": "asd",
"dst_file_name": "f",
"put_callback": None,
}
meta = SnowflakeFileMeta(**meta)
with mock.patch(
"snowflake.connector.s3_util.SnowflakeS3Util.extract_bucket_name_and_path"
):
assert SnowflakeS3Util.upload_file("f", meta, {}, 4, 67108864) is None
assert (
"snowflake.connector.s3_util",
logging.DEBUG,
"Failed to upload a file: f, err: ExpiredToken. Renewing AWS Token and Retrying",
) in caplog.record_tuples
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_get_s3_file_object_http_400_error():
"""Tests Get S3 file object with HTTP 400 error.
Looks like HTTP 400 is returned when AWS token expires and S3.Object.load is called.
"""
load_method = MagicMock(
side_effect=botocore.exceptions.ClientError(
{'Error': {'Code': '400', 'Message': 'Bad Request'}},
operation_name='mock load'))
s3object = MagicMock(load=load_method)
client = Mock()
client.Object.return_value = s3object
client.load.return_value = None
type(client).s3path = PropertyMock(return_value='s3://testbucket/')
client_meta = {
'cloud_client': client,
'stage_info': {
'location': 'sfc-teststage/rwyitestacco/users/1234/',
'locationType': 'S3',
}
}
meta = {
'name': 'data1.txt.gz',
'stage_location_type': 'S3',
'src_file_name': path.join(THIS_DIR, '../data', 'put_get_1.txt'),
'client_meta': SFResourceMeta(**client_meta),
}
meta = SnowflakeFileMeta(**meta)
filename = "/path1/file2.txt"
akey = SnowflakeS3Util.get_file_header(meta, filename)
assert akey is None
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_upload_put_timeout(tmp_path, caplog):
"""Tests whether timeout error is handled correctly when uploading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
f_name = str(tmp_path / 'some_file.txt')
resp = requests.Response()
meta = SnowflakeFileMeta(
name=f_name,
src_file_name=f_name,
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with open(f_name, 'w') as f:
f.write(random_string(15))
with mock.patch('snowflake.connector.vendored.requests.put'
if vendored_request else 'requests.put',
side_effect=requests.exceptions.Timeout(response=resp)):
SnowflakeGCSUtil.upload_file(f_name, meta, None, 99, 64000)
assert isinstance(meta.last_error, requests.exceptions.Timeout)
assert meta.result_status == ResultStatus.NEED_RETRY
assert all([
log in caplog.record_tuples
for log in [('snowflake.connector.gcs_util', logging.DEBUG,
'GCS file upload Timeout Error: ')]
])
def test_download_syscall_error(caplog, error_no, result_status):
"""Tests whether a syscall error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
mock_resource = MagicMock()
mock_resource.download_file.side_effect = OpenSSL.SSL.SysCallError(
error_no)
client_meta = {
"cloud_client": mock_resource,
"stage_info": {
"location": "loc"
},
}
meta = {
"name": "f",
"stage_location_type": "S3",
"client_meta": SFResourceMeta(**client_meta),
"sha256_digest": "asd",
"src_file_name": "f",
"src_file_size": 99,
"get_callback_output_stream": None,
"show_progress_bar": False,
"get_callback": None,
}
meta = SnowflakeFileMeta(**meta)
with mock.patch(
"snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object",
return_value=mock_resource,
):
SnowflakeS3Util._native_download_file(meta, "f", 4)
assert meta.last_error is mock_resource.download_file.side_effect
assert meta.result_status == result_status
def test_get_s3_file_object_http_400_error():
"""Tests Get S3 file object with HTTP 400 error.
Looks like HTTP 400 is returned when AWS token expires and S3.Object.load is called.
"""
load_method = MagicMock(side_effect=botocore.exceptions.ClientError(
{"Error": {
"Code": "400",
"Message": "Bad Request"
}},
operation_name="mock load",
))
s3object = MagicMock(load=load_method)
client = Mock()
client.Object.return_value = s3object
client.load.return_value = None
type(client).s3path = PropertyMock(return_value="s3://testbucket/")
client_meta = {
"cloud_client": client,
"stage_info": {
"location": "sfc-teststage/rwyitestacco/users/1234/",
"locationType": "S3",
},
}
meta = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"client_meta": SFResourceMeta(**client_meta),
}
meta = SnowflakeFileMeta(**meta)
filename = "/path1/file2.txt"
akey = SnowflakeS3Util.get_file_header(meta, filename)
assert akey is None
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_download_expiry_error(caplog):
"""Tests whether token expiry error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
mock_resource = MagicMock()
mock_resource.download_file.side_effect = botocore.exceptions.ClientError(
{"Error": {
"Code": "ExpiredToken",
"Message": "Just testing"
}}, "Testing")
client_meta = {
"cloud_client": mock_resource,
"stage_info": {
"location": "loc"
},
}
meta_dict = {
"name": "f",
"src_file_name": "f",
"stage_location_type": "S3",
"sha256_digest": "asd",
"client_meta": SFResourceMeta(**client_meta),
"src_file_size": 99,
"get_callback_output_stream": None,
"show_progress_bar": False,
"get_callback": None,
}
meta = SnowflakeFileMeta(**meta_dict)
with mock.patch(
"snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object",
return_value=mock_resource,
):
SnowflakeS3Util._native_download_file(meta, "f", 4)
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_put_with_invalid_token(tmpdir, conn_cnx):
"""[s3] SNOW-6154: Uses invalid combination of AWS credential."""
# create a data file
fname = str(tmpdir.join("test_put_get_with_aws_token.txt.gz"))
with gzip.open(fname, "wb") as f:
f.write("123,test1\n456,test2".encode(UTF8))
table_name = random_string(5, "snow6154_")
with conn_cnx() as cnx:
try:
cnx.cursor().execute(
f"create or replace table {table_name} (a int, b string)")
ret = cnx.cursor()._execute_helper(
f"put file://{fname} @%{table_name}")
stage_info = ret["data"]["stageInfo"]
stage_info["location"]
stage_credentials = stage_info["creds"]
creds = StorageCredential(stage_credentials, cnx,
"COMMAND WILL NOT BE USED")
statinfo = os.stat(fname)
meta = SnowflakeFileMeta(
name=os.path.basename(fname),
src_file_name=fname,
src_file_size=statinfo.st_size,
stage_location_type="S3",
encryption_material=None,
dst_file_name=os.path.basename(fname),
sha256_digest="None",
)
client = SnowflakeS3RestClient(meta, creds, stage_info, 8388608)
client.get_file_header(meta.name) # positive case
# negative case, no aws token
token = stage_info["creds"]["AWS_TOKEN"]
del stage_info["creds"]["AWS_TOKEN"]
with pytest.raises(requests.HTTPError,
match=".*Forbidden for url.*"):
client.get_file_header(meta.name)
# negative case, wrong location
stage_info["creds"]["AWS_TOKEN"] = token
s3path = client.s3location.path
bad_path = os.path.dirname(os.path.dirname(s3path)) + "/"
_s3location = S3Location(client.s3location.bucket_name, bad_path)
client.s3location = _s3location
client.chunks = [b"this is a chunk"]
client.num_of_chunks = 1
client.retry_count[0] = 0
client.data_file = fname
with pytest.raises(requests.HTTPError,
match=".*Forbidden for url.*"):
client.upload_chunk(0)
finally:
cnx.cursor().execute(f"drop table if exists {table_name}")
def test_get_file_header_none_with_presigned_url(tmp_path):
"""Tests whether default file handle created by get_file_header is as expected."""
meta = SnowflakeFileMeta(
name=str(tmp_path / 'some_file'),
src_file_name=str(tmp_path / 'some_file'),
stage_location_type='GCS',
presigned_url='www.example.com',
)
file_header = SnowflakeGCSUtil.get_file_header(meta, 'file')
assert file_header.digest is None
assert file_header.content_length is None
assert file_header.encryption_metadata is None
def test_upload_retry_errors(errno, tmpdir):
"""Tests whether retryable errors are handled correctly when upploading."""
f_name = str(tmpdir.join("some_file.txt"))
resp = requests.Response()
resp.status_code = errno
meta = SnowflakeFileMeta(
name=f_name,
src_file_name=f_name,
stage_location_type="GCS",
presigned_url="some_url",
sha256_digest="asd",
)
if RequestExceedMaxRetryError is not None:
mock_connection = mock.create_autospec(SnowflakeConnection)
client = SnowflakeGCSRestClient(
meta,
StorageCredential({}, mock_connection, ""),
{},
mock_connection,
"",
)
with open(f_name, "w") as f:
f.write(random_string(15))
if RequestExceedMaxRetryError is None:
with mock.patch(
"snowflake.connector.vendored.requests.put"
if vendored_request
else "requests.put",
side_effect=requests.exceptions.HTTPError(response=resp),
):
SnowflakeGCSUtil.upload_file(f_name, meta, None, 99, 64000)
assert isinstance(meta.last_error, requests.exceptions.HTTPError)
assert meta.result_status == ResultStatus.NEED_RETRY
else:
client.data_file = f_name
if vendored_request:
with mock.patch.dict(
METHODS,
{"PUT": lambda *a, **kw: resp},
):
with pytest.raises(RequestExceedMaxRetryError):
# Retry quickly during unit tests
client.SLEEP_UNIT = 0.0
client.upload_chunk(0)
else:
# Old Driver test specific code
with mock.patch("requests.put"):
SnowflakeGCSUtil.upload_file(f_name, meta, None, 99, 64000)
assert isinstance(meta.last_error, requests.exceptions.HTTPError)
assert meta.result_status == ResultStatus.NEED_RETRY
def test_upload_expiry_error():
"""Tests whether token expiry error is handled as expected when uploading."""
meta_info = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
meta = SnowflakeFileMeta(**meta_info)
creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": "", "AWS_TOKEN": ""}
rest_client = SnowflakeS3RestClient(
meta,
StorageCredential(
creds,
MagicMock(autospec=SnowflakeConnection),
"PUT file:/tmp/file.txt @~",
),
{
"locationType": "AWS",
"location": "bucket/path",
"creds": creds,
"region": "test",
"endPoint": None,
},
8 * megabyte,
)
resp = MagicMock(
autospec=Response,
status_code=400,
text=f"<Error><Code>{EXPIRED_TOKEN}</Code></Error>",
)
from snowflake.connector.storage_client import METHODS
with mock.patch.dict(METHODS, PUT=MagicMock(return_value=resp)):
exc = Exception("stop execution")
with mock.patch.object(rest_client.credentials,
"update",
side_effect=exc):
with mock.patch(
"snowflake.connector.storage_client.SnowflakeStorageClient.preprocess"
):
rest_client.prepare_upload()
with pytest.raises(Exception) as caught_exc:
rest_client.upload_chunk(0)
assert caught_exc.value is exc
def test_download_uncaught_exception(tmp_path):
"""Tests whether non-retryable errors are handled correctly when downloading."""
resp = requests.Response()
resp.status_code = 501
meta = SnowflakeFileMeta(
name=str(tmp_path / 'some_file'),
src_file_name=str(tmp_path / 'some_file'),
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with mock.patch('snowflake.connector.vendored.requests.get'
if vendored_request else 'requests.get',
side_effect=requests.exceptions.HTTPError(response=resp)):
with pytest.raises(requests.exceptions.HTTPError):
SnowflakeGCSUtil._native_download_file(meta, str(tmp_path), 99)
def test_download_retry_errors(errno, tmp_path):
"""Tests whether retryable errors are handled correctly when downloading."""
resp = requests.Response()
resp.status_code = errno
meta = SnowflakeFileMeta(
name=str(tmp_path / 'some_file'),
src_file_name=str(tmp_path / 'some_file'),
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with mock.patch('snowflake.connector.vendored.requests.get'
if vendored_request else 'requests.get',
side_effect=requests.exceptions.HTTPError(response=resp)):
SnowflakeGCSUtil._native_download_file(meta, str(tmp_path), 99)
assert isinstance(meta.last_error, requests.exceptions.HTTPError)
assert meta.result_status == ResultStatus.NEED_RETRY
def test_download_retry_errors(errno, tmp_path):
"""Tests whether retryable errors are handled correctly when downloading."""
resp = requests.Response()
resp.status_code = errno
if errno == 403:
pytest.skip("This behavior has changed in the move from SDKs")
meta_info = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
meta = SnowflakeFileMeta(**meta_info)
creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": ""}
cnx = mock.MagicMock(autospec=SnowflakeConnection)
rest_client = SnowflakeGCSRestClient(
meta,
StorageCredential(
creds,
cnx,
"GET file:/tmp/file.txt @~",
),
{
"locationType": "AWS",
"location": "bucket/path",
"creds": creds,
"region": "test",
"endPoint": None,
},
cnx,
"GET file:///tmp/file.txt @~",
)
from snowflake.connector.storage_client import METHODS
rest_client.SLEEP_UNIT = 0
with mock.patch.dict(METHODS, GET=mock.MagicMock(return_value=resp)):
with pytest.raises(
RequestExceedMaxRetryError,
match="GET with url .* failed for exceeding maximum retries",
):
rest_client.download_chunk(0)
def test_get_file_header_none_with_presigned_url(tmp_path):
"""Tests whether default file handle created by get_file_header is as expected."""
meta = SnowflakeFileMeta(
name=str(tmp_path / "some_file"),
src_file_name=str(tmp_path / "some_file"),
stage_location_type="GCS",
presigned_url="www.example.com",
)
storage_credentials = Mock()
storage_credentials.creds = {}
stage_info = Mock()
connection = Mock()
client = SnowflakeGCSRestClient(
meta, storage_credentials, stage_info, connection, ""
)
file_header = client.get_file_header(meta.name)
assert file_header is None
def test_upload_one_file_to_s3_econnreset():
"""Tests Upload one file to S3 with retry on errno.ECONNRESET.
Notes:
The last attempted max_currency should not be changed.
"""
for error_code in [errno.ECONNRESET,
errno.ETIMEDOUT,
errno.EPIPE,
-1]:
upload_file = MagicMock(
side_effect=OpenSSL.SSL.SysCallError(
error_code, 'mock err. connection aborted'))
s3object = MagicMock(metadata=defaultdict(str), upload_file=upload_file)
client = Mock()
client.Object.return_value = s3object
initial_parallel = 100
client_meta = {
'stage_info': {
'location': 'sfc-teststage/rwyitestacco/users/1234/',
'locationType': 'S3',
},
'cloud_client': client,
}
upload_meta = {
'name': 'data1.txt.gz',
'stage_location_type': 'S3',
'no_sleeping_time': True,
'parallel': initial_parallel,
'put_callback': None,
'put_callback_output_stream': None,
SHA256_DIGEST: '123456789abcdef',
'client_meta': SFResourceMeta(**client_meta),
'dst_file_name': 'data1.txt.gz',
'src_file_name': path.join(THIS_DIR, '../data', 'put_get_1.txt'),
'overwrite': True,
}
upload_meta['real_src_file_name'] = upload_meta['src_file_name']
upload_meta['upload_size'] = os.stat(upload_meta['src_file_name']).st_size
meta = SnowflakeFileMeta(**upload_meta)
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert 'last_max_concurrency' not in upload_meta
def test_download_retry_exceeded_error():
"""Tests whether a retry exceeded error is handled as expected when downloading."""
meta_info = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": "path/to/put_get_1.txt",
"overwrite": True,
}
meta = SnowflakeFileMeta(**meta_info)
creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": "", "AWS_TOKEN": ""}
rest_client = SnowflakeS3RestClient(
meta,
StorageCredential(
creds,
MagicMock(autospec=SnowflakeConnection),
"GET file:/tmp/file.txt @~",
),
{
"locationType": "AWS",
"location": "bucket/path",
"creds": creds,
"region": "test",
"endPoint": None,
},
8 * megabyte,
)
rest_client.SLEEP_UNIT = 0
resp = Response()
resp.status_code = 500 # Use a transient error code
from snowflake.connector.storage_client import METHODS
with mock.patch.dict(METHODS, GET=MagicMock(return_value=resp)):
with mock.patch.object(rest_client.credentials, "update"):
with pytest.raises(
RequestExceedMaxRetryError,
match=
r"GET with url .* failed for exceeding maximum retries",
):
rest_client.download_chunk(0)
def test_download_unknown_error(caplog):
"""Tests whether an unknown error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
mock_resource = MagicMock()
mock_resource.download_file.side_effect = botocore.exceptions.ClientError(
{"Error": {
"Code": "unknown",
"Message": "Just testing"
}}, "Testing")
client_meta = {
"cloud_client": mock_resource,
"stage_info": {
"location": "loc"
},
}
meta = {
"name": "f",
"src_file_name": "f",
"stage_location_type": "S3",
"client_meta": SFResourceMeta(**client_meta),
"sha256_digest": "asd",
"src_file_size": 99,
"get_callback_output_stream": None,
"show_progress_bar": False,
"get_callback": None,
}
meta = SnowflakeFileMeta(**meta)
with mock.patch(
"snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object",
return_value=mock_resource,
):
with pytest.raises(
botocore.exceptions.ClientError,
match=
r"An error occurred \(unknown\) when calling the Testing operation: Just testing",
):
SnowflakeS3Util._native_download_file(meta, "f", 4)
assert (
"snowflake.connector.s3_util",
logging.DEBUG,
"Failed to download a file: f, err: An error occurred (unknown) when "
"calling the Testing operation: Just testing",
) in caplog.record_tuples
def test_upload_get_timeout(tmp_path, caplog):
"""Tests whether timeout error is handled correctly when downloading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
resp = requests.Response()
meta = SnowflakeFileMeta(
name=str(tmp_path / 'some_file'),
src_file_name=str(tmp_path / 'some_file'),
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with mock.patch('snowflake.connector.vendored.requests.get'
if vendored_request else 'requests.get',
side_effect=requests.exceptions.Timeout(response=resp)):
SnowflakeGCSUtil._native_download_file(meta, str(tmp_path), 99)
assert isinstance(meta.last_error, requests.exceptions.Timeout)
assert meta.result_status == ResultStatus.NEED_RETRY
assert ('snowflake.connector.gcs_util', logging.DEBUG,
'GCS file download Timeout Error: ') in caplog.record_tuples
def test_upload_uncaught_exception(tmpdir):
"""Tests whether non-retryable errors are handled correctly when uploading."""
f_name = str(tmpdir.join('some_file.txt'))
resp = requests.Response()
resp.status_code = 501
meta = SnowflakeFileMeta(
name=f_name,
src_file_name=f_name,
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with open(f_name, 'w') as f:
f.write(random_string(15))
with mock.patch('snowflake.connector.vendored.requests.put'
if vendored_request else 'requests.put',
side_effect=requests.exceptions.HTTPError(response=resp)):
with pytest.raises(requests.exceptions.HTTPError):
SnowflakeGCSUtil.upload_file(f_name, meta, None, 99, 64000)
def test_download_timeout(tmp_path, caplog):
"""Tests whether timeout error is handled correctly when downloading."""
timeout_exc = requests.exceptions.Timeout(response=requests.Response())
meta_info = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
meta = SnowflakeFileMeta(**meta_info)
creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": ""}
cnx = mock.MagicMock(autospec=SnowflakeConnection)
rest_client = SnowflakeGCSRestClient(
meta,
StorageCredential(
creds,
cnx,
"GET file:/tmp/file.txt @~",
),
{
"locationType": "AWS",
"location": "bucket/path",
"creds": creds,
"region": "test",
"endPoint": None,
},
cnx,
"GET file:///tmp/file.txt @~",
)
from snowflake.connector.storage_client import METHODS
rest_client.SLEEP_UNIT = 0
with mock.patch.dict(METHODS, GET=mock.MagicMock(side_effect=timeout_exc)):
exc = Exception("stop execution")
with mock.patch.object(rest_client.credentials, "update", side_effect=exc):
with pytest.raises(RequestExceedMaxRetryError):
rest_client.download_chunk(0)
def test_upload_one_file_to_s3_econnreset():
"""Tests Upload one file to S3 with retry on errno.ECONNRESET.
Notes:
The last attempted max_currency should not be changed.
"""
for error_code in [errno.ECONNRESET, errno.ETIMEDOUT, errno.EPIPE, -1]:
upload_file = MagicMock(side_effect=OpenSSL.SSL.SysCallError(
error_code, "mock err. connection aborted"))
s3object = MagicMock(metadata=defaultdict(str),
upload_file=upload_file)
client = Mock()
client.Object.return_value = s3object
initial_parallel = 100
client_meta = {
"stage_info": {
"location": "sfc-teststage/rwyitestacco/users/1234/",
"locationType": "S3",
},
"cloud_client": client,
}
upload_meta = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"parallel": initial_parallel,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"client_meta": SFResourceMeta(**client_meta),
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
upload_meta["real_src_file_name"] = upload_meta["src_file_name"]
upload_meta["upload_size"] = os.stat(
upload_meta["src_file_name"]).st_size
meta = SnowflakeFileMeta(**upload_meta)
with pytest.raises(OpenSSL.SSL.SysCallError):
SnowflakeRemoteStorageUtil.upload_one_file(meta)
assert upload_file.call_count == DEFAULT_MAX_RETRY
assert "last_max_concurrency" not in upload_meta
def test_upload_retry_errors(errno, tmpdir):
"""Tests whether retryable errors are handled correctly when upploading."""
f_name = str(tmpdir.join('some_file.txt'))
resp = requests.Response()
resp.status_code = errno
meta = SnowflakeFileMeta(
name=f_name,
src_file_name=f_name,
stage_location_type='GCS',
presigned_url='some_url',
sha256_digest='asd',
)
with open(f_name, 'w') as f:
f.write(random_string(15))
with mock.patch('snowflake.connector.vendored.requests.put'
if vendored_request else 'requests.put',
side_effect=requests.exceptions.HTTPError(response=resp)):
SnowflakeGCSUtil.upload_file(f_name, meta, None, 99, 64000)
assert isinstance(meta.last_error, requests.exceptions.HTTPError)
assert meta.result_status == ResultStatus.NEED_RETRY
def test_get_header_unknown_error(caplog):
"""Tests whether unexpected errors are handled as expected when getting header."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
meta_info = {
"name": "data1.txt.gz",
"stage_location_type": "S3",
"no_sleeping_time": True,
"put_callback": None,
"put_callback_output_stream": None,
SHA256_DIGEST: "123456789abcdef",
"dst_file_name": "data1.txt.gz",
"src_file_name": path.join(THIS_DIR, "../data", "put_get_1.txt"),
"overwrite": True,
}
meta = SnowflakeFileMeta(**meta_info)
creds = {"AWS_SECRET_KEY": "", "AWS_KEY_ID": "", "AWS_TOKEN": ""}
rest_client = SnowflakeS3RestClient(
meta,
StorageCredential(
creds,
MagicMock(autospec=SnowflakeConnection),
"PUT file:/tmp/file.txt @~",
),
{
"locationType": "AWS",
"location": "bucket/path",
"creds": creds,
"region": "test",
"endPoint": None,
},
8 * megabyte,
)
resp = Response()
# dont' use transient error codes
resp.status_code = 555
from snowflake.connector.storage_client import METHODS
with mock.patch.dict(METHODS, HEAD=MagicMock(return_value=resp)):
with pytest.raises(HTTPError, match="555 Server Error"):
rest_client.get_file_header("file.txt")
def test_upload_expiry_error(caplog):
"""Tests whether token expiry error is handled as expected when uploading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
mock_resource, mock_object = MagicMock(), MagicMock()
mock_resource.Object.return_value = mock_object
mock_object.upload_file.side_effect = botocore.exceptions.ClientError(
{'Error': {'Code': 'ExpiredToken', 'Message': 'Just testing'}}, 'Testing')
client_meta = {
'cloud_client': mock_resource,
'stage_info': {'location': 'loc'},
}
meta = {'name': 'f',
'src_file_name': 'f',
'stage_location_type': 'S3',
'client_meta': SFResourceMeta(**client_meta),
'sha256_digest': 'asd',
'dst_file_name': 'f',
'put_callback': None}
meta = SnowflakeFileMeta(**meta)
with mock.patch('snowflake.connector.s3_util.SnowflakeS3Util.extract_bucket_name_and_path'):
assert SnowflakeS3Util.upload_file('f', meta, None, 4, 67108864) is None
assert ('snowflake.connector.s3_util', logging.DEBUG, 'AWS Token expired. Renew and retry') in caplog.record_tuples
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_download_expiry_error(caplog):
"""Tests whether token expiry error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
mock_resource = MagicMock()
mock_resource.download_file.side_effect = botocore.exceptions.ClientError(
{'Error': {'Code': 'ExpiredToken', 'Message': 'Just testing'}}, 'Testing')
client_meta = {
'cloud_client': mock_resource,
'stage_info': {'location': 'loc'},
}
meta_dict = {'name': 'f',
'src_file_name': 'f',
'stage_location_type': 'S3',
'sha256_digest': 'asd',
'client_meta': SFResourceMeta(**client_meta),
'src_file_size': 99,
'get_callback_output_stream': None,
'show_progress_bar': False,
'get_callback': None}
meta = SnowflakeFileMeta(**meta_dict)
with mock.patch('snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object', return_value=mock_resource):
SnowflakeS3Util._native_download_file(meta, 'f', 4)
assert meta.result_status == ResultStatus.RENEW_TOKEN
def test_download_syscall_error(caplog, error_no, result_status):
"""Tests whether a syscall error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, 'snowflake.connector')
mock_resource = MagicMock()
mock_resource.download_file.side_effect = OpenSSL.SSL.SysCallError(error_no)
client_meta = {
'cloud_client': mock_resource,
'stage_info': {'location': 'loc'},
}
meta = {'name': 'f',
'stage_location_type': 'S3',
'client_meta': SFResourceMeta(**client_meta),
'sha256_digest': 'asd',
'src_file_name': 'f',
'src_file_size': 99,
'get_callback_output_stream': None,
'show_progress_bar': False,
'get_callback': None}
meta = SnowflakeFileMeta(**meta)
with mock.patch('snowflake.connector.s3_util.SnowflakeS3Util._get_s3_object', return_value=mock_resource):
SnowflakeS3Util._native_download_file(meta, 'f', 4)
assert meta.last_error is mock_resource.download_file.side_effect
assert meta.result_status == result_status
