From d67f5f30ed86af6a39f788524fe93f64556df85a Mon Sep 17 00:00:00 2001 From: Sri Harsha CH Date: Sat, 18 Nov 2023 12:47:27 +0000 Subject: [PATCH] feat(spanner): update support for autoscaling config field --- google/cloud/spanner_v1/client.py | 4 +++ google/cloud/spanner_v1/instance.py | 21 ++++++++--- tests/system/test_instance_api.py | 47 +++++++++++++++++++++++++ tests/unit/test_instance.py | 54 +++++++++++++++++++++++++++-- 4 files changed, 120 insertions(+), 6 deletions(-) diff --git a/google/cloud/spanner_v1/client.py b/google/cloud/spanner_v1/client.py index 1eb1681d57..972fb93ac9 100644 --- a/google/cloud/spanner_v1/client.py +++ b/google/cloud/spanner_v1/client.py @@ -341,6 +341,10 @@ def instance( :type labels: dict (str -> str) or None :param labels: (Optional) User-assigned labels for this instance. + :type autoscaling_config: + :class:`~google.cloud.spanner_admin_instance_v1.types.AutoscalingConfig` + :param autoscaling_config: (Optional) The autoscaling configuration for this instance. + :rtype: :class:`~google.cloud.spanner_v1.instance.Instance` :returns: an instance owned by this client. """ diff --git a/google/cloud/spanner_v1/instance.py b/google/cloud/spanner_v1/instance.py index 06bcd9e2f4..fcb89108dd 100644 --- a/google/cloud/spanner_v1/instance.py +++ b/google/cloud/spanner_v1/instance.py @@ -379,7 +379,7 @@ def reload(self): self._update_from_pb(instance_pb) - def update(self): + def update(self, fields=None): """Update this instance. See @@ -397,6 +397,10 @@ def update(self): before calling :meth:`update`. + :type fields: Sequence[str] + :param fields: a list of fields to update. Ex: ["config", "display_name", + "processing_units", "labels","autoscaling_config"] + :rtype: :class:`google.api_core.operation.Operation` :returns: an operation instance :raises NotFound: if the instance does not exist @@ -409,12 +413,21 @@ def update(self): node_count=self._node_count, processing_units=self._processing_units, labels=self.labels, + autoscaling_config=self._autoscaling_config, ) + # default field paths to update + paths = [ + "config", + "display_name", + "processing_units", + "labels", + "autoscaling_config", + ] + if fields is not None: + paths = fields # Always update only processing_units, not nodes - field_mask = FieldMask( - paths=["config", "display_name", "processing_units", "labels"] - ) + field_mask = FieldMask(paths=paths) metadata = _metadata_with_prefix(self.name) future = api.update_instance( diff --git a/tests/system/test_instance_api.py b/tests/system/test_instance_api.py index e1623cacee..bcbb7db77e 100644 --- a/tests/system/test_instance_api.py +++ b/tests/system/test_instance_api.py @@ -185,3 +185,50 @@ def test_create_instance_with_autoscaling_config( assert instance == instance_alt assert instance.display_name == instance_alt.display_name assert instance.autoscaling_config == instance_alt.autoscaling_config + + +def test_update_instance_with_autoscaling_config( + not_emulator, + spanner_client, + shared_instance, + shared_instance_id, + instance_operation_timeout, +): + from google.cloud.spanner_admin_instance_v1 import ( + AutoscalingConfig as AutoscalingConfigPB, + ) + + autoscaling_config = AutoscalingConfigPB( + autoscaling_limits=AutoscalingConfigPB.AutoscalingLimits( + min_nodes=1, + max_nodes=2, + ), + autoscaling_targets=AutoscalingConfigPB.AutoscalingTargets( + high_priority_cpu_utilization_percent=65, + storage_utilization_percent=95, + ), + ) + assert shared_instance.autoscaling_config != autoscaling_config + + old_node_count = shared_instance.node_count + shared_instance.autoscaling_config = autoscaling_config + # Update only the autoscaling_config field. This is to ensure that + # node_count or processing_unit field is not considered during update, + # which might otherwise throw an error. + operation = shared_instance.update(fields=["autoscaling_config"]) + + # We want to make sure the operation completes. + operation.result(instance_operation_timeout) # raises on failure / timeout. + + # Create a new instance instance and reload it. + instance_alt = spanner_client.instance(shared_instance_id, None) + assert instance_alt.autoscaling_config != autoscaling_config + + instance_alt.reload() + assert instance_alt.autoscaling_config == autoscaling_config + + # Make sure to put the instance back the way it was for the + # other test cases. + shared_instance.node_count = old_node_count + shared_instance.autoscaling_config = None + shared_instance.update() diff --git a/tests/unit/test_instance.py b/tests/unit/test_instance.py index daf07c957b..e1e419c3ba 100644 --- a/tests/unit/test_instance.py +++ b/tests/unit/test_instance.py @@ -38,7 +38,13 @@ class TestInstance(unittest.TestCase): DATABASE_ID = "database_id" DATABASE_NAME = "%s/databases/%s" % (INSTANCE_NAME, DATABASE_ID) LABELS = {"test": "true"} - FIELD_MASK = ["config", "display_name", "processing_units", "labels"] + FIELD_MASK = [ + "config", + "display_name", + "processing_units", + "labels", + "autoscaling_config", + ] def _getTargetClass(self): from google.cloud.spanner_v1.instance import Instance @@ -558,7 +564,14 @@ def test_update_success_with_processing_units(self): instance, field_mask, metadata = api._updated_instance self.assertEqual( - field_mask.paths, ["config", "display_name", "processing_units", "labels"] + field_mask.paths, + [ + "config", + "display_name", + "processing_units", + "labels", + "autoscaling_config", + ], ) self.assertEqual(instance.name, self.INSTANCE_NAME) self.assertEqual(instance.config, self.CONFIG_NAME) @@ -567,6 +580,43 @@ def test_update_success_with_processing_units(self): self.assertEqual(instance.labels, self.LABELS) self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)]) + def test_update_success_with_autoscaling_config(self): + from google.cloud.spanner_admin_instance_v1 import ( + AutoscalingConfig as AutoscalingConfigPB, + ) + + autoscaling_config = AutoscalingConfigPB( + autoscaling_limits=AutoscalingConfigPB.AutoscalingLimits( + min_nodes=1, + max_nodes=2, + ), + autoscaling_targets=AutoscalingConfigPB.AutoscalingTargets( + high_priority_cpu_utilization_percent=65, + storage_utilization_percent=95, + ), + ) + + op_future = _FauxOperationFuture() + client = _Client(self.PROJECT) + api = client.instance_admin_api = _FauxInstanceAdminAPI( + _update_instance_response=op_future + ) + instance = self._make_one( + self.INSTANCE_ID, + client, + autoscaling_config=autoscaling_config, + ) + + future = instance.update(fields=["autoscaling_config"]) + + self.assertIs(future, op_future) + + instance, field_mask, metadata = api._updated_instance + self.assertEqual(field_mask.paths, ["autoscaling_config"]) + self.assertEqual(instance.name, self.INSTANCE_NAME) + self.assertEqual(instance.autoscaling_config, autoscaling_config) + self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)]) + def test_delete_grpc_error(self): from google.api_core.exceptions import Unknown