Merge pull request #1235 from dhermes/bigtable-cluster-from-pb

Adding Bigtable Cluster.from_pb factory.

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<project>[^/]+)/'
+                              r'zones/(?P<zone>[^/]+)/clusters/'
+                              r'(?P<cluster_id>[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 <exceptions.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 <exceptions.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

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