Skip to content

Commit 5166b05

Browse files
dhermesdhermes
committed
Adding cluster argument to Bigtable HappyBase connection.
1 parent 8960445 commit 5166b05

File tree

2 files changed

+213
-22
lines changed

2 files changed

+213
-22
lines changed

‎gcloud/bigtable/happybase/connection.py‎

Lines changed: 87 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
import six
1919

20+
from gcloud.bigtable.client import Client
21+
2022

2123
# Constants reproduced here for HappyBase compatibility, though values
2224
# are all null.
@@ -30,6 +32,47 @@
3032
DEFAULT_PROTOCOL = None
3133

3234

35+
def _get_cluster(timeout=None):
36+
"""Gets cluster for the default project.
37+
38+
Creates a client with the inferred credentials and project ID from
39+
the local environment. Then uses :meth:`.Client.list_clusters` to
40+
get the unique cluster owned by the project.
41+
42+
If the request fails for any reason, or if there isn't exactly one cluster
43+
owned by the project, then this function will fail.
44+
45+
:type timeout: int
46+
:param timeout: (Optional) The socket timeout in milliseconds.
47+
48+
:rtype: :class:`gcloud.bigtable.cluster.Cluster`
49+
:returns: The unique cluster owned by the project inferred from
50+
the environment.
51+
:raises: :class:`ValueError <exceptions.ValueError>` if their is a failed
52+
zone or any number of clusters other than one.
53+
"""
54+
client_kwargs = {'admin': True}
55+
if timeout is not None:
56+
client_kwargs['timeout_seconds'] = timeout / 1000.0
57+
client = Client(**client_kwargs)
58+
try:
59+
client.start()
60+
clusters, failed_zones = client.list_clusters()
61+
finally:
62+
client.stop()
63+
64+
if len(failed_zones) != 0:
65+
raise ValueError('Determining cluster via ListClusters encountered '
66+
'failed zones.')
67+
if len(clusters) == 0:
68+
raise ValueError('This client doesn\'t have access to any clusters.')
69+
if len(clusters) > 1:
70+
raise ValueError('This client has access to more than one cluster. '
71+
'Please directly pass the cluster you\'d '
72+
'like to use.')
73+
return clusters[0]
74+
75+
3376
class Connection(object):
3477
"""Connection to Cloud Bigtable backend.
3578
@@ -78,28 +121,27 @@ class Connection(object):
78121
HappyBase, but irrelevant for Cloud Bigtable since the
79122
protocol is fixed.
80123
124+
:type cluster: :class:`gcloud.bigtable.cluster.Cluster`
125+
:param cluster: (Optional) A Cloud Bigtable cluster. The instance also
126+
owns a client for making gRPC requests to the Cloud
127+
Bigtable API. If not passed in, defaults to creating client
128+
with ``admin=True`` and using the ``timeout`` here for the
129+
``timeout_seconds`` argument to the :class:`.Client``
130+
constructor. The credentials for the client
131+
will be the implicit ones loaded from the environment.
132+
Then that client is used to retrieve all the clusters
133+
owned by the client's project.
134+
81135
:raises: :class:`ValueError <exceptions.ValueError>` if any of the unused
82136
parameters are specified with a value other than the defaults.
83137
"""
84138

85139
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
86140
autoconnect=True, table_prefix=None,
87141
table_prefix_separator='_', compat=DEFAULT_COMPAT,
88-
transport=DEFAULT_TRANSPORT, protocol=DEFAULT_PROTOCOL):
89-
if host is not DEFAULT_HOST:
90-
raise ValueError('Host cannot be set for gcloud HappyBase module')
91-
if port is not DEFAULT_PORT:
92-
raise ValueError('Port cannot be set for gcloud HappyBase module')
93-
if compat is not DEFAULT_COMPAT:
94-
raise ValueError('Compat cannot be set for gcloud '
95-
'HappyBase module')
96-
if transport is not DEFAULT_TRANSPORT:
97-
raise ValueError('Transport cannot be set for gcloud '
98-
'HappyBase module')
99-
if protocol is not DEFAULT_PROTOCOL:
100-
raise ValueError('Protocol cannot be set for gcloud '
101-
'HappyBase module')
102-
142+
transport=DEFAULT_TRANSPORT, protocol=DEFAULT_PROTOCOL,
143+
cluster=None):
144+
self._reject_legacy_args(host, port, compat, transport, protocol)
103145
if table_prefix is not None:
104146
if not isinstance(table_prefix, six.string_types):
105147
raise TypeError('table_prefix must be a string', 'received',
@@ -110,7 +152,36 @@ def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
110152
'received', table_prefix_separator,
111153
type(table_prefix_separator))
112154

