Merge pull request #2616 from redapple/mediapipeline-redirect-fix-continued

[MRG] Allow redirections in media files downloads

docs/topics/media-pipeline.rst

@@ -322,6 +322,19 @@ By default, there are no size constraints, so all images are processed.
 
 .. _topics-media-pipeline-override:
 
+Allowing redirections
+---------------------
+
+.. setting:: MEDIA_ALLOW_REDIRECTS
+
+By default media pipelines ignore redirects, i.e. an HTTP redirection
+to a media file URL request will mean the media download is considered failed.
+
+To handle media redirections, set this settings to ``True``:
+
+    MEDIA_ALLOW_REDIRECTS = True
+
+
 Extending the Media Pipelines
 =============================
 

scrapy/pipelines/files.py

@@ -249,7 +249,7 @@ class FilesPipeline(MediaPipeline):
             resolve('FILES_RESULT_FIELD'), self.FILES_RESULT_FIELD
         )
 
-        super(FilesPipeline, self).__init__(download_func=download_func)
+        super(FilesPipeline, self).__init__(download_func=download_func, settings=settings)
 
     @classmethod
     def from_settings(cls, settings):

scrapy/pipelines/media.py

@@ -1,10 +1,13 @@
 from __future__ import print_function
 
+import functools
 import logging
 from collections import defaultdict
 from twisted.internet.defer import Deferred, DeferredList
 from twisted.python.failure import Failure
 
+from scrapy.settings import Settings
+from scrapy.utils.datatypes import SequenceExclude
 from scrapy.utils.defer import mustbe_deferred, defer_result
 from scrapy.utils.request import request_fingerprint
 from scrapy.utils.misc import arg_to_iter
@@ -24,9 +27,23 @@ class MediaPipeline(object):
             self.downloaded = {}
             self.waiting = defaultdict(list)
 
-    def __init__(self, download_func=None):
+    def __init__(self, download_func=None, settings=None):
         self.download_func = download_func
 
+        if isinstance(settings, dict) or settings is None:
+            settings = Settings(settings)
+        resolve = functools.partial(self._key_for_pipe,
+                                    base_class_name="MediaPipeline",
+                                    settings=settings)
+        self.allow_redirects = settings.getbool(
+            resolve('MEDIA_ALLOW_REDIRECTS'), False
+        )
+        self._handle_statuses(self.allow_redirects)
+
+    def _handle_statuses(self, allow_redirects):
+        self.handle_httpstatus_list = None
+        if allow_redirects:
+            self.handle_httpstatus_list = SequenceExclude(range(300, 400))
 
     def _key_for_pipe(self, key, base_class_name=None,
                       settings=None):
@@ -93,6 +110,12 @@ class MediaPipeline(object):
         )
         return dfd.addBoth(lambda _: wad)  # it must return wad at last
 
+    def _modify_media_request(self, request):
+        if self.handle_httpstatus_list:
+            request.meta['handle_httpstatus_list'] = self.handle_httpstatus_list
+        else:
+            request.meta['handle_httpstatus_all'] = True
+
     def _check_media_to_download(self, result, request, info):
         if result is not None:
             return result
@@ -103,7 +126,7 @@ class MediaPipeline(object):
                 callback=self.media_downloaded, callbackArgs=(request, info),
                 errback=self.media_failed, errbackArgs=(request, info))
         else:
