From ad83ffdf1f4d69ffb62b243429e7b59d0930524c Mon Sep 17 00:00:00 2001 From: Victor Torres Date: Wed, 6 Feb 2019 18:32:46 -0200 Subject: [PATCH] refactoring --- scrapy/extensions/feedexport.py | 19 ++++++-- tests/test_feedexport.py | 84 +++++++++++++++++++++++++++++++++ 2 files changed, 98 insertions(+), 5 deletions(-) diff --git a/scrapy/extensions/feedexport.py b/scrapy/extensions/feedexport.py index eb0802261..ca30322be 100644 --- a/scrapy/extensions/feedexport.py +++ b/scrapy/extensions/feedexport.py @@ -93,7 +93,7 @@ class FileFeedStorage(object): class S3FeedStorage(BlockingFeedStorage): - def __init__(self, uri, access_key=None, secret_key=None): + def __init__(self, uri, access_key=None, secret_key=None, acl=None): # BEGIN Backwards compatibility for initialising without keys (and # without using from_crawler) no_defaults = access_key is None and secret_key is None @@ -118,7 +118,7 @@ class S3FeedStorage(BlockingFeedStorage): self.secret_key = u.password or secret_key self.is_botocore = is_botocore() self.keyname = u.path[1:] # remove first "/" - self.policy = settings.get('FEED_STORAGE_S3_ACL', 'private') + self.acl = acl if self.is_botocore: import botocore.session session = botocore.session.get_session() @@ -132,19 +132,28 @@ class S3FeedStorage(BlockingFeedStorage): @classmethod def from_crawler(cls, crawler, uri): return cls(uri, crawler.settings['AWS_ACCESS_KEY_ID'], - crawler.settings['AWS_SECRET_ACCESS_KEY']) + crawler.settings['AWS_SECRET_ACCESS_KEY'], + crawler.settings.get('FEED_STORAGE_S3_ACL')) def _store_in_thread(self, file): file.seek(0) if self.is_botocore: + kwargs = dict() + if self.acl: + kwargs.update(dict(ACL=self.acl)) + self.s3_client.put_object( Bucket=self.bucketname, Key=self.keyname, Body=file, - ACL=self.policy) + **kwargs) else: conn = self.connect_s3(self.access_key, self.secret_key) bucket = conn.get_bucket(self.bucketname, validate=False) key = bucket.new_key(self.keyname) - key.set_contents_from_file(file, policy=self.policy) + kwargs = dict() + if self.acl: + kwargs.update(dict(policy=self.acl)) + + key.set_contents_from_file(file, **kwargs) key.close() diff --git a/tests/test_feedexport.py b/tests/test_feedexport.py index e46c8c14e..b07635cb0 100644 --- a/tests/test_feedexport.py +++ b/tests/test_feedexport.py @@ -18,6 +18,7 @@ from tests import mock from tests.mockserver import MockServer from w3lib.url import path_to_file_uri +import botocore.client import scrapy from scrapy.exporters import CsvItemExporter from scrapy.extensions.feedexport import ( @@ -186,6 +187,89 @@ class S3FeedStorageTest(unittest.TestCase): content = get_s3_content_and_delete(u.hostname, u.path[1:]) self.assertEqual(content, expected_content) + def test_init_without_acl(self): + storage = S3FeedStorage( + 's3://mybucket/export.csv', + 'access_key', + 'secret_key' + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, None) + + def test_init_with_acl(self): + storage = S3FeedStorage( + 's3://mybucket/export.csv', + 'access_key', + 'secret_key', + 'custom-acl' + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, 'custom-acl') + + def test_from_crawler_without_acl(self): + settings = { + 'AWS_ACCESS_KEY_ID': 'access_key', + 'AWS_SECRET_ACCESS_KEY': 'secret_key', + } + crawler = get_crawler(settings_dict=settings) + storage = S3FeedStorage.from_crawler( + crawler, + 's3://mybucket/export.csv' + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, None) + + def test_from_crawler_with_acl(self): + settings = { + 'AWS_ACCESS_KEY_ID': 'access_key', + 'AWS_SECRET_ACCESS_KEY': 'secret_key', + 'FEED_STORAGE_S3_ACL': 'custom-acl', + } + crawler = get_crawler(settings_dict=settings) + storage = S3FeedStorage.from_crawler( + crawler, + 's3://mybucket/export.csv' + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, 'custom-acl') + + def test_store_in_thread_without_acl(self): + storage = S3FeedStorage( + 's3://mybucket/export.csv', + 'access_key', + 'secret_key', + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, None) + + with mock.patch('botocore.client.BaseClient._make_api_call') as _make_api_call_mock: + storage._store_in_thread(BytesIO(b'test file')) + operation_name, api_params = _make_api_call_mock.call_args[0] + self.assertEqual(operation_name, 'PutObject') + self.assertNotIn('ACL', api_params) + + def test_store_in_thread_with_acl(self): + storage = S3FeedStorage( + 's3://mybucket/export.csv', + 'access_key', + 'secret_key', + 'custom-acl' + ) + self.assertEqual(storage.access_key, 'access_key') + self.assertEqual(storage.secret_key, 'secret_key') + self.assertEqual(storage.acl, 'custom-acl') + + with mock.patch('botocore.client.BaseClient._make_api_call') as _make_api_call_mock: + storage._store_in_thread(BytesIO(b'test file')) + operation_name, api_params = _make_api_call_mock.call_args[0] + self.assertEqual(operation_name, 'PutObject') + self.assertEqual(api_params.get('ACL'), 'custom-acl') + class StdoutFeedStorageTest(unittest.TestCase):