Fix KeyError: 'ContentRange' when received full content from S3 (#789)

* Fix `KeyError: 'ContentRange'` when received full content from S3 * Add `HTTPStatusCode` to s3 test fixture * Use single quote * use constants from http package * HttpStatus -> HTTPStatus --------- Co-authored-by: Michael Penkov <[email protected]>
piskvorky · Feb 21, 2024 · 2cba857 · 2cba857
1 parent 269c3a2
commit 2cba857
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 3 deletions.
diff --git a/smart_open/s3.py b/smart_open/s3.py
@@ -7,6 +7,7 @@
 #
 """Implements file-like objects for reading and writing from/to AWS S3."""
 
+import http
 import io
 import functools
 import logging
@@ -486,9 +487,17 @@ def _open_body(self, start=None, stop=None):
                 self,
                 response['ResponseMetadata']['RetryAttempts'],
             )
-            _, start, stop, length = smart_open.utils.parse_content_range(response['ContentRange'])
+            #
+            # range request may not always return partial content, see:
+            # https://developer.mozilla.org/en-US/docs/Web/HTTP/Range_requests#partial_request_responses
+            #
+            status_code = response['ResponseMetadata']['HTTPStatusCode']
+            if status_code == http.HTTPStatus.PARTIAL_CONTENT:
+                _, start, stop, length = smart_open.utils.parse_content_range(response['ContentRange'])
+                self._position = start
+            elif status_code == http.HTTPStatus.OK:
+                length = response["ContentLength"]
             self._content_length = length
-            self._position = start
             self._body = response['Body']
 
     def read(self, size=-1):

diff --git a/smart_open/tests/test_s3.py b/smart_open/tests/test_s3.py
@@ -156,7 +156,7 @@ def get_object(self, *args, **kwargs):
             'ContentLength': self._datasize,
             'ContentRange': 'bytes 0-%d/%d' % (self._datasize, self._datasize),
             'Body': self._body,
-            'ResponseMetadata': {'RetryAttempts': 1},
+            'ResponseMetadata': {'RetryAttempts': 1, 'HTTPStatusCode': 206},
         }