Skip to content

Commit

Permalink
feat: drop database protection
Browse files Browse the repository at this point in the history
  • Loading branch information
aayushimalik committed Mar 16, 2023
1 parent 35953e5 commit e4163d0
Show file tree
Hide file tree
Showing 3 changed files with 116 additions and 0 deletions.
53 changes: 53 additions & 0 deletions google/cloud/spanner_v1/database.py
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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)
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.
Expand Down
34 changes: 34 additions & 0 deletions tests/system/test_database_api.py
Expand Up @@ -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()
29 changes: 29 additions & 0 deletions tests/unit/test_database.py
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand Down

0 comments on commit e4163d0

Please sign in to comment.