Add 'Bucket.lock_retention_policy' API wrapper. (#448)

Toward #445.
googleapis · Sep 7, 2018 · 2f22ae1 · 2f22ae1
1 parent 3c7054e
commit 2f22ae1
Show file tree

Hide file tree

Showing 2 changed files with 142 additions and 0 deletions.
diff --git a/storage/google/cloud/storage/bucket.py b/storage/google/cloud/storage/bucket.py
@@ -1778,3 +1778,37 @@ def generate_upload_policy(
         }
 
         return fields
+
+    def lock_retention_policy(self, client=None):
+        """Lock the bucket's retention policy.
+
+        :raises ValueError:
+            if the bucket has no metageneration (i.e., new or never reloaded);
+            if the bucket has no retention policy assigned;
+            if the bucket's retention policy is already locked.
+        """
+        if 'metageneration' not in self._properties:
+            raise ValueError(
+                "Bucket has no retention policy assigned: try 'reload'?")
+
+        policy = self._properties.get('retentionPolicy')
+
+        if policy is None:
+            raise ValueError(
+                "Bucket has no retention policy assigned: try 'reload'?")
+
+        if policy.get('isLocked'):
+            raise ValueError("Bucket's retention policy is already locked.")
+
+        client = self._require_client(client)
+
+        query_params = {'metageneration': self.metageneration}
+
+        if self.user_project is not None:
+            query_params['userProject'] = self.user_project
+
+        path = '/b/{}/lockRetentionPolicy'.format(self.name)
+        api_response = client._connection.api_request(
+            method='POST', path=path, query_params=query_params,
+            _target_object=self)
+        self._set_properties(api_response)
diff --git a/storage/tests/unit/test_bucket.py b/storage/tests/unit/test_bucket.py
@@ -2408,6 +2408,114 @@ def test_generate_upload_policy_bad_credentials(self):
         with self.assertRaises(AttributeError):
             bucket.generate_upload_policy([])
 
+    def test_lock_retention_policy_no_policy_set(self):
+        credentials = object()
+        connection = _Connection()
+        connection.credentials = credentials
+        client = _Client(connection)
+        name = 'name'
+        bucket = self._make_one(client=client, name=name)
+        bucket._properties['metageneration'] = 1234
+
+        with self.assertRaises(ValueError):
+            bucket.lock_retention_policy()
+
+    def test_lock_retention_policy_no_metageneration(self):
+        credentials = object()
+        connection = _Connection()
+        connection.credentials = credentials
+        client = _Client(connection)
+        name = 'name'
+        bucket = self._make_one(client=client, name=name)
+        bucket._properties['retentionPolicy'] = {
+            'effectiveTime': '2018-03-01T16:46:27.123456Z',
+            'retentionPeriod': 86400 * 100,  # 100 days
+        }
+
+        with self.assertRaises(ValueError):
+            bucket.lock_retention_policy()
+
+    def test_lock_retention_policy_already_locked(self):
+        credentials = object()
+        connection = _Connection()
+        connection.credentials = credentials
+        client = _Client(connection)
+        name = 'name'
+        bucket = self._make_one(client=client, name=name)
+        bucket._properties['metageneration'] = 1234
+        bucket._properties['retentionPolicy'] = {
+            'effectiveTime': '2018-03-01T16:46:27.123456Z',
+            'isLocked': True,
+            'retentionPeriod': 86400 * 100,  # 100 days
+        }
+
+        with self.assertRaises(ValueError):
+            bucket.lock_retention_policy()
+
+    def test_lock_retention_policy_ok(self):
+        name = 'name'
+        response = {
+            'name': name,
+            'metageneration': 1235,
+            'retentionPolicy': {
+                'effectiveTime': '2018-03-01T16:46:27.123456Z',
+                'isLocked': True,
+                'retentionPeriod': 86400 * 100,  # 100 days
+            },
+        }
+        credentials = object()
+        connection = _Connection(response)
+        connection.credentials = credentials
+        client = _Client(connection)
+        bucket = self._make_one(client=client, name=name)
+        bucket._properties['metageneration'] = 1234
+        bucket._properties['retentionPolicy'] = {
+            'effectiveTime': '2018-03-01T16:46:27.123456Z',
+            'retentionPeriod': 86400 * 100,  # 100 days
+        }
+
+        bucket.lock_retention_policy()
+
+        kw, = connection._requested
+        self.assertEqual(kw['method'], 'POST')
+        self.assertEqual(kw['path'], '/b/{}/lockRetentionPolicy'.format(name))
+        self.assertEqual(kw['query_params'], {'metageneration': 1234})
+
+    def test_lock_retention_policy_w_user_project(self):
+        name = 'name'
+        user_project = 'user-project-123'
+        response = {
+            'name': name,
+            'metageneration': 1235,
+            'retentionPolicy': {
+                'effectiveTime': '2018-03-01T16:46:27.123456Z',
+                'isLocked': True,
+                'retentionPeriod': 86400 * 100,  # 100 days
+            },
+        }
+        credentials = object()
+        connection = _Connection(response)
+        connection.credentials = credentials
+        client = _Client(connection)
+        bucket = self._make_one(
+            client=client, name=name, user_project=user_project)
+        bucket._properties['metageneration'] = 1234
+        bucket._properties['retentionPolicy'] = {
+            'effectiveTime': '2018-03-01T16:46:27.123456Z',
+            'retentionPeriod': 86400 * 100,  # 100 days
+        }
+
+        bucket.lock_retention_policy()
+
+        kw, = connection._requested
+        self.assertEqual(kw['method'], 'POST')
+        self.assertEqual(kw['path'], '/b/{}/lockRetentionPolicy'.format(name))
+        self.assertEqual(
+            kw['query_params'], {
+                'metageneration': 1234,
+                'userProject': user_project,
+            })
+
 
 class _Connection(object):
     _delete_bucket = False