-            request.meta['handle_httpstatus_all'] = True
+            self._modify_media_request(request)
             dfd = self.crawler.engine.download(request, info.spider)
             dfd.addCallbacks(
                 callback=self.media_downloaded, callbackArgs=(request, info),

tests/mockserver.py

@@ -5,13 +5,17 @@ from subprocess import Popen, PIPE
 
 from twisted.web.server import Site, NOT_DONE_YET
 from twisted.web.resource import Resource
+from twisted.web.static import File
 from twisted.web.test.test_webclient import PayloadResource
 from twisted.web.server import GzipEncoderFactory
 from twisted.web.resource import EncodingResourceWrapper
+from twisted.web.util import redirectTo
 from twisted.internet import reactor, ssl
 from twisted.internet.task import deferLater
 
+
 from scrapy.utils.python import to_bytes, to_unicode
+from tests import tests_datadir
 
 
 def getarg(request, name, default=None, type=None):
@@ -120,6 +124,16 @@ class Echo(LeafResource):
         return to_bytes(json.dumps(output))
 
 
+class RedirectTo(LeafResource):
+
+    def render(self, request):
+        goto = getarg(request, b'goto', b'/')
+        # we force the body content, otherwise Twisted redirectTo()
+        # returns HTML with <meta http-equiv="refresh"
+        redirectTo(goto, request)
+        return b'redirecting...'
+
+
 class Partial(LeafResource):
 
     def render_GET(self, request):
@@ -160,6 +174,8 @@ class Root(Resource):
         self.putChild(b"echo", Echo())
         self.putChild(b"payload", PayloadResource())
         self.putChild(b"xpayload", EncodingResourceWrapper(PayloadResource(), [GzipEncoderFactory()]))
+        self.putChild(b"files", File(os.path.join(tests_datadir, 'test_site/files/')))
+        self.putChild(b"redirect-to", RedirectTo())
 
     def getChild(self, name, request):
         return self

tests/test_pipeline_crawl.py

@@ -0,0 +1,182 @@
+# -*- coding: utf-8 -*-
+import os
+import shutil
+
+from testfixtures import LogCapture
+from twisted.internet import defer
+from twisted.trial.unittest import TestCase
+from w3lib.url import add_or_replace_parameter
+
+from scrapy.crawler import CrawlerRunner
+from scrapy import signals
+from tests.mockserver import MockServer
+from tests.spiders import SimpleSpider
+
+
+class MediaDownloadSpider(SimpleSpider):
+    name = 'mediadownload'
+
+    def _process_url(self, url):
+        return url
+
+    def parse(self, response):
+        self.logger.info(response.headers)
+        self.logger.info(response.text)
+        item = {
+            self.media_key: [],
+            self.media_urls_key: [
+                self._process_url(response.urljoin(href))
+                    for href in response.xpath('''
+                        //table[thead/tr/th="Filename"]
+                            /tbody//a/@href
+                        ''').extract()],
+        }
+        yield item
+
+
+class BrokenLinksMediaDownloadSpider(MediaDownloadSpider):
+    name = 'brokenmedia'
+
+    def _process_url(self, url):
+        return url + '.foo'
+
+
+class RedirectedMediaDownloadSpider(MediaDownloadSpider):
+    name = 'redirectedmedia'
+
+    def _process_url(self, url):
+        return add_or_replace_parameter(
+                    'http://localhost:8998/redirect-to',
+                    'goto', url)
+
+
+class FileDownloadCrawlTestCase(TestCase):
+    pipeline_class = 'scrapy.pipelines.files.FilesPipeline'
+    store_setting_key = 'FILES_STORE'
+    media_key = 'files'
+    media_urls_key = 'file_urls'
+    expected_checksums = set([
+        '5547178b89448faf0015a13f904c936e',
+        'c2281c83670e31d8aaab7cb642b824db',
+        'ed3f6538dc15d4d9179dae57319edc5f'])
+
+    def setUp(self):
+        self.mockserver = MockServer()
+        self.mockserver.__enter__()
+
+        # prepare a directory for storing files
+        self.tmpmediastore = self.mktemp()
+        os.mkdir(self.tmpmediastore)
+        self.settings = {
+            'ITEM_PIPELINES': {self.pipeline_class: 1},
+            self.store_setting_key: self.tmpmediastore,
+        }
+        self.runner = CrawlerRunner(self.settings)
+        self.items = []
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpmediastore)
+        self.items = []
+        self.mockserver.__exit__(None, None, None)
+
+    def _on_item_scraped(self, item):
+        self.items.append(item)
+
+    def _create_crawler(self, spider_class, **kwargs):
+        crawler = self.runner.create_crawler(spider_class, **kwargs)
+        crawler.signals.connect(self._on_item_scraped, signals.item_scraped)
+        return crawler
+
+    def _assert_files_downloaded(self, items, logs):
+        self.assertEqual(len(items), 1)
+        self.assertIn(self.media_key, items[0])
+
+        # check that logs show the expected number of successful file downloads
+        file_dl_success = 'File (downloaded): Downloaded file from'
+        self.assertEqual(logs.count(file_dl_success), 3)
+
+        # check that the images/files checksums are what we know they should be
+        if self.expected_checksums is not None:
+            checksums = set(
+                i['checksum']
+                    for item in items
+                        for i in item[self.media_key])
+            self.assertEqual(checksums, self.expected_checksums)
+
+        # check that the image files where actually written to the media store
+        for item in items:
+            for i in item[self.media_key]:
+                self.assertTrue(
+                    os.path.exists(
+                        os.path.join(self.tmpmediastore, i['path'])))
+
+    def _assert_files_download_failure(self, crawler, items, code, logs):
+
+        # check that the item does NOT have the "images/files" field populated
+        self.assertEqual(len(items), 1)
+        self.assertIn(self.media_key, items[0])
+        self.assertFalse(items[0][self.media_key])
+
+        # check that there was 1 successful fetch and 3 other responses with non-200 code
+        self.assertEqual(crawler.stats.get_value('downloader/request_method_count/GET'), 4)
+        self.assertEqual(crawler.stats.get_value('downloader/response_count'), 4)
+        self.assertEqual(crawler.stats.get_value('downloader/response_status_count/200'), 1)
+        self.assertEqual(crawler.stats.get_value('downloader/response_status_count/%d' % code), 3)
+
+        # check that logs do show the failure on the file downloads
+        file_dl_failure = 'File (code: %d): Error downloading file from' % code
+        self.assertEqual(logs.count(file_dl_failure), 3)
+
+        # check that no files were written to the media store
+        self.assertEqual(os.listdir(self.tmpmediastore), [])
+
+    @defer.inlineCallbacks
+    def test_download_media(self):
+        crawler = self._create_crawler(MediaDownloadSpider)
+        with LogCapture() as log:
+            yield crawler.crawl("http://localhost:8998/files/images/",
+                media_key=self.media_key,
+                media_urls_key=self.media_urls_key)
+        self._assert_files_downloaded(self.items, str(log))
+
+    @defer.inlineCallbacks
+    def test_download_media_wrong_urls(self):
+        crawler = self._create_crawler(BrokenLinksMediaDownloadSpider)
+        with LogCapture() as log:
+            yield crawler.crawl("http://localhost:8998/files/images/",
+                media_key=self.media_key,
+                media_urls_key=self.media_urls_key)
+        self._assert_files_download_failure(crawler, self.items, 404, str(log))
+
+    @defer.inlineCallbacks
+    def test_download_media_redirected_default_failure(self):
+        crawler = self._create_crawler(RedirectedMediaDownloadSpider)
+        with LogCapture() as log:
+            yield crawler.crawl("http://localhost:8998/files/images/",
+                media_key=self.media_key,
+                media_urls_key=self.media_urls_key)
+        self._assert_files_download_failure(crawler, self.items, 302, str(log))
+
+    @defer.inlineCallbacks
+    def test_download_media_redirected_allowed(self):
+        settings = dict(self.settings)
+        settings.update({'MEDIA_ALLOW_REDIRECTS': True})
+        self.runner = CrawlerRunner(settings)
+
+        crawler = self._create_crawler(RedirectedMediaDownloadSpider)
+        with LogCapture() as log:
+            yield crawler.crawl("http://localhost:8998/files/images/",
+                media_key=self.media_key,
+                media_urls_key=self.media_urls_key)
+        self._assert_files_downloaded(self.items, str(log))
+        self.assertEqual(crawler.stats.get_value('downloader/response_status_count/302'), 3)
+
+
+class ImageDownloadCrawlTestCase(FileDownloadCrawlTestCase):
+    pipeline_class = 'scrapy.pipelines.images.ImagesPipeline'
+    store_setting_key = 'IMAGES_STORE'
+    media_key = 'images'
+    media_urls_key = 'image_urls'
+
+    # somehow checksums for images are different for Python 3.3
+    expected_checksums = None