113-
self.timeout = timeout
114155
self.autoconnect = autoconnect
115156
self.table_prefix = table_prefix
116157
self.table_prefix_separator = table_prefix_separator
158+
159+
if cluster is None:
160+
self._cluster = _get_cluster(timeout=timeout)
161+
else:
162+
if timeout is not None:
163+
raise ValueError('Timeout cannot be used when an existing '
164+
'cluster is passed')
165+
self._cluster = cluster.copy()
166+
167+
@staticmethod
168+
def _reject_legacy_args(host, port, compat, transport, protocol):
169+
"""Check legacy HappyBase arguments and raise if set.
170+
171+
:raises: :class:`ValueError <exceptions.ValueError>` if any of the
172+
legacy parameters are specified with a value other than
173+
the defaults.
174+
"""
175+
if host is not DEFAULT_HOST:
176+
raise ValueError('Host cannot be set for gcloud HappyBase module')
177+
if port is not DEFAULT_PORT:
178+
raise ValueError('Port cannot be set for gcloud HappyBase module')
179+
if compat is not DEFAULT_COMPAT:
180+
raise ValueError('Compat cannot be set for gcloud '
181+
'HappyBase module')
182+
if transport is not DEFAULT_TRANSPORT:
183+
raise ValueError('Transport cannot be set for gcloud '
184+
'HappyBase module')
185+
if protocol is not DEFAULT_PROTOCOL:
186+
raise ValueError('Protocol cannot be set for gcloud '
187+
'HappyBase module')

‎gcloud/bigtable/happybase/test_connection.py‎

Lines changed: 126 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,60 @@
1616
import unittest2
1717

1818

19+
class Test__get_cluster(unittest2.TestCase):
20+
21+
def _callFUT(self, timeout=None):
22+
from gcloud.bigtable.happybase.connection import _get_cluster
23+
return _get_cluster(timeout=timeout)
24+
25+
def _helper(self, timeout=None, clusters=(), failed_zones=()):
26+
from functools import partial
27+
from gcloud._testing import _Monkey
28+
from gcloud.bigtable.happybase import connection as MUT
29+
30+
client_with_clusters = partial(_Client, clusters=clusters,
31+
failed_zones=failed_zones)
32+
with _Monkey(MUT, Client=client_with_clusters):
33+
result = self._callFUT(timeout=timeout)
34+
35+
# If we've reached this point, then _callFUT didn't fail, so we know
36+
# there is exactly one cluster.
37+
cluster, = clusters
38+
self.assertEqual(result, cluster)
39+
client = cluster.client
40+
self.assertEqual(client.args, ())
41+
expected_kwargs = {'admin': True}
42+
if timeout is not None:
43+
expected_kwargs['timeout_seconds'] = timeout / 1000.0
44+
self.assertEqual(client.kwargs, expected_kwargs)
45+
self.assertEqual(client.start_calls, 1)
46+
self.assertEqual(client.stop_calls, 1)
47+
48+
def test_default(self):
49+
cluster = _Cluster()
50+
self._helper(clusters=[cluster])
51+
52+
def test_with_timeout(self):
53+
cluster = _Cluster()
54+
self._helper(timeout=2103, clusters=[cluster])
55+
56+
def test_with_no_clusters(self):
57+
with self.assertRaises(ValueError):
58+
self._helper()
59+
60+
def test_with_too_many_clusters(self):
61+
clusters = [_Cluster(), _Cluster()]
62+
with self.assertRaises(ValueError):
63+
self._helper(clusters=clusters)
64+
65+
def test_with_failed_zones(self):
66+
cluster = _Cluster()
67+
failed_zone = 'us-central1-c'
68+
with self.assertRaises(ValueError):
69+
self._helper(clusters=[cluster],
70+
failed_zones=[failed_zone])
71+
72+
1973
class TestConnection(unittest2.TestCase):
2074

