Add version_id transport parameter for fetching a specific S3 object version

smart_open/s3.py

@@ -70,6 +70,7 @@ def open(
         bucket_id,
         key_id,
         mode,
+        version_id=None,
         buffer_size=DEFAULT_BUFFER_SIZE,
         min_part_size=DEFAULT_MIN_PART_SIZE,
         session=None,
@@ -97,6 +98,8 @@ def open(
     multipart_upload_kwargs: dict, optional
         Additional parameters to pass to boto3's initiate_multipart_upload function.
         For writing only.
+    version_id: str, optional
+        Version of the object. If None, will fetch the most recent version.
 
     """
     logger.debug('%r', locals())
@@ -112,6 +115,7 @@ def open(
         fileobj = SeekableBufferedInputBase(
             bucket_id,
             key_id,
+            version_id=version_id,
             buffer_size=buffer_size,
             session=session,
             resource_kwargs=resource_kwargs,
@@ -127,10 +131,21 @@ def open(
         )
     else:
         assert False, 'unexpected mode: %r' % mode
-
     return fileobj
 
 
+def _get(s3_object, version=None, **kwargs):
+    if version is not None:
+        kwargs['VersionId'] = version
+    try:
+        return s3_object.get(**kwargs)
+    except botocore.client.ClientError as error:
+        raise IOError(str(error) + "\n" +
+                      '%r does not exist in the bucket %r, version_id %r'
+                      'or is forbidden for access' % (s3_object.key, s3_object.bucket_name, version)
+                      )
+
+
 class RawReader(object):
     """Read an S3 object."""
     def __init__(self, s3_object):
@@ -147,9 +162,10 @@ def read(self, size=-1):
 class SeekableRawReader(object):
     """Read an S3 object."""
 
-    def __init__(self, s3_object, content_length):
+    def __init__(self, s3_object, content_length, version_id=None):
         self._object = s3_object
         self._content_length = content_length
+        self._version_id = version_id
         self.seek(0)
 
     def seek(self, position):
@@ -177,7 +193,7 @@ def seek(self, position):
             #
             self._body = io.BytesIO()
         else:
-            self._body = self._object.get(Range=range_string)['Body']
+            self._body = _get(self._object, self._version_id, Range=range_string)['Body']
 
     def read(self, size=-1):
         if self._position >= self._content_length:
@@ -191,7 +207,7 @@ def read(self, size=-1):
 
 
 class BufferedInputBase(io.BufferedIOBase):
-    def __init__(self, bucket, key, buffer_size=DEFAULT_BUFFER_SIZE,
+    def __init__(self, bucket, key, version_id=None, buffer_size=DEFAULT_BUFFER_SIZE,
                  line_terminator=BINARY_NEWLINE, session=None, resource_kwargs=None):
         if session is None:
             session = boto3.Session()
@@ -202,6 +218,7 @@ def __init__(self, bucket, key, buffer_size=DEFAULT_BUFFER_SIZE,
         self._object = s3.Object(bucket, key)
         self._raw_reader = RawReader(self._object)
         self._content_length = self._object.content_length
+        self._content_length = _get(self._object, self._version_id)['ContentLength']
         self._current_pos = 0
         self._buffer = smart_open.bytebuffer.ByteBuffer(buffer_size)
         self._eof = False
@@ -328,24 +345,18 @@ class SeekableBufferedInputBase(BufferedInputBase):
 
     Implements the io.BufferedIOBase interface of the standard library."""
 
-    def __init__(self, bucket, key, buffer_size=DEFAULT_BUFFER_SIZE,
+    def __init__(self, bucket, key, version_id=None, buffer_size=DEFAULT_BUFFER_SIZE,
                  line_terminator=BINARY_NEWLINE, session=None, resource_kwargs=None):
         if session is None:
             session = boto3.Session()
         if resource_kwargs is None:
             resource_kwargs = {}
         s3 = session.resource('s3', **resource_kwargs)
         self._object = s3.Object(bucket, key)
+        self._version_id = version_id
+        self._content_length = _get(self._object, self._version_id)['ContentLength']
 
-        try:
-            self._content_length = self._object.content_length
-        except botocore.client.ClientError:
-            raise ValueError(
-                '%r does not exist in the bucket %r, '
-                'or is forbidden for access' % (key, bucket)
-            )
-
-        self._raw_reader = SeekableRawReader(self._object, self._content_length)
+        self._raw_reader = SeekableRawReader(self._object, self._content_length, self._version_id)
         self._current_pos = 0
         self._buffer = smart_open.bytebuffer.ByteBuffer(buffer_size)
         self._eof = False

smart_open/tests/test_s3_version.py

@@ -0,0 +1,101 @@
+# -*- coding: utf-8 -*-
+import gzip
+import io
+import logging
+import os
+import time
+import unittest
+import uuid
+import warnings
+
+import boto.s3.bucket
+import boto3
+import botocore.client
+import mock
+import moto
+import six
+
+import smart_open
+
+BUCKET_NAME = 'test-smartopen-{}'.format(uuid.uuid4().hex)
+KEY_NAME = 'test-key'
+WRITE_KEY_NAME = 'test-write-key'
+DISABLE_MOCKS = os.environ.get('SO_DISABLE_MOCKS') == "1"
+
+
+logger = logging.getLogger(__name__)
+
+
+def maybe_mock_s3(func):
+    if DISABLE_MOCKS:
+        return func
+    else:
+        return moto.mock_s3(func)
+
+
+@maybe_mock_s3
+def setUpModule():
+    '''Called once by unittest when initializing this module.  Sets up the
+    test S3 bucket.
+
+    '''
+    boto3.resource('s3').create_bucket(Bucket=BUCKET_NAME)
+    boto3.resource('s3').BucketVersioning(BUCKET_NAME).enable()
+
+
+def put_to_bucket(contents, num_attempts=12, sleep_time=5):
+    # fake (or not) connection, bucket and key
+    logger.debug('%r', locals())
+
+    #
+    # In real life, it can take a few seconds for the bucket to become ready.
+    # If we try to write to the key while the bucket while it isn't ready, we
+    # will get a ClientError: NoSuchBucket.
+    #
+    for attempt in range(num_attempts):
+        try:
+            boto3.resource('s3').Object(BUCKET_NAME, KEY_NAME).put(Body=contents)
+            return
+        except botocore.exceptions.ClientError as err:
+            logger.error('caught %r, retrying', err)
+            time.sleep(sleep_time)
+
+    assert False, 'failed to create bucket %s after %d attempts' % (BUCKET_NAME, num_attempts)
+
+
+@maybe_mock_s3
+class TestVersionId(unittest.TestCase):
+
+    def test_good_id(self):
+        """Does version_id into s3 work correctly?"""
+        test_string_ver0 = u"String version 1.0".encode('utf8')
+        test_string_ver1 = u"String version 2.0".encode('utf8')
+        # write into key
+        with smart_open.s3.BufferedOutputBase(BUCKET_NAME, WRITE_KEY_NAME) as fout:
+            fout.write(test_string_ver0)
+        with smart_open.s3.BufferedOutputBase(BUCKET_NAME, WRITE_KEY_NAME) as fout:
+            fout.write(test_string_ver1)
+
+        versions = boto3.resource('s3').Bucket(BUCKET_NAME).object_versions.filter(Prefix=WRITE_KEY_NAME)
+        check_version = list(versions)[0].get()['VersionId']
+        with smart_open.s3.SeekableBufferedInputBase(BUCKET_NAME, WRITE_KEY_NAME,check_version) as fin:
+            expected = fin.read()
+        self.assertEqual(expected, test_string_ver0)
+
+    def test_bad_id(self):
+        """Does version_id exception into s3 work correctly?"""
+        test_string_ver0 = u"String version 1.0".encode('utf8')
+        test_string_ver1 = u"String version 2.0".encode('utf8')
+        # write into key
+        with smart_open.s3.BufferedOutputBase(BUCKET_NAME, WRITE_KEY_NAME) as fout:
+            fout.write(test_string_ver0)
+        with smart_open.s3.BufferedOutputBase(BUCKET_NAME, WRITE_KEY_NAME) as fout:
+            fout.write(test_string_ver1)
+        versions = boto3.resource('s3').Bucket(BUCKET_NAME).object_versions.filter(Prefix=WRITE_KEY_NAME)
+        check_version = list(versions)[0].get()['VersionId']
+
+        with self.assertRaises(IOError):
+            smart_open.s3.open(BUCKET_NAME, WRITE_KEY_NAME, 'rb', version_id=check_version+check_version)
+
+if __name__ == '__main__':
+    unittest.main()