tests/test_pipeline_media.py

@@ -6,6 +6,7 @@ from twisted.internet import reactor
 from twisted.internet.defer import Deferred, inlineCallbacks
 
 from scrapy.http import Request, Response
+from scrapy.settings import Settings
 from scrapy.spiders import Spider
 from scrapy.utils.request import request_fingerprint
 from scrapy.pipelines.media import MediaPipeline
@@ -22,10 +23,12 @@ def _mocked_download_func(request, info):
 class BaseMediaPipelineTestCase(unittest.TestCase):
 
     pipeline_class = MediaPipeline
+    settings = None
 
     def setUp(self):
         self.spider = Spider('media.com')
-        self.pipe = self.pipeline_class(download_func=_mocked_download_func)
+        self.pipe = self.pipeline_class(download_func=_mocked_download_func,
+                                        settings=Settings(self.settings))
         self.pipe.open_spider(self.spider)
         self.info = self.pipe.spiderinfo
 
@@ -82,6 +85,11 @@ class BaseMediaPipelineTestCase(unittest.TestCase):
         new_item = yield self.pipe.process_item(item, self.spider)
         assert new_item is item
 
+    def test_modify_media_request(self):
+        request = Request('http://url')
+        self.pipe._modify_media_request(request)
+        assert request.meta == {'handle_httpstatus_all': True}
+
 
 class MockedMediaPipeline(MediaPipeline):
 
