add timeout override to create_bitstream

Author: jr-rk

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,19 +872,21 @@ 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)
             self.session.headers.update({'X-XSRF-Token': t})
             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

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):