Add support for generating CLI commands to object_storage service.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=288604096

CHANGES.next.md

@@ -391,6 +391,7 @@
   https://github.com/GoogleCloudPlatform/spark-bigquery-connector. Support is
   currently limited to GCP Dataproc provider, but it could run anywhere
   provided some auth is plumbed through.
+- Add support for generating CLI download commands for object_storage_service.
 
 ### Bug fixes and maintenance updates:
 - Moved GPU-related specs from GceVmSpec to BaseVmSpec

perfkitbenchmarker/object_storage_service.py

@@ -110,6 +110,41 @@ def Copy(self, src_url, dst_url):
     """
     pass
 
+  @abc.abstractmethod
+  def CopyToBucket(self, src_path, bucket, object_path):
+    """Copy a local file to a bucket.
+
+    Args:
+      src_path: string, the local source path.
+      bucket: string, the destination bucket.
+      object_path: string, the object's path in the bucket.
+    """
+    pass
+
+  @abc.abstractmethod
+  def MakeRemoteCliDownloadUrl(self, bucket, object_path):
+    """Creates a download url for an object in a bucket.
+
+    This is used by GenerateCliDownloadFileCommand().
+
+    Args:
+      bucket: string, the name of the bucket.
+      object_path: string, the path of the object in the bucket.
+    """
+    pass
+
+  @abc.abstractmethod
+  def GenerateCliDownloadFileCommand(self, src_url, local_path):
+    """Generates a CLI command to copy src_url to local_path.
+
+    This is suitable for use in scripts e.g. startup scripts.
+
+    Args:
+      src_url: string, the source url path.
+      local_path: string, the local path.
+    """
+    pass
+
   @abc.abstractmethod
   def List(self, bucket):
     """List providers, buckets, or objects.

perfkitbenchmarker/providers/aws/s3.py

@@ -14,6 +14,7 @@
 
 """Contains classes/functions related to S3."""
 
+import posixpath
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import flags
 from perfkitbenchmarker import linux_packages
@@ -67,6 +68,20 @@ def Copy(self, src_url, dst_url):
     """See base class."""
     vm_util.IssueCommand(['aws', 's3', 'cp', src_url, dst_url])
 
+  def CopyToBucket(self, src_path, bucket, object_path):
+    """See base class."""
+    dst_url = self.MakeRemoteCliDownloadUrl(bucket, object_path)
+    vm_util.IssueCommand(['aws', 's3', 'cp', src_path, dst_url])
+
+  def MakeRemoteCliDownloadUrl(self, bucket, object_path):
+    """See base class."""
+    path = posixpath.join(bucket, object_path)
+    return 's3://' + path
+
+  def GenerateCliDownloadFileCommand(self, src_url, local_path):
+    """See base class."""
+    return 'aws s3 cp "%s" "%s"' % (src_url, local_path)
+
   def List(self, buckets):
     """See base class."""
     stdout, _, _ = vm_util.IssueCommand(['aws', 's3', 'ls', buckets])

perfkitbenchmarker/providers/azure/azure_blob_storage.py

@@ -14,6 +14,7 @@
 
 """Contains classes/functions related to Azure Blob Storage."""
 
+import datetime
 from perfkitbenchmarker import errors
 from perfkitbenchmarker import flags
 from perfkitbenchmarker import linux_packages
@@ -131,6 +132,42 @@ def Copy(self, src_url, dst_url):
     """See base class."""
     raise NotImplementedError()
 
+  def CopyToBucket(self, src_path, bucket, object_path):
+    vm_util.IssueCommand(['az', 'storage', 'blob', 'upload',
+                          '--account-name', self.storage_account.name,
+                          '--file', src_path,
+                          '--container', bucket,
+                          '--name', object_path])
+
+  def _GenerateDownloadToken(self, bucket, object_path):
+    blob_store_expiry = datetime.datetime.utcnow() + datetime.timedelta(
+        days=365)
+    stdout, _, _ = vm_util.IssueCommand([
+        'az', 'storage', 'blob', 'generate-sas',
+        '--account-name', self.storage_account.name,
+        '--container-name', bucket,
+        '--name', object_path,
+        '--expiry', blob_store_expiry.strftime('%Y-%m-%dT%H:%M:%SZ'),
+        '--permissions', 'r'
+    ])
+    token = stdout.strip('\n').strip('"')
+    return token
+
+  def MakeRemoteCliDownloadUrl(self, bucket, object_path):
+    """See base class."""
+    token = self._GenerateDownloadToken(bucket, object_path)
+    url = 'https://{acc}.blob.core.windows.net/{con}/{src}?{tkn}'.format(
+        acc=self.storage_account.name,
+        con=bucket,
+        src=object_path,
+        tkn=token)
+    return url
+
+  def GenerateCliDownloadFileCommand(self, src_url, dst_url):
+    """See base class."""
+    return 'wget -O {dst_url} "{src_url}"'.format(src_url=src_url,
+                                                  dst_url=dst_url)
+
   def List(self, buckets):
     """See base class."""
     raise NotImplementedError()

