feat: add support for updateDatabase in Cloud Spanner (#914)

* feat: drop database protection

Co-authored-by: Rajat Bhatta <[email protected]>

Author: aayushimalik

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
@@ -127,6 +128,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
@@ -141,6 +145,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
@@ -159,6 +164,8 @@ def __init__(
         self._database_dialect = database_dialect
         self._database_role = database_role
         self._route_to_leader_enabled = self._instance._client.route_to_leader_enabled
+        self._enable_drop_protection = enable_drop_protection
+        self._reconciling = False
 
         if pool is None:
             pool = BurstyPool(database_role=database_role)
@@ -332,6 +339,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.
@@ -461,14 +491,16 @@ 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
+        self._reconciling = response.reconciling
 
     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
@@ -492,6 +524,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):

samples/samples/snippets.py

@@ -196,6 +196,27 @@ def create_database(instance_id, database_id):
 # [END spanner_create_database]
 
 
+# [START spanner_update_database]
+def update_database(instance_id, database_id):
+    """Updates the drop protection setting for a database."""
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+
+    db = instance.database(database_id)
+    db.enable_drop_protection = True
+
+    operation = db.update(["enable_drop_protection"])
+
+    print("Waiting for update operation for {} to complete...".format(
+      db.name))
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print("Updated database {}.".format(db.name))
+
+
+# [END spanner_update_database]
+
+
 # [START spanner_create_database_with_encryption_key]
 def create_database_with_encryption_key(instance_id, database_id, kms_key_name):
     """Creates a database with tables using a Customer Managed Encryption Key (CMEK)."""

samples/samples/snippets_test.py

@@ -154,6 +154,19 @@ def test_create_instance_with_processing_units(capsys, lci_instance_id):
     retry_429(instance.delete)()
 
 
+def test_update_database(capsys, instance_id, sample_database):
+    snippets.update_database(
+        instance_id, sample_database.database_id
+    )
+    out, _ = capsys.readouterr()
+    assert "Updated database {}.".format(sample_database.name) in out
+
+    # Cleanup
+    sample_database.enable_drop_protection = False
+    op = sample_database.update(["enable_drop_protection"])
+    op.result()
+
+
 def test_create_database_with_encryption_config(
   capsys, instance_id, cmek_database_id, kms_key_name
 ):

tests/system/test_database_api.py

@@ -562,3 +562,41 @@ 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
 
@@ -760,6 +762,8 @@ def test_reload_success(self):
             encryption_config=encryption_config,
             encryption_info=encryption_info,
             default_leader=default_leader,
+            reconciling=True,
+            enable_drop_protection=True,
         )
         api.get_database.return_value = db_pb
         instance = _Instance(self.INSTANCE_NAME, client=client)
@@ -776,6 +780,8 @@ def test_reload_success(self):
         self.assertEqual(database._encryption_config, encryption_config)
         self.assertEqual(database._encryption_info, encryption_info)
         self.assertEqual(database._default_leader, default_leader)
+        self.assertEqual(database._reconciling, True)
+        self.assertEqual(database._enable_drop_protection, True)
 
         api.get_database_ddl.assert_called_once_with(
             database=self.DATABASE_NAME,
@@ -892,6 +898,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