Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
63 changes: 63 additions & 0 deletions gcloud/bigtable/cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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
83 changes: 83 additions & 0 deletions gcloud/bigtable/test_cluster.py
Original file line number Diff line number Diff line change
Expand Up @@ -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