2175
def _getTargetClass(self):
@@ -26,24 +80,47 @@ def _makeOne(self, *args, **kwargs):
2680
return self._getTargetClass()(*args, **kwargs)
2781

2882
def test_constructor_defaults(self):
29-
connection = self._makeOne()
30-
self.assertEqual(connection.timeout, None)
83+
cluster = _Cluster() # Avoid implicit environ check.
84+
connection = self._makeOne(cluster=cluster)
85+
3186
self.assertTrue(connection.autoconnect)
87+
self.assertEqual(connection._cluster, cluster)
3288
self.assertEqual(connection.table_prefix, None)
3389
self.assertEqual(connection.table_prefix_separator, '_')
3490

91+
def test_constructor_missing_cluster(self):
92+
from gcloud._testing import _Monkey
93+
from gcloud.bigtable.happybase import connection as MUT
94+
95+
cluster = _Cluster()
96+
timeout = object()
97+
get_cluster_called = []
98+
99+
def mock_get_cluster(timeout):
100+
get_cluster_called.append(timeout)
101+
return cluster
102+
103+
with _Monkey(MUT, _get_cluster=mock_get_cluster):
104+
connection = self._makeOne(autoconnect=False, cluster=None,
105+
timeout=timeout)
106+
self.assertEqual(connection.table_prefix, None)
107+
self.assertEqual(connection.table_prefix_separator, '_')
108+
self.assertEqual(connection._cluster, cluster)
109+
110+
self.assertEqual(get_cluster_called, [timeout])
111+
35112
def test_constructor_explicit(self):
36-
timeout = 12345
37113
autoconnect = False
38114
table_prefix = 'table-prefix'
39115
table_prefix_separator = 'sep'
116+
cluster_copy = _Cluster()
117+
cluster = _Cluster(copies=[cluster_copy])
40118

41119
connection = self._makeOne(
42-
timeout=timeout,
43120
autoconnect=autoconnect,
44121
table_prefix=table_prefix,
45-
table_prefix_separator=table_prefix_separator)
46-
self.assertEqual(connection.timeout, timeout)
122+
table_prefix_separator=table_prefix_separator,
123+
cluster=cluster)
47124
self.assertFalse(connection.autoconnect)
48125
self.assertEqual(connection.table_prefix, table_prefix)
49126
self.assertEqual(connection.table_prefix_separator,
@@ -69,6 +146,10 @@ def test_constructor_with_protocol(self):
69146
with self.assertRaises(ValueError):
70147
self._makeOne(protocol=object())
71148

149+
def test_constructor_with_timeout_and_cluster(self):
150+
with self.assertRaises(ValueError):
151+
self._makeOne(cluster=object(), timeout=object())
152+
72153
def test_constructor_non_string_prefix(self):
73154
table_prefix = object()
74155

@@ -82,3 +163,42 @@ def test_constructor_non_string_prefix_separator(self):
82163
with self.assertRaises(TypeError):
83164
self._makeOne(autoconnect=False,
84165
table_prefix_separator=table_prefix_separator)
166+
167+
168+
class _Client(object):
169+
170+
def __init__(self, *args, **kwargs):
171+
self.clusters = kwargs.pop('clusters', [])
172+
for cluster in self.clusters:
173+
cluster.client = self
174+
self.failed_zones = kwargs.pop('failed_zones', [])
175+
self.args = args
176+
self.kwargs = kwargs
177+
self.start_calls = 0
178+
self.stop_calls = 0
179+
180+
def start(self):
181+
self.start_calls += 1
182+
183+
def stop(self):
184+
self.stop_calls += 1
185+
186+
def list_clusters(self):
187+
return self.clusters, self.failed_zones
188+
189+
190+
class _Cluster(object):
191+
192+
def __init__(self, copies=(), list_tables_result=()):
193+
self.copies = list(copies)
194+
# Included to support Connection.__del__
195+
self._client = _Client()
196+
self.list_tables_result = list_tables_result
197+
198+
def copy(self):
199+
if self.copies:
200+
result = self.copies[0]
201+
self.copies[:] = self.copies[1:]
202+
return result
203+
else:
204+
return self

0 commit comments

Comments
 (0)