feat: drop database protection

aayushimalik · aayushimalik · commit e4163d03f525 · 2023-03-16T11:19:12.000Z
diff --git a/google/cloud/spanner_v1/database.py b/google/cloud/spanner_v1/database.py
@@ -29,6 +29,7 @@
 from google.api_core import gapic_v1
 from google.iam.v1 import iam_policy_pb2
 from google.iam.v1 import options_pb2
+from google.protobuf.field_mask_pb2 import FieldMask
 
 from google.cloud.spanner_admin_database_v1 import CreateDatabaseRequest
 from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
@@ -124,6 +125,9 @@ class Database(object):
         (Optional) database dialect for the database
     :type database_role: str or None
     :param database_role: (Optional) user-assigned database_role for the session.
+    :type drop_protection_enabled: boolean
+    :param drop_protection_enabled: (Optional) Represents whether the database
+        has drop protection enabled or not.
     """
 
     _spanner_api = None
@@ -138,6 +142,7 @@ def __init__(
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
+        drop_protection_enabled=False,
     ):
         self.database_id = database_id
         self._instance = instance
@@ -155,6 +160,8 @@ def __init__(
         self._encryption_config = encryption_config
         self._database_dialect = database_dialect
         self._database_role = database_role
+        self.drop_protection_enabled = drop_protection_enabled
+        self._reconciling = False
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -328,6 +335,15 @@ def database_role(self):
         """
         return self._database_role
 
+    @property
+    def reconciling(self):
+        """Whether the database is currently reconciling.
+
+        :rtype: boolean
+        :returns: a boolean representing whether the database is reconciling
+        """
+        return self._reconciling
+
     @property
     def logger(self):
         """Logger used by the database.
@@ -457,6 +473,7 @@ def reload(self):
         self._encryption_info = response.encryption_info
         self._default_leader = response.default_leader
         self._database_dialect = response.database_dialect
+        self.drop_protection_enabled = response.enable_drop_protection
 
     def update_ddl(self, ddl_statements, operation_id=""):
         """Update DDL for this database.
@@ -488,6 +505,42 @@ def update_ddl(self, ddl_statements, operation_id=""):
         future = api.update_database_ddl(request=request, metadata=metadata)
         return future
 
+    def update(self):
+        """Update this database.
+
+        See
+        TODO: insert documentation once ready
+
+        .. note::
+
+            Updates the `drop_protection_enabled` field. To change this value
+            before updating, set it via
+
+            .. code:: python
+
+                database.drop_protection_enabled = True
+
+           before calling :meth:`update`.
+
+        :rtype: :class:`google.api_core.operation.Operation`
+        :returns: an operation instance
+        :raises NotFound: if the database does not exist
+        """
+        api = self._instance._client.database_admin_api
+        database_pb = DatabasePB(
+            name=self.name, enable_drop_protection=self.drop_protection_enabled
+        )
+
+        # Only support updating drop protection for now.
+        field_mask = FieldMask(paths=["enable_drop_protection"])
+        metadata = _metadata_with_prefix(self.name)
+
+        future = api.update_database(
+            database=database_pb, update_mask=field_mask, metadata=metadata
+        )
+
+        return future
+
     def drop(self):
         """Drop this database.
 
diff --git a/tests/system/test_database_api.py b/tests/system/test_database_api.py
@@ -562,3 +562,37 @@ def _unit_of_work(transaction, name):
         rows = list(after.read(sd.COUNTERS_TABLE, sd.COUNTERS_COLUMNS, sd.ALL))
 
     assert len(rows) == 2
+
+def test_update_database(
+    not_emulator,
+    shared_database,
+    shared_instance,
+    database_operation_timeout
+):
+    old_protection = shared_database.drop_protection_enabled
+    new_protection = True
+    shared_database.drop_protection_enabled = new_protection
+    operation = shared_database.update()
+
+    # We want to make sure the operation completes.
+    operation.result(database_operation_timeout)  # raises on failure / timeout.
+
+    # Create a new database instance and reload it.
+    database_alt = shared_instance.database(shared_database.name.split("/")[-1])
+    assert database_alt.drop_protection_enabled != new_protection
+
+    database_alt.reload()
+    assert database_alt.drop_protection_enabled == new_protection
+
+    with pytest.raises(exceptions.FailedPrecondition,
+                       match='`enable_drop_protection` setting'):
+        database_alt.drop()
+
+    with pytest.raises(exceptions.FailedPrecondition,
+                       match='drop protection enabled'):
+        shared_instance.delete()
+
+    # Make sure to put the database back the way it was for the
+    # other test cases.
+    shared_database.drop_protection_enabled = old_protection
+    shared_database.update()
diff --git a/tests/unit/test_database.py b/tests/unit/test_database.py
@@ -17,8 +17,10 @@
 
 import mock
 from google.api_core import gapic_v1
+from google.cloud.spanner_admin_database_v1 import Database as DatabasePB
 from google.cloud.spanner_v1.param_types import INT64
 from google.api_core.retry import Retry
+from google.protobuf.field_mask_pb2 import FieldMask
 
 from google.cloud.spanner_v1 import RequestOptions
 
@@ -877,6 +879,33 @@ def test_update_ddl_w_operation_id(self):
             metadata=[("google-cloud-resource-prefix", database.name)],
         )
 
+    def test_update_success(self):
+        op_future = object()
+        client = _Client()
+        api = client.database_admin_api = self._make_database_admin_api()
+        api.update_database.return_value = op_future
+
+        instance = _Instance(self.INSTANCE_NAME, client=client)
+        pool = _Pool()
+        database = self._make_one(
+            self.DATABASE_ID, instance, drop_protection_enabled=True, pool=pool
+        )
+
+        future = database.update()
+
+        self.assertIs(future, op_future)
+
+        expected_database = DatabasePB(name=database.name,
+                                       enable_drop_protection=True)
+
+        field_mask = FieldMask(paths=["enable_drop_protection"])
+
+        api.update_database.assert_called_once_with(
+            database=expected_database,
+            update_mask=field_mask,
+            metadata=[("google-cloud-resource-prefix", database.name)],
+        )
+
     def test_drop_grpc_error(self):
         from google.api_core.exceptions import Unknown
 