perfkitbenchmarker/providers/azure/azure_network.py

@@ -131,7 +131,7 @@ def AddTag(self, key, value):
     ]
     _, _, retcode = vm_util.IssueCommand(tag_cmd, raise_on_failure=False)
     if retcode:
-      raise errors.resource.CreationError('Error tagging Azure resource group.')
+      raise errors.Resource.CreationError('Error tagging Azure resource group.')
 
 
 class AzureStorageAccount(resource.BaseResource):

perfkitbenchmarker/providers/gcp/gcs.py

@@ -65,6 +65,20 @@ def Copy(self, src_url, dst_url):
     """See base class."""
     vm_util.IssueCommand(['gsutil', 'cp', src_url, dst_url])
 
+  def CopyToBucket(self, src_path, bucket, object_path):
+    """See base class."""
+    dst_url = self.MakeRemoteCliDownloadUrl(bucket, object_path)
+    vm_util.IssueCommand(['gsutil', 'cp', src_path, dst_url])
+
+  def MakeRemoteCliDownloadUrl(self, bucket, object_path):
+    """See base class."""
+    path = posixpath.join(bucket, object_path)
+    return 'gs://' + path
+
+  def GenerateCliDownloadFileCommand(self, src_url, local_path):
+    """See base class."""
+    return 'gsutil cp "%s" "%s"' % (src_url, local_path)
+
   def List(self, buckets):
     """See base class."""
     stdout, _, _ = vm_util.IssueCommand(['gsutil', 'ls', buckets])

perfkitbenchmarker/providers/openstack/swift.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+"""Module containing classes related to Swift Storage Service."""
 
 import os
 
@@ -75,6 +76,18 @@ def Copy(self, src_url, dst_url):
     """See base class."""
     raise NotImplementedError()
 
+  def CopyToBucket(self, src_path, bucket, object_path):
+    """See base class."""
+    raise NotImplementedError()
+
+  def MakeRemoteCliDownloadUrl(self, bucket, object_path):
+    """See base class."""
+    raise NotImplementedError()
+
+  def GenerateCliDownloadFileCommand(self, src_url, local_path):
+    """See base class."""
+    raise NotImplementedError()
+
   def List(self, buckets):
     """See base class."""
     raise NotImplementedError()

tests/providers/openstack/swift_test.py

@@ -0,0 +1,45 @@
+"""Tests for perfkitbenchmarker.providers.openstack.swift."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import os
+import unittest
+import mock
+from perfkitbenchmarker import vm_util
+from perfkitbenchmarker.providers.openstack import swift
+
+
+class SwiftTest(unittest.TestCase):
+
+  def setUp(self):
+    super(SwiftTest, self).setUp()
+    p = mock.patch(swift.__name__ + '.FLAGS')
+    self.mock_flags = p.start()
+    self.addCleanup(p.stop)
+    self.mock_flags.openstack_swift_insecure = False
+
+  @mock.patch.dict(os.environ, {'OS_AUTH_URL': 'OS_AUTH_URL',
+                                'OS_TENANT_NAME': 'OS_TENANT_NAME',
+                                'OS_USERNAME': 'OS_USERNAME',
+                                'OS_PASSWORD': 'OS_PASSWORD'})
+  def testMakeBucket(self):
+    swift_storage_service = swift.SwiftStorageService()
+    swift_storage_service.PrepareService('location')
+
+    with mock.patch(vm_util.__name__ + '.IssueCommand',
+                    return_value=('stdout', 'stderr', 0)) as mock_util:
+      swift_storage_service.MakeBucket('new_bucket')
+      mock_util.assert_called_with(['swift',
+                                    '--os-auth-url', 'OS_AUTH_URL',
+                                    '--os-tenant-name', 'OS_TENANT_NAME',
+                                    '--os-username', 'OS_USERNAME',
+                                    '--os-password', 'OS_PASSWORD',
+                                    'post',
+                                    'new_bucket'],
+                                   raise_on_failure=False)
+
+
+if __name__ == '__main__':
+  unittest.main()