remote: http: raise exception when download response with error status code

This commit is contained in:
Paweł Redzyński
2019-11-29 15:30:55 +01:00
parent 510ff8a9da
commit 72abfb0572
5 changed files with 43 additions and 5 deletions
+8
View File
@@ -348,3 +348,11 @@ class NoOutputInExternalRepoError(DvcException):
relpath(path, external_repo_path), external_repo_url
)
)
class HTTPErrorStatusCodeException(DvcException):
def __init__(self, code, reason):
super(HTTPErrorStatusCodeException, self).__init__(
"Server responded with error status code: '{}' and message: "
"'{}'".format(code, reason)
)
+7 -2
View File
@@ -8,6 +8,7 @@ from funcy import cached_property, wrap_prop
from dvc.config import Config
from dvc.config import ConfigError
from dvc.exceptions import DvcException
from dvc.exceptions import HTTPErrorStatusCodeException
from dvc.progress import Tqdm
from dvc.remote.base import RemoteBASE
from dvc.scheme import Schemes
@@ -37,7 +38,11 @@ class RemoteHTTP(RemoteBASE):
)
def _download(self, from_info, to_file, name=None, no_progress_bar=False):
request = self._request("GET", from_info.url, stream=True)
response = self._request("GET", from_info.url, stream=True)
if response.status_code != 200:
raise HTTPErrorStatusCodeException(
response.status_code, response.reason
)
with Tqdm(
total=None if no_progress_bar else self._content_length(from_info),
leave=False,
@@ -46,7 +51,7 @@ class RemoteHTTP(RemoteBASE):
disable=no_progress_bar,
) as pbar:
with open(to_file, "wb") as fd:
for chunk in request.iter_content(chunk_size=self.CHUNK_SIZE):
for chunk in response.iter_content(chunk_size=self.CHUNK_SIZE):
fd.write(chunk)
fd.flush()
pbar.update(len(chunk))
+2 -1
View File
@@ -43,6 +43,7 @@ from tests.func.test_data_cloud import _should_test_ssh
from tests.func.test_data_cloud import get_ssh_url
from tests.func.test_data_cloud import TEST_AWS_REPO_BUCKET
from tests.func.test_data_cloud import TEST_GCP_REPO_BUCKET
from tests.utils.httpd import ContentMD5Handler
from tests.utils.httpd import StaticFileServer
@@ -1176,7 +1177,7 @@ class TestReproExternalHTTP(TestReproExternalBase):
self.dvc.remove("imported_file.dvc")
with StaticFileServer(handler="Content-MD5") as httpd:
with StaticFileServer(handler_class=ContentMD5Handler) as httpd:
import_url = urljoin(self.get_remote(httpd.server_port), self.FOO)
import_output = "imported_file"
import_stage = self.dvc.imp_url(import_url, import_output)
+24
View File
@@ -1,7 +1,11 @@
import pytest
from BaseHTTPServer import BaseHTTPRequestHandler
from dvc.config import ConfigError
from dvc.exceptions import HTTPErrorStatusCodeException
from dvc.path_info import URLInfo
from dvc.remote.http import RemoteHTTP
from tests.utils.httpd import StaticFileServer
def test_no_traverse_compatibility(dvc_repo):
@@ -13,3 +17,23 @@ def test_no_traverse_compatibility(dvc_repo):
with pytest.raises(ConfigError):
RemoteHTTP(dvc_repo, config)
@pytest.mark.parametrize("response_code", [404, 403, 500])
def test_download_fails_on_error_code(response_code, dvc_repo):
class ErrorStatusRequestHandler(BaseHTTPRequestHandler):
def do_GET(self):
self.send_response(response_code, message="Error message")
self.end_headers()
with StaticFileServer(ErrorStatusRequestHandler) as httpd:
url = "http://localhost:{}/".format(httpd.server_port)
config = {"url": url}
remote = RemoteHTTP(dvc_repo, config)
import os
with pytest.raises(HTTPErrorStatusCodeException):
remote._download(
URLInfo(os.path.join(url, "file.txt")), "file.txt"
)
+2 -2
View File
@@ -45,9 +45,9 @@ class ContentMD5Handler(TestRequestHandler):
class StaticFileServer:
_lock = threading.Lock()
def __init__(self, handler="etag"):
def __init__(self, handler_class=ETagHandler):
self._lock.acquire()
handler_class = ETagHandler if handler == "etag" else ContentMD5Handler
self.response_handler = handler_class
self._httpd = HTTPServer(("localhost", 0), handler_class)
self._thread = None