feat: add support for updateDatabase in Cloud Spanner

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 enable_drop_protection: boolean
+    :param enable_drop_protection: (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,
+        enable_drop_protection=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._enable_drop_protection = enable_drop_protection
+        self._reconciling = False
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -328,6 +335,29 @@ 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 enable_drop_protection(self):
+        """Whether the database has drop protection enabled.
+
+        :rtype: boolean
+        :returns: a boolean representing whether the database has drop
+            protection enabled
+        """
+        return self._enable_drop_protection
+
+    @enable_drop_protection.setter
+    def enable_drop_protection(self, value):
+        self._enable_drop_protection = value
+
     @property
     def logger(self):
         """Logger used by the database.
@@ -457,14 +487,15 @@ def reload(self):
         self._encryption_info = response.encryption_info
         self._default_leader = response.default_leader
         self._database_dialect = response.database_dialect
+        self._enable_drop_protection = response.enable_drop_protection
 
     def update_ddl(self, ddl_statements, operation_id=""):
         """Update DDL for this database.
 
         Apply any configured schema from :attr:`ddl_statements`.
 
         See
-        https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.UpdateDatabase
+        https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.UpdateDatabaseDdl
 
         :type ddl_statements: Sequence[str]
         :param ddl_statements: a list of DDL statements to use on this database
@@ -488,6 +519,46 @@ def update_ddl(self, ddl_statements, operation_id=""):
         future = api.update_database_ddl(request=request, metadata=metadata)
         return future
 
+    def update(self, fields):
+        """Update this database.
+
+        See
+        https://cloud.google.com/spanner/reference/rpc/google.spanner.admin.database.v1#google.spanner.admin.database.v1.DatabaseAdmin.UpdateDatabase
+
+        .. note::
+
+            Updates the specified fields of a Cloud Spanner database. Currently,
+            only the `enable_drop_protection` field supports updates. To change
+            this value before updating, set it via
+
+            .. code:: python
+
+                database.enable_drop_protection = True
+
+           before calling :meth:`update`.
+
+        :type fields: Sequence[str]
+        :param fields: a list of fields to 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._enable_drop_protection
+        )
+
+        # Only support updating drop protection for now.
+        field_mask = FieldMask(paths=fields)
+        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.
 

google/cloud/spanner_v1/instance.py

@@ -432,6 +432,7 @@ def database(
         encryption_config=None,
         database_dialect=DatabaseDialect.DATABASE_DIALECT_UNSPECIFIED,
         database_role=None,
+        enable_drop_protection=False,
     ):
         """Factory to create a database within this instance.
 
@@ -467,6 +468,10 @@ def database(
         :param database_dialect:
             (Optional) database dialect for the database
 
+        :type enable_drop_protection: boolean
+        :param enable_drop_protection: (Optional) Represents whether the database
+            has drop protection enabled or not.
+
         :rtype: :class:`~google.cloud.spanner_v1.database.Database`
         :returns: a database owned by this instance.
         """
@@ -479,6 +484,7 @@ def database(
             encryption_config=encryption_config,
             database_dialect=database_dialect,
             database_role=database_role,
+            enable_drop_protection=enable_drop_protection,
         )
 
     def list_databases(self, page_size=None):

tests/system/test_database_api.py

@@ -562,3 +562,42 @@ 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_success(
+    not_emulator, shared_database, shared_instance, database_operation_timeout
+):
+    old_protection = shared_database.enable_drop_protection
+    new_protection = True
+    shared_database.enable_drop_protection = new_protection
+    operation = shared_database.update(["enable_drop_protection"])
+
+    # 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.enable_drop_protection != new_protection
+
+    database_alt.reload()
+    assert database_alt.enable_drop_protection == new_protection
+
+    with pytest.raises(exceptions.FailedPrecondition):
+        database_alt.drop()
+
+    with pytest.raises(exceptions.FailedPrecondition):
+        shared_instance.delete()
+
+    # Make sure to put the database back the way it was for the
+    # other test cases.
+    shared_database.enable_drop_protection = old_protection
+    shared_database.update(["enable_drop_protection"])
+
+def test_update_database_invalid(
+    not_emulator, shared_database
+):
+    shared_database.enable_drop_protection = True
+
+    # Empty `fields` is not supported.
+    with pytest.raises(exceptions.InvalidArgument):
+        shared_database.update([])

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,32 @@ 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, enable_drop_protection=True, pool=pool
+        )
+
+        future = database.update(["enable_drop_protection"])
+
+        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