diff --git a/gcloud/bigtable/cluster.py b/gcloud/bigtable/cluster.py index e2d8420b28e7..677ca72c8eb3 100644 --- a/gcloud/bigtable/cluster.py +++ b/gcloud/bigtable/cluster.py @@ -15,9 +15,40 @@ """User friendly container for Google Cloud Bigtable Cluster.""" +import re + from gcloud.bigtable.table import Table +_CLUSTER_NAME_RE = re.compile(r'^projects/(?P[^/]+)/' + r'zones/(?P[^/]+)/clusters/' + r'(?P[a-z][-a-z0-9]*)$') + + +def _get_pb_property_value(message_pb, property_name): + """Return a message field value. + + :type message_pb: :class:`google.protobuf.message.Message` + :param message_pb: The message to check for ``property_name``. + + :type property_name: str + :param property_name: The property value to check against. + + :rtype: object + :returns: The value of ``property_name`` set on ``message_pb``. + :raises: :class:`ValueError ` if the result returned + from the ``message_pb`` does not contain the ``property_name`` + value. + """ + # Make sure `property_name` is set on the response. + # NOTE: As of proto3, HasField() only works for message fields, not for + # singular (non-message) fields. + all_fields = set([field.name for field in message_pb._fields]) + if property_name not in all_fields: + raise ValueError('Message does not contain %s.' % (property_name,)) + return getattr(message_pb, property_name) + + class Cluster(object): """Representation of a Google Cloud Bigtable Cluster. @@ -60,3 +91,35 @@ def table(self, table_id): :returns: The table owned by this cluster. """ return Table(table_id, self) + + def _update_from_pb(self, cluster_pb): + self.display_name = _get_pb_property_value(cluster_pb, 'display_name') + self.serve_nodes = _get_pb_property_value(cluster_pb, 'serve_nodes') + + @classmethod + def from_pb(cls, cluster_pb, client): + """Creates a cluster instance from a protobuf. + + :type cluster_pb: :class:`bigtable_cluster_data_pb2.Cluster` + :param cluster_pb: A cluster protobuf object. + + :type client: :class:`.client.Client` + :param client: The client that owns the cluster. + + :rtype: :class:`Cluster` + :returns: The cluster parsed from the protobuf response. + :raises: :class:`ValueError ` if the cluster + name does not match :data:`_CLUSTER_NAME_RE` or if the parsed + project ID does not match the project ID on the client. + """ + match = _CLUSTER_NAME_RE.match(cluster_pb.name) + if match is None: + raise ValueError('Cluster protobuf name was not in the ' + 'expected format.', cluster_pb.name) + if match.group('project') != client.project: + raise ValueError('Project ID on cluster does not match the ' + 'project ID on the client') + + result = cls(match.group('zone'), match.group('cluster_id'), client) + result._update_from_pb(cluster_pb) + return result diff --git a/gcloud/bigtable/test_cluster.py b/gcloud/bigtable/test_cluster.py index 8e3855070e3e..6f19dadad1ff 100644 --- a/gcloud/bigtable/test_cluster.py +++ b/gcloud/bigtable/test_cluster.py @@ -65,3 +65,86 @@ def test_table_factory(self): self.assertTrue(isinstance(table, Table)) self.assertEqual(table.table_id, table_id) self.assertEqual(table._cluster, cluster) + + def test_from_pb_success(self): + from gcloud.bigtable._generated import ( + bigtable_cluster_data_pb2 as data_pb2) + + project = 'PROJECT' + zone = 'zone' + cluster_id = 'cluster-id' + client = _Client(project=project) + + cluster_name = ('projects/' + project + '/zones/' + zone + + '/clusters/' + cluster_id) + cluster_pb = data_pb2.Cluster( + name=cluster_name, + display_name=cluster_id, + serve_nodes=3, + ) + + klass = self._getTargetClass() + cluster = klass.from_pb(cluster_pb, client) + self.assertTrue(isinstance(cluster, klass)) + self.assertEqual(cluster._client, client) + self.assertEqual(cluster.zone, zone) + self.assertEqual(cluster.cluster_id, cluster_id) + + def test_from_pb_bad_cluster_name(self): + from gcloud.bigtable._generated import ( + bigtable_cluster_data_pb2 as data_pb2) + + cluster_name = 'INCORRECT_FORMAT' + cluster_pb = data_pb2.Cluster(name=cluster_name) + + klass = self._getTargetClass() + with self.assertRaises(ValueError): + klass.from_pb(cluster_pb, None) + + def test_from_pb_project_mistmatch(self): + from gcloud.bigtable._generated import ( + bigtable_cluster_data_pb2 as data_pb2) + + project = 'PROJECT' + zone = 'zone' + cluster_id = 'cluster-id' + alt_project = 'ALT_PROJECT' + client = _Client(project=alt_project) + + self.assertNotEqual(project, alt_project) + + cluster_name = ('projects/' + project + '/zones/' + zone + + '/clusters/' + cluster_id) + cluster_pb = data_pb2.Cluster(name=cluster_name) + + klass = self._getTargetClass() + with self.assertRaises(ValueError): + klass.from_pb(cluster_pb, client) + + +class Test__get_pb_property_value(unittest2.TestCase): + + def _callFUT(self, message_pb, property_name): + from gcloud.bigtable.cluster import _get_pb_property_value + return _get_pb_property_value(message_pb, property_name) + + def test_it(self): + from gcloud.bigtable._generated import ( + bigtable_cluster_data_pb2 as data_pb2) + serve_nodes = 119 + cluster_pb = data_pb2.Cluster(serve_nodes=serve_nodes) + result = self._callFUT(cluster_pb, 'serve_nodes') + self.assertEqual(result, serve_nodes) + + def test_with_value_unset_on_pb(self): + from gcloud.bigtable._generated import ( + bigtable_cluster_data_pb2 as data_pb2) + cluster_pb = data_pb2.Cluster() + with self.assertRaises(ValueError): + self._callFUT(cluster_pb, 'serve_nodes') + + +class _Client(object): + + def __init__(self, project): + self.project = project