17696: KeepClient uses default storage classes when none is required explicitly

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

diff --git a/sdk/python/arvados/keep.py b/sdk/python/arvados/keep.py
index 86b1d91b8246ef20ed860cdb516ca5d4a53624b9..9dfe0436dec9bdf22eb71ad9bfe2e8a201ee3ab6 100644 (file)
--- a/sdk/python/arvados/keep.py
+++ b/sdk/python/arvados/keep.py
@@ -842,6 +842,7 @@ class KeepClient(object):
         self.hits_counter = Counter()
         self.misses_counter = Counter()
         self._storage_classes_unsupported_warning = False
+        self._default_classes = []
 
         if local_store:
             self.local_store = local_store
@@ -882,6 +883,12 @@ class KeepClient(object):
                 self._writable_services = None
                 self.using_proxy = None
                 self._static_services_list = False
+                try:
+                    self._default_classes = [
+                        k for k, v in self.api_client.config()['StorageClasses'].items() if v['Default']]
+                except KeyError:
+                    # We're talking to an old cluster
+                    pass
 
     def current_timeout(self, attempt_number):
         """Return the appropriate timeout to use for this client.
@@ -1174,7 +1181,7 @@ class KeepClient(object):
                 "failed to read {} after {}".format(loc_s, loop.attempts_str()), service_errors, label="service")
 
     @retry.retry_method
-    def put(self, data, copies=2, num_retries=None, request_id=None, classes=[]):
+    def put(self, data, copies=2, num_retries=None, request_id=None, classes=None):
         """Save data in Keep.
 
         This method will get a list of Keep services from the API server, and
@@ -1195,6 +1202,8 @@ class KeepClient(object):
           be written.
         """
 
+        classes = classes or self._default_classes
+
         if not isinstance(data, bytes):
             data = data.encode()
 

diff --git a/sdk/python/tests/arvados_testutil.py b/sdk/python/tests/arvados_testutil.py
index f251ea654b5fed6d0c5f9c837c720aa733b02baf..d9b3ca86c4f9055dde2fa9b54ad63ed65d16d755 100644 (file)
--- a/sdk/python/tests/arvados_testutil.py
+++ b/sdk/python/tests/arvados_testutil.py
@@ -190,7 +190,13 @@ class MockStreamReader(object):
 
 class ApiClientMock(object):
     def api_client_mock(self):
-        return mock.MagicMock(name='api_client_mock')
+        api_mock = mock.MagicMock(name='api_client_mock')
+        api_mock.config.return_value = {
+            'StorageClasses': {
+                'default': {'Default': True}
+            }
+        }
+        return api_mock
 
     def mock_keep_services(self, api_mock=None, status=200, count=12,
                            service_type='disk',

diff --git a/sdk/python/tests/test_keep_client.py b/sdk/python/tests/test_keep_client.py
index 0eefa586d9c436413e2e1934d9cf401e4ed17467..b1c42fd2b3a1475934a0c6090e12139750210f46 100644 (file)
--- a/sdk/python/tests/test_keep_client.py
+++ b/sdk/python/tests/test_keep_client.py
@@ -540,26 +540,49 @@ class KeepStorageClassesTestCase(unittest.TestCase, tutil.ApiClientMock):
         self.data = b'xyzzy'
         self.locator = '1271ed5ef305aadabc605b1609e24c52'
 
+    def test_multiple_default_storage_classes_req_header(self):
+        api_mock = self.api_client_mock()
+        api_mock.config.return_value = {
+            'StorageClasses': {
+                'foo': { 'Default': True },
+                'bar': { 'Default': True },
+                'baz': { 'Default': False }
+            }
+        }
+        api_client = self.mock_keep_services(api_mock=api_mock, count=2)
+        keep_client = arvados.KeepClient(api_client=api_client)
+        resp_hdr = {
+            'x-keep-storage-classes-confirmed': 'foo=1, bar=1',
+            'x-keep-replicas-stored': 1
+        }
+        with tutil.mock_keep_responses(self.locator, 200, **resp_hdr) as mock:
+            keep_client.put(self.data, copies=1)
+            req_hdr = mock.responses[0]
+            self.assertIn(
+                'X-Keep-Storage-Classes: bar, foo', req_hdr.getopt(pycurl.HTTPHEADER))
+
     def test_storage_classes_req_header(self):
+        self.assertEqual(
+            self.api_client.config()['StorageClasses'],
+            {'default': {'Default': True}})
         cases = [
             # requested, expected
             [['foo'], 'X-Keep-Storage-Classes: foo'],
             [['bar', 'foo'], 'X-Keep-Storage-Classes: bar, foo'],
-            [[], None],
+            [[], 'X-Keep-Storage-Classes: default'],
+            [None, 'X-Keep-Storage-Classes: default'],
         ]
         for req_classes, expected_header in cases:
             headers = {'x-keep-replicas-stored': 1}
-            if len(req_classes) > 0:
+            if req_classes is None or len(req_classes) == 0:
+                confirmed_hdr = 'default=1'
+            elif len(req_classes) > 0:
                 confirmed_hdr = ', '.join(["{}=1".format(cls) for cls in req_classes])
-                headers.update({'x-keep-storage-classes-confirmed': confirmed_hdr})
+            headers.update({'x-keep-storage-classes-confirmed': confirmed_hdr})
             with tutil.mock_keep_responses(self.locator, 200, **headers) as mock:
                 self.keep_client.put(self.data, copies=1, classes=req_classes)
-                resp = mock.responses[0]
-                if expected_header is not None:
-                    self.assertIn(expected_header, resp.getopt(pycurl.HTTPHEADER))
-                else:
-                    for hdr in resp.getopt(pycurl.HTTPHEADER):
-                        self.assertNotRegex(hdr, r'^X-Keep-Storage-Classes.*')
+                req_hdr = mock.responses[0]
+                self.assertIn(expected_header, req_hdr.getopt(pycurl.HTTPHEADER))
 
     def test_partial_storage_classes_put(self):
         headers = {
@@ -1368,6 +1391,8 @@ class KeepClientAPIErrorTest(unittest.TestCase):
                     return "abc"
                 elif r == "insecure":
                     return False
+                elif r == "config":
+                    return lambda: {}
                 else:
                     raise arvados.errors.KeepReadError()
         keep_client = arvados.KeepClient(api_client=ApiMock(),