diff --git a/dspace_rest_client/client.py b/dspace_rest_client/client.py index 9cacb08..061680f 100644 --- a/dspace_rest_client/client.py +++ b/dspace_rest_client/client.py @@ -103,6 +103,9 @@ def __init__(self, api_endpoint=API_ENDPOINT, username=USERNAME, password=PASSWO :param api_endpoint: base path to DSpace REST API, eg. http://localhost:8080/server/api :param username: username with appropriate privileges to perform operations on REST API :param password: password for the above username + :param timeout: default per-request timeout in seconds, used by every request unless a + method call overrides it (eg. create_bitstream's own timeout argument). + None (default) falls back to DEFAULT_TIMEOUT (60s). """ self.session = requests.Session() self.API_ENDPOINT = api_endpoint @@ -830,7 +833,7 @@ def get_bitstreams(self, uuid=None, bundle=None, page=0, size=20, sort=None): bitstreams.append(Bitstream(bitstream_resource)) return bitstreams - def create_bitstream(self, bundle=None, name=None, path=None, mime=None, metadata=None, retry=False): + def create_bitstream(self, bundle=None, name=None, path=None, mime=None, metadata=None, retry=False, timeout=None): """ Upload a file and create a bitstream for a specified parent bundle, from the uploaded file and the supplied metadata. @@ -845,6 +848,10 @@ def create_bitstream(self, bundle=None, name=None, path=None, mime=None, metadat @param metadata: Full metadata JSON @param retry: A 'retried' indicator. If the first attempt fails due to an expired or missing auth token, the request will retry once, after the token is refreshed. (default: False) + @param timeout: Per-call timeout in seconds for this upload, overriding self.timeout - useful for + large files that need longer than the client's default. None (default) falls back + to self.timeout. Preserved across the CSRF-retry recursion, so it still applies + to the retried request. @return: constructed Bitstream object from the API response, or None if the operation failed. """ # TODO: It is probably wise to allow the bundle UUID to be simply passed as an alternative to having the full @@ -865,7 +872,8 @@ def create_bitstream(self, bundle=None, name=None, path=None, mime=None, metadat h.update({'Content-Encoding': 'gzip', 'User-Agent': self.USER_AGENT}) req = Request('POST', url, data=payload, headers=h, files=files) prepared_req = self.session.prepare_request(req) - r = self.session.send(prepared_req, proxies=self.proxies, timeout=self.timeout) + r = self.session.send(prepared_req, proxies=self.proxies, + timeout=timeout if timeout is not None else self.timeout) if 'DSPACE-XSRF-TOKEN' in r.headers: t = r.headers['DSPACE-XSRF-TOKEN'] _logger.debug('Updating token to ' + t) @@ -873,11 +881,12 @@ def create_bitstream(self, bundle=None, name=None, path=None, mime=None, metadat self.session.cookies.update({'X-XSRF-Token': t}) if not retry and r.status_code in (401, 403): r_json = parse_json(r) - if 'message' in r_json and 'CSRF token' in r_json['message']: + if 'message' in (r_json or {}) and 'CSRF token' in r_json['message']: _logger.debug("Retrying request with updated CSRF token") else: self.authenticate() - return self.create_bitstream(bundle, name, path, mime, metadata, True) + return self.create_bitstream(bundle=bundle, name=name, path=path, mime=mime, + metadata=metadata, retry=True, timeout=timeout) if r.status_code == 201 or r.status_code == 200: # Success diff --git a/tests/test_client_write.py b/tests/test_client_write.py index 9e4815c..a8bbec3 100644 --- a/tests/test_client_write.py +++ b/tests/test_client_write.py @@ -192,6 +192,26 @@ def test_server_error_returns_none(self): bundle=bundle, name="a.pdf", path=self.path, mime="application/pdf")) + def test_csrf_retry_preserves_custom_timeout(self): + # the retry recursion used to drop a caller-supplied timeout override, + # silently falling back to the client default on the retried request + c = make_client() + bundle = Bundle(bundle_json("bnd")) + with requests_mock.Mocker() as m: + m.post(f"{API}/core/bundles/bnd/bitstreams", [ + {"status_code": 403, "json": {"message": "CSRF token invalid"}}, + {"status_code": 201, "json": bitstream_json("bsnew", "a.pdf", size=20)}, + ]) + bs = c.create_bitstream( + bundle=bundle, name="a.pdf", path=self.path, + mime="application/pdf", timeout=900) + self.assertIsInstance(bs, Bitstream) + # initial attempt + exactly one CSRF retry, both bounded by the + # caller's override rather than the client's flat default + self.assertEqual(len(m.request_history), 2) + self.assertEqual(m.request_history[0].timeout, 900) + self.assertEqual(m.request_history[1].timeout, 900) + class TestCreateClarinAllowances(unittest.TestCase):