From 7e75562f0e97cb21dbf2a94be92e356a9b17b54d Mon Sep 17 00:00:00 2001
From: Taylor McKinnon <mail@tmacs.space>
Date: Thu, 9 Nov 2023 13:54:06 -0800
Subject: [PATCH] request bucket attributes from bucketd

---
 lib/reindex/s3_bucketd.py | 40 ++++++++++++++++++++++++++++-----------
 1 file changed, 29 insertions(+), 11 deletions(-)

diff --git a/lib/reindex/s3_bucketd.py b/lib/reindex/s3_bucketd.py
index f0fee300..91051544 100644
--- a/lib/reindex/s3_bucketd.py
+++ b/lib/reindex/s3_bucketd.py
@@ -65,7 +65,8 @@ def __init__(self, bucket):
 class BucketDClient:
 
     '''Performs Listing calls against bucketd'''
-    __url_format = '{addr}/default/bucket/{bucket}'
+    __url_attribute_format = '{addr}/default/attributes/{bucket}'
+    __url_bucket_format = '{addr}/default/bucket/{bucket}'
     __headers = {"x-scal-request-uids": "utapi-reindex-list-buckets"}
 
     def __init__(self, bucketd_addr=None, max_retries=2, only_latest_when_locked=False):
@@ -103,7 +104,7 @@ def _list_bucket(self, bucket, **kwargs):
         parameters value. On the first request the function will be called with
         `None` and should return its initial value. Return `None` for the param to be excluded.
         '''
-        url = self.__url_format.format(addr=self._bucketd_addr, bucket=bucket)
+        url = self.__url_bucket_format.format(addr=self._bucketd_addr, bucket=bucket)
         static_params = {k: v for k, v in kwargs.items() if not callable(v)}
         dynamic_params = {k: v for k, v in kwargs.items() if callable(v)}
         is_truncated = True # Set to True for first loop
@@ -137,6 +138,26 @@ def _list_bucket(self, bucket, **kwargs):
             else:
                 is_truncated = len(payload) > 0
 
+    def _get_bucket_attributes(self, name):
+        url = self.__url_attribute_format.format(addr=self._bucketd_addr, bucket=name)
+        try:
+            resp = self._do_req(url)
+            if resp.status_code == 200:
+                return resp.json()
+            else:
+                _log.error('Error getting bucket attributes bucket:%s status_code:%s'%(name, resp.status_code))
+                raise InvalidListing(name)
+        except ValueError as e:
+            _log.exception(e)
+            _log.error('Invalid attributes response body! bucket:%s'%name)
+        except MaxRetriesReached:
+            _log.error('Max retries reached getting bucket attributes bucket:%s'%name)
+            raise
+        except Exception as e:
+            _log.exception(e)
+            _log.error('Unhandled exception getting bucket attributes bucket:%s'%name)
+            raise
+
     def list_buckets(self, name = None):
 
         def get_next_marker(p):
@@ -153,17 +174,14 @@ def get_next_marker(p):
             buckets = []
             for result in payload['Contents']:
                 match = re.match("(\w+)..\|..(\w+.*)", result['key'])
-                bucket_md = json.loads(result['value'])
-                obj_lock_enabled = bucket_md.get('objectLockEnabled', False)
-                # obj_lock_config = bucket_md.get('objectLockConfiguration', {})
-                # obj_lock_mode = obj_lock_config.get('rule', {}).get('mode')
-                # if mode == 'GOVERNANCE' or mode == 'COMPLIANCE':
-                bucket = Bucket(*match.groups(), object_locked=obj_lock_enabled)
-                if obj_lock_enabled:
-                    bucket = bucket._replace(object_locked=True)
+                bucket = Bucket(*match.groups())
                 if name is None or bucket.name == name:
+                    # We need to get the attributes for each bucket to determine if it is locked
+                    if self._only_latest_when_locked:
+                        bucket_attrs = self._get_bucket_attributes(bucket.name)
+                        obj_lock_enabled = bucket_attrs.get('objectLockEnabled', False)
+                        bucket = bucket._replace(object_locked=obj_lock_enabled)
                     buckets.append(bucket)
-                print(result)
 
             if buckets:
                 yield buckets