@@ -249,3 +257,61 @@ class MediaPipelineTestCase(BaseMediaPipelineTestCase):
         self.assertEqual(new_item['results'], [(True, 'ITSME')])
         self.assertEqual(self.pipe._mockcalled, \
                 ['get_media_requests', 'media_to_download', 'item_completed'])
+
+
+class MediaPipelineAllowRedirectSettingsTestCase(unittest.TestCase):
+
+    def _assert_request_no3xx(self, pipeline_class, settings):
+        pipe = pipeline_class(settings=Settings(settings))
+        request = Request('http://url')
+        pipe._modify_media_request(request)
+
+        self.assertIn('handle_httpstatus_list', request.meta)
+        for status, check in [
+                (200, True),
+
+                # These are the status codes we want
+                # the downloader to handle itself
+                (301, False),
+                (302, False),
+                (302, False),
+                (307, False),
+                (308, False),
+
+                # we still want to get 4xx and 5xx
+                (400, True),
+                (404, True),
+                (500, True)]:
+            if check:
+                self.assertIn(status, request.meta['handle_httpstatus_list'])
+            else:
+                self.assertNotIn(status, request.meta['handle_httpstatus_list'])
+
+    def test_standard_setting(self):
+        self._assert_request_no3xx(
+            MediaPipeline,
+            {
+                'MEDIA_ALLOW_REDIRECTS': True
+            })
+
+    def test_subclass_standard_setting(self):
+
+        class UserDefinedPipeline(MediaPipeline):
+            pass
+
+        self._assert_request_no3xx(
+            UserDefinedPipeline,
+            {
+                'MEDIA_ALLOW_REDIRECTS': True
+            })
+
+    def test_subclass_specific_setting(self):
+
+        class UserDefinedPipeline(MediaPipeline):
+            pass
+
+        self._assert_request_no3xx(
+            UserDefinedPipeline,
+            {
+                'USERDEFINEDPIPELINE_MEDIA_ALLOW_REDIRECTS': True
+            })