Skip to content

Commit

Permalink
feat: add support for low-cost instances (#313)
Browse files Browse the repository at this point in the history
* Add LCI implementation

* Update google/cloud/spanner_v1/instance.py

Co-authored-by: larkee <31196561+larkee@users.noreply.github.com>

* Fix docstring format

* Update google/cloud/spanner_v1/instance.py

Co-authored-by: larkee <31196561+larkee@users.noreply.github.com>

Co-authored-by: larkee <31196561+larkee@users.noreply.github.com>
  • Loading branch information
zoercai and larkee authored Jun 23, 2021
1 parent b8b24e1 commit 44aa7cc
Show file tree
Hide file tree
Showing 5 changed files with 177 additions and 18 deletions.
9 changes: 7 additions & 2 deletions google/cloud/spanner_v1/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,6 @@
from google.cloud.client import ClientWithProject
from google.cloud.spanner_v1 import __version__
from google.cloud.spanner_v1._helpers import _merge_query_options, _metadata_with_prefix
from google.cloud.spanner_v1.instance import DEFAULT_NODE_COUNT
from google.cloud.spanner_v1.instance import Instance
from google.cloud.spanner_v1 import ExecuteSqlRequest
from google.cloud.spanner_admin_instance_v1 import ListInstanceConfigsRequest
Expand Down Expand Up @@ -294,8 +293,9 @@ def instance(
instance_id,
configuration_name=None,
display_name=None,
node_count=DEFAULT_NODE_COUNT,
node_count=None,
labels=None,
processing_units=None,
):
"""Factory to create a instance associated with this client.
Expand All @@ -320,6 +320,10 @@ def instance(
:param node_count: (Optional) The number of nodes in the instance's
cluster; used to set up the instance's cluster.
:type processing_units: int
:param processing_units: (Optional) The number of processing units
allocated to this instance.
:type labels: dict (str -> str) or None
:param labels: (Optional) User-assigned labels for this instance.
Expand All @@ -334,6 +338,7 @@ def instance(
display_name,
self._emulator_host,
labels,
processing_units,
)

def list_instances(self, filter_="", page_size=None):
Expand Down
83 changes: 74 additions & 9 deletions google/cloud/spanner_v1/instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
"""User friendly container for Cloud Spanner Instance."""

import google.api_core.operation
from google.api_core.exceptions import InvalidArgument
import re

from google.cloud.spanner_admin_instance_v1 import Instance as InstancePB
Expand All @@ -41,6 +42,7 @@
)

DEFAULT_NODE_COUNT = 1
PROCESSING_UNITS_PER_NODE = 1000

_OPERATION_METADATA_MESSAGES = (
backup.Backup,
Expand Down Expand Up @@ -95,6 +97,10 @@ class Instance(object):
:type node_count: int
:param node_count: (Optional) Number of nodes allocated to the instance.
:type processing_units: int
:param processing_units: (Optional) The number of processing units
allocated to this instance.
:type display_name: str
:param display_name: (Optional) The display name for the instance in the
Cloud Console UI. (Must be between 4 and 30
Expand All @@ -110,15 +116,29 @@ def __init__(
instance_id,
client,
configuration_name=None,
node_count=DEFAULT_NODE_COUNT,
node_count=None,
display_name=None,
emulator_host=None,
labels=None,
processing_units=None,
):
self.instance_id = instance_id
self._client = client
self.configuration_name = configuration_name
self.node_count = node_count
if node_count is not None and processing_units is not None:
if processing_units != node_count * PROCESSING_UNITS_PER_NODE:
raise InvalidArgument(
"Only one of node count and processing units can be set."
)
if node_count is None and processing_units is None:
self._node_count = DEFAULT_NODE_COUNT
self._processing_units = DEFAULT_NODE_COUNT * PROCESSING_UNITS_PER_NODE
elif node_count is not None:
self._node_count = node_count
self._processing_units = node_count * PROCESSING_UNITS_PER_NODE
else:
self._processing_units = processing_units
self._node_count = processing_units // PROCESSING_UNITS_PER_NODE
self.display_name = display_name or instance_id
self.emulator_host = emulator_host
if labels is None:
Expand All @@ -134,7 +154,8 @@ def _update_from_pb(self, instance_pb):
raise ValueError("Instance protobuf does not contain display_name")
self.display_name = instance_pb.display_name
self.configuration_name = instance_pb.config
self.node_count = instance_pb.node_count
self._node_count = instance_pb.node_count
self._processing_units = instance_pb.processing_units
self.labels = instance_pb.labels

@classmethod
Expand Down Expand Up @@ -190,6 +211,44 @@ def name(self):
"""
return self._client.project_name + "/instances/" + self.instance_id

@property
def processing_units(self):
"""Processing units used in requests.
:rtype: int
:returns: The number of processing units allocated to this instance.
"""
return self._processing_units

@processing_units.setter
def processing_units(self, value):
"""Sets the processing units for requests. Affects node_count.
:param value: The number of processing units allocated to this instance.
"""
self._processing_units = value
self._node_count = value // PROCESSING_UNITS_PER_NODE

@property
def node_count(self):
"""Node count used in requests.
:rtype: int
:returns:
The number of nodes in the instance's cluster;
used to set up the instance's cluster.
"""
return self._node_count

@node_count.setter
def node_count(self, value):
"""Sets the node count for requests. Affects processing_units.
:param value: The number of nodes in the instance's cluster.
"""
self._node_count = value
self._processing_units = value * PROCESSING_UNITS_PER_NODE

def __eq__(self, other):
if not isinstance(other, self.__class__):
return NotImplemented
Expand Down Expand Up @@ -218,7 +277,8 @@ def copy(self):
self.instance_id,
new_client,
self.configuration_name,
node_count=self.node_count,
node_count=self._node_count,
processing_units=self._processing_units,
display_name=self.display_name,
)

