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