def mocked_file_agent(*args, **kwargs):
agent = SnowflakeFileTransferAgent(*args, **kwargs)
agent.renew_expired_client = mock.MagicMock(
wraps=agent.renew_expired_client
)
mocked_file_agent.agent = agent
return agent
def test_get_empty_file(tmpdir):
"""Tests for error message when retrieving missing file."""
tmp_dir = str(tmpdir.mkdir("getfiledir"))
con = mock.MagicMock()
cursor = con.cursor()
cursor.errorhandler = Error.default_errorhandler
query = f"GET something file:\\{tmp_dir}"
ret = {
"data": {
"localLocation": tmp_dir,
"command": "DOWNLOAD",
"autoCompress": False,
"src_locations": [],
"sourceCompression": "none",
"stageInfo": {
"creds": {},
"location": "",
"locationType": "S3",
"path": "remote_loc",
},
},
"success": True,
}
sf_file_transfer_agent = SnowflakeFileTransferAgent(
cursor, query, ret, raise_put_get_error=True
)
with pytest.raises(OperationalError, match=".*the file does not exist.*$"):
sf_file_transfer_agent.execute()
assert not sf_file_transfer_agent.result()["rowset"]
def mocked_file_agent(*args, **kwargs):
agent = SnowflakeFileTransferAgent(*args, **kwargs)
agent._update_file_metas_with_presigned_url = mock.MagicMock(
wraps=agent._update_file_metas_with_presigned_url
)
mocked_file_agent.agent = agent
return agent
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
exc = requests.exceptions.HTTPError(response=resp)
with open(f_name, "w") as f:
f.write(random_string(15))
agent = SnowflakeFileTransferAgent(
mock.MagicMock(),
f"put {f_name} @~",
{
"data": {
"command": "UPLOAD",
"src_locations": [f_name],
"stageInfo": {
"locationType": "GCS",
"location": "",
"creds": {"AWS_SECRET_KEY": "", "AWS_KEY_ID": ""},
"region": "test",
"endPoint": None,
},
"localLocation": "/tmp",
}
},
)
with mock.patch(
"snowflake.connector.gcs_storage_client.SnowflakeGCSRestClient.get_file_header",
), mock.patch(
"snowflake.connector.gcs_storage_client.SnowflakeGCSRestClient._upload_chunk",
side_effect=exc,
):
agent.execute()
assert agent._file_metadata[0].error_details is exc
def test_upload_file_with_s3_upload_failed_error(tmp_path):
"""Tests Upload file with S3UploadFailedError, which could indicate AWS token expires."""
file1 = tmp_path / "file1"
with file1.open("w") as f:
f.write("test1")
rest_client = SnowflakeFileTransferAgent(
MagicMock(autospec=SnowflakeCursor),
"PUT some_file.txt",
{
"data": {
"command": "UPLOAD",
"autoCompress": False,
"src_locations": [file1],
"sourceCompression": "none",
"stageInfo": {
"creds": {
"AWS_SECRET_KEY": "secret key",
"AWS_KEY_ID": "secret id",
"AWS_TOKEN": "",
},
"location": "some_bucket",
"region": "no_region",
"locationType": "S3",
"path": "remote_loc",
"endPoint": "",
},
},
"success": True,
},
)
exc = Exception("Stop executing")
def mock_transfer_accelerate_config(
self: SnowflakeS3RestClient,
use_accelerate_endpoint: bool | None = None,
) -> bool:
self.endpoint = f"https://{self.s3location.bucket_name}.s3.awsamazon.com"
return False
with mock.patch(
"snowflake.connector.s3_storage_client.SnowflakeS3RestClient._has_expired_token",
return_value=True,
):
with mock.patch(
"snowflake.connector.s3_storage_client.SnowflakeS3RestClient.transfer_accelerate_config",
mock_transfer_accelerate_config,
):
with mock.patch(
"snowflake.connector.file_transfer_agent.StorageCredential.update",
side_effect=exc,
) as mock_update:
rest_client.execute()
assert mock_update.called
assert rest_client._results[0].error_details is exc
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()
with open(f_name, "w") as f:
f.write(random_string(15))
agent = SnowflakeFileTransferAgent(
mock.Mock(autospec=SnowflakeConnection, connection=None),
f"put {f_name} @~",
{
"data": {
"command": "UPLOAD",
"src_locations": [f_name],
"stageInfo": {
"locationType": "GCS",
"location": "",
"creds": {"AWS_SECRET_KEY": "", "AWS_KEY_ID": ""},
"region": "test",
"endPoint": None,
},
"localLocation": "/tmp",
}
},
)
mocked_put, mocked_head = mock.MagicMock(), mock.MagicMock()
mocked_put.side_effect = requests.exceptions.Timeout(response=resp)
resp = Response()
resp.status_code = 404
mocked_head.return_value = resp
SnowflakeGCSRestClient.SLEEP_UNIT = 0
from snowflake.connector.storage_client import METHODS
with mock.patch.dict(METHODS, {"PUT": mocked_put, "HEAD": mocked_head}):
agent.execute()
assert (
"snowflake.connector.storage_client",
logging.WARNING,
"PUT with url https://storage.googleapis.com//some_file.txt.gz failed for transient error: ",
) in caplog.record_tuples
assert (
"snowflake.connector.file_transfer_agent",
logging.DEBUG,
"Chunk 0 of file some_file.txt failed to transfer for unexpected exception PUT with url https://storage.googleapis.com//some_file.txt.gz failed for exceeding maximum retries.",
) in caplog.record_tuples
def test_download_unknown_error(caplog):
"""Tests whether an unknown error is handled as expected when downloading."""
caplog.set_level(logging.DEBUG, "snowflake.connector")
agent = SnowflakeFileTransferAgent(
MagicMock(),
"get @~/f /tmp",
{
"data": {
"command": "DOWNLOAD",
"src_locations": ["/tmp/a"],
"stageInfo": {
"locationType": "S3",
"location": "",
"creds": {
"AWS_SECRET_KEY": "",
"AWS_KEY_ID": "",
"AWS_TOKEN": ""
},
"region": "",
"endPoint": None,
},
"localLocation": "/tmp",
}
},
)
resp = Response()
resp.status_code = 400
resp.reason = "No, just chuck testing..."
with mock.patch(
"snowflake.connector.s3_storage_client.SnowflakeS3RestClient._send_request_with_authentication_and_retry",
return_value=resp,
), mock.patch(
"snowflake.connector.file_transfer_agent.SnowflakeFileTransferAgent._transfer_accelerate_config",
side_effect=None,
):
agent.execute()
assert (str(agent._file_metadata[0].error_details) ==
"400 Client Error: No, just chuck testing... for url: None")
assert verify_log_tuple(
"snowflake.connector.storage_client",
logging.ERROR,
re.compile("Failed to download a file: .*a"),
caplog.record_tuples,
)
def test_upload_file_with_azure_upload_failed_error(tmp_path):
"""Tests Upload file with expired Azure storage token."""
file1 = tmp_path / "file1"
with file1.open("w") as f:
f.write("test1")
rest_client = SnowflakeFileTransferAgent(
mock.MagicMock(autospec=SnowflakeCursor),
"PUT some_file.txt",
{
"data": {
"command": "UPLOAD",
"src_locations": [file1],
"sourceCompression": "none",
"stageInfo": {
"creds": {
"AZURE_SAS_TOKEN": "sas_token",
},
"location": "some_bucket",
"region": "no_region",
"locationType": "AZURE",
"path": "remote_loc",
"endPoint": "",
"storageAccount": "storage_account",
},
},
"success": True,
},
)
exc = Exception("Stop executing")
with mock.patch(
"snowflake.connector.azure_storage_client.SnowflakeAzureRestClient._has_expired_token",
return_value=True,
):
with mock.patch(
"snowflake.connector.file_transfer_agent.StorageCredential.update",
side_effect=exc,
) as mock_update:
rest_client.execute()
assert mock_update.called
assert rest_client._results[0].error_details is exc
def test_put_error(tmpdir):
"""Tests for raise_put_get_error flag (now turned on by default) in SnowflakeFileTransferAgent."""
tmp_dir = str(tmpdir.mkdir('putfiledir'))
file1 = path.join(tmp_dir, 'file1')
remote_location = path.join(tmp_dir, 'remote_loc')
with open(file1, 'w') as f:
f.write('test1')
# nobody can read now.
chmod(file1, 0o000)
con = MagicMock()
cursor = con.cursor()
cursor.errorhandler = Error.default_errorhandler
query = 'PUT something'
ret = {
'data': {
'command': 'UPLOAD',
'autoCompress': False,
'src_locations': [file1],
'sourceCompression': 'none',
'stageInfo': {
'location': remote_location,
'locationType': 'LOCAL_FS',
'path': 'remote_loc',
}
},
'success': True,
}
# no error is raised
sf_file_transfer_agent = SnowflakeFileTransferAgent(
cursor, query, ret, raise_put_get_error=False)
sf_file_transfer_agent.execute()
sf_file_transfer_agent.result()
# Permission error should be raised
sf_file_transfer_agent = SnowflakeFileTransferAgent(
cursor, query, ret, raise_put_get_error=True)
sf_file_transfer_agent.execute()
with pytest.raises(Exception):
sf_file_transfer_agent.result()
# unspecified, should fail because flag is on by default now
sf_file_transfer_agent = SnowflakeFileTransferAgent(cursor, query, ret)
sf_file_transfer_agent.execute()
with pytest.raises(Exception):
sf_file_transfer_agent.result()
chmod(file1, 0o700)
def mocked_file_agent(*args, **kwargs):
newkwargs = kwargs.copy()
newkwargs.update(multipart_threshold=10000)
agent = SnowflakeFileTransferAgent(*args, **newkwargs)
mocked_file_agent.agent = agent
return agent
def test_put_error(tmpdir):
"""Tests for raise_put_get_error flag (now turned on by default) in SnowflakeFileTransferAgent."""
tmp_dir = str(tmpdir.mkdir("putfiledir"))
file1 = path.join(tmp_dir, "file1")
remote_location = path.join(tmp_dir, "remote_loc")
with open(file1, "w") as f:
f.write("test1")
# nobody can read now.
chmod(file1, 0o000)
con = MagicMock()
cursor = con.cursor()
cursor.errorhandler = Error.default_errorhandler
query = "PUT something"
ret = {
"data": {
"command": "UPLOAD",
"autoCompress": False,
"src_locations": [file1],
"sourceCompression": "none",
"stageInfo": {
"location": remote_location,
"locationType": "LOCAL_FS",
"path": "remote_loc",
},
},
"success": True,
}
# no error is raised
sf_file_transfer_agent = SnowflakeFileTransferAgent(
cursor, query, ret, raise_put_get_error=False)
sf_file_transfer_agent.execute()
sf_file_transfer_agent.result()
# Permission error should be raised
sf_file_transfer_agent = SnowflakeFileTransferAgent(
cursor, query, ret, raise_put_get_error=True)
sf_file_transfer_agent.execute()
with pytest.raises(Exception):
sf_file_transfer_agent.result()
# unspecified, should fail because flag is on by default now
sf_file_transfer_agent = SnowflakeFileTransferAgent(cursor, query, ret)
sf_file_transfer_agent.execute()
with pytest.raises(Exception):
sf_file_transfer_agent.result()
chmod(file1, 0o700)
def mocked_file_agent(*args, **kwargs):
agent = SnowflakeFileTransferAgent(*args, **kwargs)
mocked_file_agent.agent = agent
return agent