Expand Down Expand Up @@ -250,7 +310,7 @@ def create(self):
name=self.name,
config=self.configuration_name,
display_name=self.display_name,
node_count=self.node_count,
processing_units=self._processing_units,
labels=self.labels,
)
metadata = _metadata_with_prefix(self.name)
Expand Down Expand Up @@ -306,8 +366,8 @@ def update(self):
.. note::
Updates the ``display_name``, ``node_count`` and ``labels``. To change those
values before updating, set them via
Updates the ``display_name``, ``node_count``, ``processing_units``
and ``labels``. To change those values before updating, set them via
.. code:: python
Expand All @@ -325,10 +385,15 @@ def update(self):
name=self.name,
config=self.configuration_name,
display_name=self.display_name,
node_count=self.node_count,
node_count=self._node_count,
processing_units=self._processing_units,
labels=self.labels,
)
field_mask = FieldMask(paths=["config", "display_name", "node_count", "labels"])

# Always update only processing_units, not nodes
field_mask = FieldMask(
paths=["config", "display_name", "processing_units", "labels"]
)
metadata = _metadata_with_prefix(self.name)

future = api.update_instance(
Expand Down
29 changes: 29 additions & 0 deletions tests/system/test_system.py
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,35 @@ def test_create_instance(self):
self.assertEqual(instance, instance_alt)
self.assertEqual(instance.display_name, instance_alt.display_name)

@unittest.skipIf(USE_EMULATOR, "Skipping LCI tests")
@unittest.skipUnless(CREATE_INSTANCE, "Skipping instance creation")
def test_create_instance_with_processing_nodes(self):
ALT_INSTANCE_ID = "new" + unique_resource_id("-")
PROCESSING_UNITS = 5000
instance = Config.CLIENT.instance(
instance_id=ALT_INSTANCE_ID,
configuration_name=Config.INSTANCE_CONFIG.name,
processing_units=PROCESSING_UNITS,
)
operation = instance.create()
# Make sure this instance gets deleted after the test case.
self.instances_to_delete.append(instance)

# We want to make sure the operation completes.
operation.result(
SPANNER_OPERATION_TIMEOUT_IN_SECONDS
) # raises on failure / timeout.

# Create a new instance instance and make sure it is the same.
instance_alt = Config.CLIENT.instance(
ALT_INSTANCE_ID, Config.INSTANCE_CONFIG.name
)
instance_alt.reload()

self.assertEqual(instance, instance_alt)
self.assertEqual(instance.display_name, instance_alt.display_name)
self.assertEqual(instance.processing_units, instance_alt.processing_units)

@unittest.skipIf(USE_EMULATOR, "Skipping updating instance")
def test_update_instance(self):
OLD_DISPLAY_NAME = Config.INSTANCE.display_name
Expand Down
3 changes: 3 additions & 0 deletions tests/unit/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ class TestClient(unittest.TestCase):
INSTANCE_NAME = "%s/instances/%s" % (PATH, INSTANCE_ID)
DISPLAY_NAME = "display-name"
NODE_COUNT = 5
PROCESSING_UNITS = 5000
LABELS = {"test": "true"}
TIMEOUT_SECONDS = 80

Expand Down Expand Up @@ -580,6 +581,7 @@ def test_list_instances(self):
config=self.CONFIGURATION_NAME,
display_name=self.DISPLAY_NAME,
node_count=self.NODE_COUNT,
processing_units=self.PROCESSING_UNITS,
)
]
)
Expand All @@ -597,6 +599,7 @@ def test_list_instances(self):
self.assertEqual(instance.config, self.CONFIGURATION_NAME)
self.assertEqual(instance.display_name, self.DISPLAY_NAME)
self.assertEqual(instance.node_count, self.NODE_COUNT)
self.assertEqual(instance.processing_units, self.PROCESSING_UNITS)

