Ability to specify prefix for S3 storage

depot/io/awss3.py

@@ -56,19 +56,37 @@ def public_url(self):
         return self._key.generate_url(expires_in=0, query_auth=False)
 
 
+class BucketDriver(object):
+
+    def __init__(self, bucket, prefix):
+        self.bucket = bucket
+        self.prefix = prefix
+
+    def get_key(self, key_name):
+        return self.bucket.get_key('%s%s' % (self.prefix, key_name))
+
+    def new_key(self, key_name):
+        return self.bucket.new_key('%s%s' % (self.prefix, key_name))
+
+    def list_key_names(self):
+        keys = self.bucket.list(prefix=self.prefix)
+        return [k.name[len(self.prefix):] for k in keys]
+
+
 class S3Storage(FileStorage):
     """:class:`depot.io.interfaces.FileStorage` implementation that stores files on S3.
 
     All the files are stored inside a bucket named ``bucket`` on ``host`` which Depot
     connects to using ``access_key_id`` and ``secret_access_key``. If ``host`` is
     omitted the Amazon AWS S3 storage is used. Additionally, a canned ACL policy of
     either ``private`` or ``public-read`` can be specified with the ``policy`` parameter.
-    Finally, the ``encrypt_key`` parameter can be specified to use the server side
-    encryption feature.
+    The ``encrypt_key`` parameter can be specified to use the server side
+    encryption feature. The ``prefix`` parameter can be used to store all files
+    under specified prefix.
     """
 
     def __init__(self, access_key_id, secret_access_key, bucket=None, host=None,
-                 policy=None, encrypt_key=False):
+                 policy=None, encrypt_key=False, prefix=''):
         policy = policy or CANNED_ACL_PUBLIC_READ
         assert policy in [CANNED_ACL_PUBLIC_READ, CANNED_ACL_PRIVATE], (
             "Key policy must be %s or %s" % (CANNED_ACL_PUBLIC_READ, CANNED_ACL_PRIVATE))
