Merge pull request #2874 from coreos-inc/joseph.schorr/QS-22/swift-copy-to

Add missing `copy_to` method to Swift

storage/swift.py

@@ -355,3 +355,28 @@ class SwiftStorage(BaseStorage):
     # Delete all the uploaded segments.
     for segment in SwiftStorage._segment_list_from_metadata(storage_metadata, key=_SEGMENTS_KEY):
       self.remove(segment.path)
+
+  def copy_to(self, destination, path):
+    if (self.__class__ == destination.__class__ and
+        self._swift_user == destination._swift_user and
+        self._swift_password == destination._swift_password and
+        self._auth_url == destination._auth_url and
+        self._auth_version == destination._auth_version):
+        logger.debug('Copying file from swift %s to swift %s via a Swift copy',
+                     self._swift_container, destination)
+
+        normalized_path = self._normalize_path(path)
+        target = '/%s/%s' % (destination._swift_container, normalized_path)
+
+        try:
+          self._get_connection().copy_object(self._swift_container, normalized_path, target)
+        except Exception as ex:
+          logger.exception('Could not swift copy path %s: %s', path, ex)
+          raise IOError('Failed to swift copy path %s' % path)        
+        return
+
+    # Fallback to a slower, default copy.
+    logger.debug('Copying file from swift %s to %s via a streamed copy', self._swift_container,
+                 destination)
+    with self.stream_read_file(path) as fp:
+      destination.stream_write(path, fp)

storage/test/test_swift.py

@@ -1,6 +1,7 @@
 import io
 import pytest
 import hashlib
+import copy
 
 from collections import defaultdict
 from mock import MagicMock
@@ -26,9 +27,9 @@ class MockSwiftStorage(SwiftStorage):
     return self._connection
 
 class FakeSwiftStorage(SwiftStorage):
-  def __init__(self, fail_checksum=False, *args, **kwargs):
+  def __init__(self, fail_checksum=False, connection=None, *args, **kwargs):
     super(FakeSwiftStorage, self).__init__(*args, **kwargs)
-    self._connection = FakeSwift(fail_checksum=fail_checksum)
+    self._connection = connection or FakeSwift(fail_checksum=fail_checksum)
 
   def _get_connection(self):
     return self._connection
@@ -42,6 +43,11 @@ class FakeSwift(object):
   def head_object(self, container, path):
     return self.containers[container].get(path)
 
+  def copy_object(self, container, path, target):
+    pieces = target.split('/', 2)
+    _, content = self.get_object(container, path)
+    self.put_object(pieces[1], pieces[2], content)
+
   def put_object(self, container, path, content, chunk_size=None, content_type=None, headers=None):
     if not isinstance(content, str):
       if hasattr(content, 'read'):
@@ -53,7 +59,7 @@ class FakeSwift(object):
       'content': content,
       'chunk_size': chunk_size,
       'content_type': content_type,
-      'headers': headers,
+      'headers': headers or {},
     }
 
     digest = hashlib.md5()
@@ -103,7 +109,6 @@ def test_fixed_path_concat():
   swift.exists('object/path')
   swift._get_connection().head_object.assert_called_with('container-name', 'basepath/object/path')
 
-
 def test_simple_path_concat():
   simple_concat_args = dict(base_args)
   simple_concat_args['simple_path_concat'] = True
@@ -150,6 +155,41 @@ def test_remove():
   swift.remove('somepath')
   assert not swift.exists('somepath')
 
+def test_copy_to():
+  swift = FakeSwiftStorage(**base_args)
+
+  modified_args = copy.deepcopy(base_args)
+  modified_args['swift_container'] = 'another_container'
+
+  another_swift = FakeSwiftStorage(connection=swift._connection, **modified_args)
+
+  swift.put_content('somepath', 'some content here')
+  swift.copy_to(another_swift, 'somepath')
+
+  assert swift.exists('somepath')
+  assert another_swift.exists('somepath')
+
+  assert swift.get_content('somepath') == 'some content here'
+  assert another_swift.get_content('somepath') == 'some content here'  
+
+def test_copy_to_different():
+  swift = FakeSwiftStorage(**base_args)
+
+  modified_args = copy.deepcopy(base_args)
+  modified_args['swift_user'] = 'foobarbaz'
+  modified_args['swift_container'] = 'another_container'
+
+  another_swift = FakeSwiftStorage(**modified_args)
+
+  swift.put_content('somepath', 'some content here')
+  swift.copy_to(another_swift, 'somepath')
+
+  assert swift.exists('somepath')
+  assert another_swift.exists('somepath')
+
+  assert swift.get_content('somepath') == 'some content here'
+  assert another_swift.get_content('somepath') == 'some content here'  
+
 def test_checksum():
   swift = FakeSwiftStorage(**base_args)
   swift.put_content('somepath', 'hello world!')