Skip to content

Commit 083e738

Browse files
dhermesdhermes
committed
Adding Bigtable Cluster.from_pb factory.
For creating a Cluster object from a protobuf cluster returned from the API. Also adding a helper for making sure properties are set on a protobuf object.
1 parent b54f725 commit 083e738

File tree

2 files changed

+146
-0
lines changed

2 files changed

+146
-0
lines changed

‎gcloud/bigtable/cluster.py‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,40 @@
1515
"""User friendly container for Google Cloud Bigtable Cluster."""
1616

1717

18+
import re
19+
1820
from gcloud.bigtable.table import Table
1921

2022

23+
_CLUSTER_NAME_RE = re.compile(r'^projects/(?P<project>[^/]+)/'
24+
r'zones/(?P<zone>[^/]+)/clusters/'
25+
r'(?P<cluster_id>[a-z][-a-z0-9]*)$')
26+
27+
28+
def _get_pb_property_value(message_pb, property_name):
29+
"""Return a message field value.
30+
31+
:type message_pb: :class:`google.protobuf.message.Message`
32+
:param message_pb: The message to check for ``property_name``.
33+
34+
:type property_name: str
35+
:param property_name: The property value to check against.
36+
37+
:rtype: object
38+
:returns: The value of ``property_name`` set on ``message_pb``.
39+
:raises: :class:`ValueError <exceptions.ValueError>` if the result returned
40+
from the ``message_pb`` does not contain the ``property_name``
41+
value.
42+
"""
43+
# Make sure `property_name` is set on the response.
44+
# NOTE: As of proto3, HasField() only works for message fields, not for
45+
# singular (non-message) fields.
46+
all_fields = set([field.name for field in message_pb._fields])
47+
if property_name not in all_fields:
48+
raise ValueError('Message does not contain %s.' % (property_name,))
49+
return getattr(message_pb, property_name)
50+
51+
2152
class Cluster(object):
2253
"""Representation of a Google Cloud Bigtable Cluster.
2354
@@ -60,3 +91,35 @@ def table(self, table_id):
6091
:returns: The table owned by this cluster.
6192
"""
6293
return Table(table_id, self)
94+
95+
def _update_from_pb(self, cluster_pb):
96+
self.display_name = _get_pb_property_value(cluster_pb, 'display_name')
97+
self.serve_nodes = _get_pb_property_value(cluster_pb, 'serve_nodes')
98+
99+
@classmethod
100+
def from_pb(cls, cluster_pb, client):
101+
"""Creates a cluster instance from a protobuf.
102+
103+
:type cluster_pb: :class:`bigtable_cluster_data_pb2.Cluster`
104+
:param cluster_pb: A cluster protobuf object.
105+
106+
:type client: :class:`.client.Client`
107+
:param client: The client that owns the cluster.
108+
109+
:rtype: :class:`Cluster`
110+
:returns: The cluster parsed from the protobuf response.
111+
:raises: :class:`ValueError <exceptions.ValueError>` if the cluster
112+
name does not match :data:`_CLUSTER_NAME_RE` or if the parsed
113+
project ID does not match the project ID on the client.
114+
"""
115+
match = _CLUSTER_NAME_RE.match(cluster_pb.name)
116+
if match is None:
117+
raise ValueError('Cluster protobuf name was not in the '
118+
'expected format.', cluster_pb.name)
119+
if match.group('project') != client.project:
120+
raise ValueError('Project ID on cluster does not match the '
121+
'project ID on the client')
122+
123+
result = cls(match.group('zone'), match.group('cluster_id'), client)
124+
result._update_from_pb(cluster_pb)
125+
return result

‎gcloud/bigtable/test_cluster.py‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -65,3 +65,86 @@ def test_table_factory(self):
6565
self.assertTrue(isinstance(table, Table))
6666
self.assertEqual(table.table_id, table_id)
6767
self.assertEqual(table._cluster, cluster)
68+
69+
def test_from_pb_success(self):
70+
from gcloud.bigtable._generated import (
71+
bigtable_cluster_data_pb2 as data_pb2)
72+
73+
project = 'PROJECT'
74+
zone = 'zone'
75+
cluster_id = 'cluster-id'
76+
client = _Client(project=project)
77+
78+
cluster_name = ('projects/' + project + '/zones/' + zone +
79+
'/clusters/' + cluster_id)
80+
cluster_pb = data_pb2.Cluster(
81+
name=cluster_name,
82+
display_name=cluster_id,
83+
serve_nodes=3,
84+
)
85+
86+
klass = self._getTargetClass()
87+
cluster = klass.from_pb(cluster_pb, client)
88+
self.assertTrue(isinstance(cluster, klass))
89+
self.assertEqual(cluster._client, client)
90+
self.assertEqual(cluster.zone, zone)
91+
self.assertEqual(cluster.cluster_id, cluster_id)
92+
93+
def test_from_pb_bad_cluster_name(self):
94+
from gcloud.bigtable._generated import (
95+
bigtable_cluster_data_pb2 as data_pb2)
96+
97+
cluster_name = 'INCORRECT_FORMAT'
98+
cluster_pb = data_pb2.Cluster(name=cluster_name)
99+
100+
klass = self._getTargetClass()
101+
with self.assertRaises(ValueError):
102+
klass.from_pb(cluster_pb, None)
103+
104+
def test_from_pb_project_mistmatch(self):
105+
from gcloud.bigtable._generated import (
106+
bigtable_cluster_data_pb2 as data_pb2)
107+
108+
project = 'PROJECT'
109+
zone = 'zone'
110+
cluster_id = 'cluster-id'
111+
alt_project = 'ALT_PROJECT'
112+
client = _Client(project=alt_project)
113+
114+
self.assertNotEqual(project, alt_project)
115+
116+
cluster_name = ('projects/' + project + '/zones/' + zone +
117+
'/clusters/' + cluster_id)
118+
cluster_pb = data_pb2.Cluster(name=cluster_name)
119+
120+
klass = self._getTargetClass()
121+
with self.assertRaises(ValueError):
122+
klass.from_pb(cluster_pb, client)
123+
124+
125+
class Test__get_pb_property_value(unittest2.TestCase):
126+
127+
def _callFUT(self, message_pb, property_name):
128+
from gcloud.bigtable.cluster import _get_pb_property_value
129+
return _get_pb_property_value(message_pb, property_name)
130+
131+
def test_it(self):
132+
from gcloud.bigtable._generated import (
133+
bigtable_cluster_data_pb2 as data_pb2)
134+
serve_nodes = 119
135+
cluster_pb = data_pb2.Cluster(serve_nodes=serve_nodes)
136+
result = self._callFUT(cluster_pb, 'serve_nodes')
137+
self.assertEqual(result, serve_nodes)
138+
139+
def test_with_value_unset_on_pb(self):
140+
from gcloud.bigtable._generated import (
141+
bigtable_cluster_data_pb2 as data_pb2)
142+
cluster_pb = data_pb2.Cluster()
143+
with self.assertRaises(ValueError):
144+
self._callFUT(cluster_pb, 'serve_nodes')
145+
146+
147+
class _Client(object):
148+
149+
def __init__(self, project):
150+
self.project = project

0 commit comments

Comments
 (0)