@@ -82,15 +100,15 @@ def __init__(self, access_key_id, secret_access_key, bucket=None, host=None,
         if host is not None:
             kw['host'] = host
         self._conn = S3Connection(access_key_id, secret_access_key, **kw)
-        self._bucket = self._conn.lookup(bucket)
-        if self._bucket is None:
-            self._bucket = self._conn.create_bucket(bucket)
+        bucket = self._conn.lookup(bucket) or self._conn.create_bucket(bucket)
+        self._bucket_driver = BucketDriver(bucket, prefix)
+
 
     def get(self, file_or_id):
         fileid = self.fileid(file_or_id)
         _check_file_id(fileid)
 
-        key = self._bucket.get_key(fileid)
+        key = self._bucket_driver.get_key(fileid)
         if key is None:
             raise IOError('File %s not existing' % fileid)
 
@@ -125,7 +143,7 @@ def __save_file(self, key, content, filename, content_type=None):
     def create(self, content, filename=None, content_type=None):
         content, filename, content_type = self.fileinfo(content, filename, content_type)
         new_file_id = str(uuid.uuid1())
-        key = self._bucket.new_key(new_file_id)
+        key = self._bucket_driver.new_key(new_file_id)
         self.__save_file(key, content, filename, content_type)
         return new_file_id
 
@@ -139,27 +157,27 @@ def replace(self, file_or_id, content, filename=None, content_type=None):
             filename = f.filename
             content_type = f.content_type
 
-        key = self._bucket.get_key(fileid)
+        key = self._bucket_driver.get_key(fileid)
         self.__save_file(key, content, filename, content_type)
         return fileid
 
     def delete(self, file_or_id):
         fileid = self.fileid(file_or_id)
         _check_file_id(fileid)
 
-        k = self._bucket.get_key(fileid)
+        k = self._bucket_driver.get_key(fileid)
         if k:
             k.delete()
 
     def exists(self, file_or_id):
         fileid = self.fileid(file_or_id)
         _check_file_id(fileid)
 
-        k = self._bucket.get_key(fileid)
+        k = self._bucket_driver.get_key(fileid)
         return k is not None
 
     def list(self):
-        return [key.name for key in self._bucket.list()]
+        return self._bucket_driver.list_key_names()
 
 
 def _check_file_id(file_id):

tests/test_awss3_storage.py

@@ -31,15 +31,15 @@ def setup(self):
 
     def test_fileoutside_depot(self):
         fid = str(uuid.uuid1())
-        key = self.fs._bucket.new_key(fid)
+        key = self.fs._bucket_driver.new_key(fid)
         key.set_contents_from_string(FILE_CONTENT)
 
         f = self.fs.get(fid)
         assert f.read() == FILE_CONTENT
 
     def test_invalid_modified(self):
         fid = str(uuid.uuid1())
-        key = self.fs._bucket.new_key(fid)
+        key = self.fs._bucket_driver.new_key(fid)
         key.set_metadata('x-depot-modified', 'INVALID')
         key.set_contents_from_string(FILE_CONTENT)
 
@@ -56,23 +56,23 @@ def test_creates_bucket_when_missing(self):
     def test_default_bucket_name(self):
         with mock.patch('boto.s3.connection.S3Connection.lookup', return_value='YES'):
             fs = S3Storage(*self.cred)
-            assert fs._bucket == 'YES'
+            assert fs._bucket_driver.bucket == 'YES'
 
     def test_public_url(self):
         fid = str(uuid.uuid1())
-        key = self.fs._bucket.new_key(fid)
+        key = self.fs._bucket_driver.new_key(fid)
         key.set_contents_from_string(FILE_CONTENT)
 
         f = self.fs.get(fid)
         assert '.s3.amazonaws.com' in f.public_url, f.public_url
         assert f.public_url.endswith('/%s' % fid), f.public_url
 
     def teardown(self):
-        keys = [key.name for key in self.fs._bucket]
+        keys = [key.name for key in self.fs._bucket_driver.bucket]
         if keys:
-            self.fs._bucket.delete_keys(keys)
+            self.fs._bucket_driver.bucket.delete_keys(keys)
 
         try:
-            self.fs._conn.delete_bucket(self.fs._bucket.name)
+            self.fs._conn.delete_bucket(self.fs._bucket_driver.bucket.name)
         except:
             pass

tests/test_storage_interface.py

@@ -253,6 +253,12 @@ def teardown(self):
 
 
 class TestS3FileStorage(BaseStorageTestFixture):
+
+    @classmethod
+    def get_storage(cls, access_key_id, secret_access_key, bucket_name):
+        from depot.io.awss3 import S3Storage
+        return S3Storage(access_key_id, secret_access_key, bucket_name)
+
     @classmethod
     def setup_class(cls):
         try:
@@ -269,16 +275,28 @@ def setup_class(cls):
         PID = os.getpid()
         NODE = str(uuid.uuid1()).rsplit('-', 1)[-1]
         BUCKET_NAME = 'fdtest-%s-%s-%s' % (access_key_id.lower(), NODE, PID)
-        cls.fs = S3Storage(access_key_id, secret_access_key, BUCKET_NAME)
+        cls.fs = cls.get_storage(access_key_id, secret_access_key, BUCKET_NAME)
 
     def teardown(self):
-        keys = [key.name for key in self.fs._bucket]
+        keys = [key.name for key in self.fs._bucket_driver.bucket]
         if keys:
-            self.fs._bucket.delete_keys(keys)
+            self.fs._bucket_driver.bucket.delete_keys(keys)
 
     @classmethod
     def teardown_class(cls):
         try:
-            cls.fs._conn.delete_bucket(cls.fs._bucket.name)
+            cls.fs._conn.delete_bucket(cls.fs._bucket_driver.bucket.name)
         except:
             pass
+
+
+class TestS3FileStorageWithPrefix(TestS3FileStorage):
+
+    @classmethod
+    def get_storage(cls, access_key_id, secret_access_key, bucket_name):
+        from depot.io.awss3 import S3Storage
+        return S3Storage(
+            access_key_id,
+            secret_access_key,
+            bucket_name,
+            prefix='my-prefix/')