17351: Adds support for storage classes while writing blocks to Keep.

Also, enabled support for multiple storage classes. For some reason it was
limited to just one.

Arvados-DCO-1.1-Signed-off-by: Lucas Di Pentima <lucas.dipentima@curii.com>

diff --git a/sdk/python/arvados/commands/put.py b/sdk/python/arvados/commands/put.py
index 9596a2dc2d26efcddebbd7921e32ff38ebdb7352..0920369646c6b7556babdef38fced98fee2e18b8 100644 (file)
--- a/sdk/python/arvados/commands/put.py
+++ b/sdk/python/arvados/commands/put.py
@@ -656,15 +656,11 @@ class ArvPutUploadJob(object):
                 else:
                     # The file already exist on remote collection, skip it.
                     pass
-            self._remote_collection.save(storage_classes=self.storage_classes,
-                                         num_retries=self.num_retries,
+            self._remote_collection.save(num_retries=self.num_retries,
                                          trash_at=self._collection_trash_at())
         else:
-            if self.storage_classes is None:
-                self.storage_classes = ['default']
             self._local_collection.save_new(
                 name=self.name, owner_uuid=self.owner_uuid,
-                storage_classes=self.storage_classes,
                 ensure_unique_name=self.ensure_unique_name,
                 num_retries=self.num_retries,
                 trash_at=self._collection_trash_at())
@@ -868,6 +864,7 @@ class ArvPutUploadJob(object):
                 self._remote_collection = arvados.collection.Collection(
                     update_collection,
                     api_client=self._api_client,
+                    storage_classes_desired=self.storage_classes,
                     num_retries=self.num_retries)
             except arvados.errors.ApiError as error:
                 raise CollectionUpdateError("Cannot read collection {} ({})".format(update_collection, error))
@@ -910,6 +907,7 @@ class ArvPutUploadJob(object):
             self._local_collection = arvados.collection.Collection(
                 self._state['manifest'],
                 replication_desired=self.replication_desired,
+                storage_classes_desired=(self.storage_classes or ['default']),
                 put_threads=self.put_threads,
                 api_client=self._api_client,
                 num_retries=self.num_retries)
@@ -1197,11 +1195,7 @@ def main(arguments=None, stdout=sys.stdout, stderr=sys.stderr,
     #  Split storage-classes argument
     storage_classes = None
     if args.storage_classes:
-        storage_classes = args.storage_classes.strip().split(',')
-        if len(storage_classes) > 1:
-            logger.error("Multiple storage classes are not supported currently.")
-            sys.exit(1)
-
+        storage_classes = args.storage_classes.strip().replace(' ', '').split(',')
 
     # Setup exclude regex from all the --exclude arguments provided
     name_patterns = []

diff --git a/sdk/python/tests/test_arv_put.py b/sdk/python/tests/test_arv_put.py
index e75d39d8793ffb6032a8591a37d9d5da2168ba72..caa03a3e03c02692204f3e707fe31df7a97961ad 100644 (file)
--- a/sdk/python/tests/test_arv_put.py
+++ b/sdk/python/tests/test_arv_put.py
@@ -813,11 +813,6 @@ class ArvadosPutTest(run_test_server.TestCaseWithServers,
                           self.call_main_with_args,
                           ['--project-uuid', self.Z_UUID, '--stream'])
 
-    def test_error_when_multiple_storage_classes_specified(self):
-        self.assertRaises(SystemExit,
-                          self.call_main_with_args,
-                          ['--storage-classes', 'hot,cold'])
-
     def test_error_when_excluding_absolute_path(self):
         tmpdir = self.make_tmpdir()
         self.assertRaises(SystemExit,
@@ -1315,13 +1310,16 @@ class ArvPutIntegrationTest(run_test_server.TestCaseWithServers,
 
     def test_put_collection_with_storage_classes_specified(self):
         collection = self.run_and_find_collection("", ['--storage-classes', 'hot'])
-
         self.assertEqual(len(collection['storage_classes_desired']), 1)
         self.assertEqual(collection['storage_classes_desired'][0], 'hot')
 
+    def test_put_collection_with_multiple_storage_classes_specified(self):
+        collection = self.run_and_find_collection("", ['--storage-classes', ' foo, bar  ,baz'])
+        self.assertEqual(len(collection['storage_classes_desired']), 3)
+        self.assertEqual(collection['storage_classes_desired'], ['foo', 'bar', 'baz'])
+
     def test_put_collection_without_storage_classes_specified(self):
         collection = self.run_and_find_collection("")
-
         self.assertEqual(len(collection['storage_classes_desired']), 1)
         self.assertEqual(collection['storage_classes_desired'][0], 'default')