expected_metadata = (
("google-cloud-resource-prefix", client.project_name),
Expand Down
71 changes: 64 additions & 7 deletions tests/unit/test_instance.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ class TestInstance(unittest.TestCase):
LOCATION = "projects/" + PROJECT + "/locations/" + CONFIG_NAME
DISPLAY_NAME = "display_name"
NODE_COUNT = 5
PROCESSING_UNITS = 5000
OP_ID = 8915
OP_NAME = "operations/projects/%s/instances/%soperations/%d" % (
PROJECT,
Expand All @@ -39,6 +40,7 @@ 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"]

def _getTargetClass(self):
from google.cloud.spanner_v1.instance import Instance
Expand Down Expand Up @@ -230,7 +232,7 @@ def test_create_already_exists(self):
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.INSTANCE_ID)
self.assertEqual(instance.node_count, 1)
self.assertEqual(instance.processing_units, 1000)
self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)])

def test_create_success(self):
Expand Down Expand Up @@ -258,7 +260,36 @@ def test_create_success(self):
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.DISPLAY_NAME)
self.assertEqual(instance.node_count, self.NODE_COUNT)
self.assertEqual(instance.processing_units, self.PROCESSING_UNITS)
self.assertEqual(instance.labels, self.LABELS)
self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)])

def test_create_with_processing_units(self):
op_future = _FauxOperationFuture()
client = _Client(self.PROJECT)
api = client.instance_admin_api = _FauxInstanceAdminAPI(
_create_instance_response=op_future
)
instance = self._make_one(
self.INSTANCE_ID,
client,
configuration_name=self.CONFIG_NAME,
display_name=self.DISPLAY_NAME,
processing_units=self.PROCESSING_UNITS,
labels=self.LABELS,
)

future = instance.create()

self.assertIs(future, op_future)

(parent, instance_id, instance, metadata) = api._created_instance
self.assertEqual(parent, self.PARENT)
self.assertEqual(instance_id, self.INSTANCE_ID)
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.DISPLAY_NAME)
self.assertEqual(instance.processing_units, self.PROCESSING_UNITS)
self.assertEqual(instance.labels, self.LABELS)
self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)])

Expand Down Expand Up @@ -389,9 +420,7 @@ def test_update_not_found(self):
instance.update()

instance, field_mask, metadata = api._updated_instance
self.assertEqual(
field_mask.paths, ["config", "display_name", "node_count", "labels"]
)
self.assertEqual(field_mask.paths, self.FIELD_MASK)
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.INSTANCE_ID)
Expand All @@ -417,14 +446,42 @@ def test_update_success(self):

self.assertIs(future, op_future)

instance, field_mask, metadata = api._updated_instance
self.assertEqual(field_mask.paths, self.FIELD_MASK)
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.DISPLAY_NAME)
self.assertEqual(instance.node_count, self.NODE_COUNT)
self.assertEqual(instance.labels, self.LABELS)
self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)])

def test_update_success_with_processing_units(self):
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,
configuration_name=self.CONFIG_NAME,
processing_units=self.PROCESSING_UNITS,
display_name=self.DISPLAY_NAME,
labels=self.LABELS,
)

future = instance.update()

self.assertIs(future, op_future)

instance, field_mask, metadata = api._updated_instance
self.assertEqual(
field_mask.paths, ["config", "display_name", "node_count", "labels"]
field_mask.paths, ["config", "display_name", "processing_units", "labels"]
)
self.assertEqual(instance.name, self.INSTANCE_NAME)
self.assertEqual(instance.config, self.CONFIG_NAME)
self.assertEqual(instance.display_name, self.DISPLAY_NAME)
self.assertEqual(instance.node_count, self.NODE_COUNT)
self.assertEqual(instance.processing_units, self.PROCESSING_UNITS)
self.assertEqual(instance.labels, self.LABELS)
self.assertEqual(metadata, [("google-cloud-resource-prefix", instance.name)])

Expand Down

0 comments on commit 44aa7cc

Please sign in to comment.