From d35fae7bcc48429292d2b8efe2f30cbab2d0c52e Mon Sep 17 00:00:00 2001
From: Patrick Clay <pclay@google.com>
Date: Thu, 12 Dec 2024 13:39:18 -0800
Subject: [PATCH] Wire --object_storage_storage_class through s3.py

PiperOrigin-RevId: 705616161
---
 .../object_storage_service_benchmark.py               |  5 ++++-
 .../object_storage_interface.py                       | 11 +++++------
 .../scripts/object_storage_api_test_scripts/s3.py     |  9 ++++++++-
 3 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/perfkitbenchmarker/linux_benchmarks/object_storage_service_benchmark.py b/perfkitbenchmarker/linux_benchmarks/object_storage_service_benchmark.py
index 4015afdefb..eee9bb7fd0 100644
--- a/perfkitbenchmarker/linux_benchmarks/object_storage_service_benchmark.py
+++ b/perfkitbenchmarker/linux_benchmarks/object_storage_service_benchmark.py
@@ -98,7 +98,7 @@
     'Storage multiregion for GCS in object storage benchmark.',
 )
 
-flags.DEFINE_string(
+OBJECT_STORAGE_STORAGE_CLASS = flags.DEFINE_string(
     'object_storage_storage_class',
     None,
     'Storage class to use in object storage benchmark.',
@@ -1940,6 +1940,9 @@ def Run(benchmark_spec):
 
   metadata.update(service.Metadata(vms[0]))
 
+  if OBJECT_STORAGE_STORAGE_CLASS.value:
+    metadata['object_storage_class'] = OBJECT_STORAGE_STORAGE_CLASS.value
+
   results = []
   test_script_path = '/tmp/run/%s' % API_TEST_SCRIPT
   try:
diff --git a/perfkitbenchmarker/scripts/object_storage_api_test_scripts/object_storage_interface.py b/perfkitbenchmarker/scripts/object_storage_api_test_scripts/object_storage_interface.py
index 196ee014ca..657a5895cd 100644
--- a/perfkitbenchmarker/scripts/object_storage_api_test_scripts/object_storage_interface.py
+++ b/perfkitbenchmarker/scripts/object_storage_api_test_scripts/object_storage_interface.py
@@ -15,10 +15,9 @@
 """The generic superclass for object storage API providers."""
 
 import abc
-import six
 
 
-class ObjectStorageServiceBase(six.with_metaclass(abc.ABCMeta, object)):
+class ObjectStorageServiceBase(metaclass=abc.ABCMeta):
   """Our interface to an object storage service."""
 
   @abc.abstractmethod
@@ -88,7 +87,7 @@ def BulkDeleteObjects(self, bucket, objects_to_delete, delay_time):
     pass
 
   @abc.abstractmethod
-  def WriteObjectFromBuffer(self, bucket, object, stream, size):
+  def WriteObjectFromBuffer(self, bucket, object_name, stream, size):
     """Write an object to a bucket.
 
     Exceptions are propagated to the caller, which can decide whether
@@ -97,7 +96,7 @@ def WriteObjectFromBuffer(self, bucket, object, stream, size):
 
     Args:
       bucket: the name of the bucket to write to.
-      object: the name of the object.
+      object_name: the name of the object.
       stream: a read()-able and seek()-able stream to transfer.
       size: the number of bytes to transfer.
 
@@ -108,7 +107,7 @@ def WriteObjectFromBuffer(self, bucket, object, stream, size):
     pass
 
   @abc.abstractmethod
-  def ReadObject(self, bucket, object):
+  def ReadObject(self, bucket, object_name):
     """Read an object.
 
     Exceptions are propagated to the caller, which can decide whether
@@ -116,7 +115,7 @@ def ReadObject(self, bucket, object):
 
     Args:
       bucket: the name of the bucket.
-      object: the name of the object.
+      object_name: the name of the object.
 
     Returns:
       A tuple of (start_time, latency)
diff --git a/perfkitbenchmarker/scripts/object_storage_api_test_scripts/s3.py b/perfkitbenchmarker/scripts/object_storage_api_test_scripts/s3.py
index 83b6e514aa..08400d32b4 100644
--- a/perfkitbenchmarker/scripts/object_storage_api_test_scripts/s3.py
+++ b/perfkitbenchmarker/scripts/object_storage_api_test_scripts/s3.py
@@ -80,7 +80,14 @@ def WriteObjectFromBuffer(self, bucket, object_name, stream, size):
     start_time = time.time()
     stream.seek(0)
     obj = six.BytesIO(stream.read(size))
-    self.client.put_object(Body=obj, Bucket=bucket, Key=object_name)
+    # boto requires additional parameters to be not passed or valid.
+    # Passing StorageClass=None causes a type error. So use kwargs.
+    extra_kwargs = {}
+    if FLAGS.object_storage_class:
+      extra_kwargs['StorageClass'] = FLAGS.object_storage_class
+    self.client.put_object(
+        Body=obj, Bucket=bucket, Key=object_name, **extra_kwargs
+    )
     latency = time.time() - start_time
     return start_time, latency