Skip to content

Commit b30d3f4

Browse files
dhermesdhermes
committed
Merge pull request #1405 from dhermes/bigtable-add-cluster-to-happybase
Adding cluster argument to Bigtable HappyBase connection.
2 parents ae36d52 + 9c40113 commit b30d3f4

2 files changed

Lines changed: 248 additions & 60 deletions

File tree

‎gcloud/bigtable/happybase/connection.py‎

Lines changed: 100 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,12 @@
1515
"""Google Cloud Bigtable HappyBase connection module."""
1616

1717

18+
import warnings
19+
1820
import six
1921

22+
from gcloud.bigtable.client import Client
23+
2024

2125
# Constants reproduced here for HappyBase compatibility, though values
2226
# are all null.
@@ -29,6 +33,50 @@
2933
DEFAULT_COMPAT = None
3034
DEFAULT_PROTOCOL = None
3135

36+
_LEGACY_ARGS = frozenset(('host', 'port', 'compat', 'transport', 'protocol'))
37+
_WARN = warnings.warn
38+
39+
40+
def _get_cluster(timeout=None):
41+
"""Gets cluster for the default project.
42+
43+
Creates a client with the inferred credentials and project ID from
44+
the local environment. Then uses :meth:`.Client.list_clusters` to
45+
get the unique cluster owned by the project.
46+
47+
If the request fails for any reason, or if there isn't exactly one cluster
48+
owned by the project, then this function will fail.
49+
50+
:type timeout: int
51+
:param timeout: (Optional) The socket timeout in milliseconds.
52+
53+
:rtype: :class:`gcloud.bigtable.cluster.Cluster`
54+
:returns: The unique cluster owned by the project inferred from
55+
the environment.
56+
:raises: :class:`ValueError <exceptions.ValueError>` if their is a failed
57+
zone or any number of clusters other than one.
58+
"""
59+
client_kwargs = {'admin': True}
60+
if timeout is not None:
61+
client_kwargs['timeout_seconds'] = timeout / 1000.0
62+
client = Client(**client_kwargs)
63+
try:
64+
client.start()
65+
clusters, failed_zones = client.list_clusters()
66+
finally:
67+
client.stop()
68+
69+
if len(failed_zones) != 0:
70+
raise ValueError('Determining cluster via ListClusters encountered '
71+
'failed zones.')
72+
if len(clusters) == 0:
73+
raise ValueError('This client doesn\'t have access to any clusters.')
74+
if len(clusters) > 1:
75+
raise ValueError('This client has access to more than one cluster. '
76+
'Please directly pass the cluster you\'d '
77+
'like to use.')
78+
return clusters[0]
79+
3280

3381
class Connection(object):
3482
"""Connection to Cloud Bigtable backend.
@@ -41,13 +89,10 @@ class Connection(object):
4189
:class:`Credentials <oauth2client.client.Credentials>` stored on the
4290
client.
4391
44-
:type host: :data:`NoneType <types.NoneType>`
45-
:param host: Unused parameter. Provided for compatibility with HappyBase,
46-
but irrelevant for Cloud Bigtable since it has a fixed host.
47-
48-
:type port: :data:`NoneType <types.NoneType>`
49-
:param port: Unused parameter. Provided for compatibility with HappyBase,
50-
but irrelevant for Cloud Bigtable since it has a fixed host.
92+
The arguments ``host``, ``port``, ``compat``, ``transport`` and
93+
``protocol`` are allowed (as keyword arguments) for compatibility with
94+
HappyBase. However, they will not be used in anyway, and will cause a
95+
warning if passed.
5196
5297
:type timeout: int
5398
:param timeout: (Optional) The socket timeout in milliseconds.
@@ -63,43 +108,28 @@ class Connection(object):
63108
:param table_prefix_separator: (Optional) Separator used with
64109
``table_prefix``. Defaults to ``_``.
65110
66-
:type compat: :data:`NoneType <types.NoneType>`
67-
:param compat: Unused parameter. Provided for compatibility with
68-
HappyBase, but irrelevant for Cloud Bigtable since there
69-
is only one version.
70-
71-
:type transport: :data:`NoneType <types.NoneType>`
72-
:param transport: Unused parameter. Provided for compatibility with
73-
HappyBase, but irrelevant for Cloud Bigtable since the
74-
transport is fixed.
75-
76-
:type protocol: :data:`NoneType <types.NoneType>`
77-
:param protocol: Unused parameter. Provided for compatibility with
78-
HappyBase, but irrelevant for Cloud Bigtable since the
79-
protocol is fixed.
111+
:type cluster: :class:`gcloud.bigtable.cluster.Cluster`
112+
:param cluster: (Optional) A Cloud Bigtable cluster. The instance also
113+
owns a client for making gRPC requests to the Cloud
114+
Bigtable API. If not passed in, defaults to creating client
115+
with ``admin=True`` and using the ``timeout`` here for the
116+
``timeout_seconds`` argument to the :class:`.Client``
117+
constructor. The credentials for the client
118+
will be the implicit ones loaded from the environment.
119+
Then that client is used to retrieve all the clusters
120+
owned by the client's project.
121+
122+
:type kwargs: dict
123+
:param kwargs: Remaining keyword arguments. Provided for HappyBase
124+
compatibility.
80125
81126
:raises: :class:`ValueError <exceptions.ValueError>` if any of the unused
82127
parameters are specified with a value other than the defaults.
83128
"""
84129

85-
def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
86-
autoconnect=True, table_prefix=None,
87-
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-
130+
def __init__(self, timeout=None, autoconnect=True, table_prefix=None,
131+
table_prefix_separator='_', cluster=None, **kwargs):
132+
self._handle_legacy_args(kwargs)
103133
if table_prefix is not None:
104134
if not isinstance(table_prefix, six.string_types):
105135
raise TypeError('table_prefix must be a string', 'received',
@@ -110,7 +140,37 @@ def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None,
110140
'received', table_prefix_separator,
111141
type(table_prefix_separator))
112142

113-
self.timeout = timeout
114143
self.autoconnect = autoconnect
115144
self.table_prefix = table_prefix
116145
self.table_prefix_separator = table_prefix_separator
146+
147+
if cluster is None:
148+
self._cluster = _get_cluster(timeout=timeout)
149+
else:
150+
if timeout is not None:
151+
raise ValueError('Timeout cannot be used when an existing '
152+
'cluster is passed')
153+
self._cluster = cluster.copy()
154+
155+
@staticmethod
156+
def _handle_legacy_args(arguments_dict):
157+
"""Check legacy HappyBase arguments and warn if set.
158+
159+
:type arguments_dict: dict
160+
:param arguments_dict: Unused keyword arguments.
161+
162+
:raises: :class:`TypeError <exceptions.TypeError>` if a keyword other
163+
than ``host``, ``port``, ``compat``, ``transport`` or
164+
``protocol`` is used.
165+
"""
166+
common_args = _LEGACY_ARGS.intersection(six.iterkeys(arguments_dict))
167+
if common_args:
168+
all_args = ', '.join(common_args)
169+
message = ('The HappyBase legacy arguments %s were used. These '
170+
'arguments are unused by gcloud.' % (all_args,))
171+
_WARN(message)
172+
for arg_name in common_args:
173+
arguments_dict.pop(arg_name)
174+
if arguments_dict:
175+
unexpected_names = arguments_dict.keys()
176+
raise TypeError('Received unexpected arguments', unexpected_names)

‎gcloud/bigtable/happybase/test_connection.py‎

Lines changed: 148 additions & 20 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,48 +80,83 @@ 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,
50127
table_prefix_separator)
51128

52-
def test_constructor_with_host(self):
53-
with self.assertRaises(ValueError):
54-
self._makeOne(host=object())
129+
def test_constructor_with_unknown_argument(self):
130+
cluster = _Cluster()
131+
with self.assertRaises(TypeError):
132+
self._makeOne(cluster=cluster, unknown='foo')
55133

56-
def test_constructor_with_port(self):
57-
with self.assertRaises(ValueError):
58-
self._makeOne(port=object())
134+
def test_constructor_with_legacy_args(self):
135+
from gcloud._testing import _Monkey
136+
from gcloud.bigtable.happybase import connection as MUT
59137

60-
def test_constructor_with_compat(self):
61-
with self.assertRaises(ValueError):
62-
self._makeOne(compat=object())
138+
warned = []
63139

64-
def test_constructor_with_transport(self):
65-
with self.assertRaises(ValueError):
66-
self._makeOne(transport=object())
140+
def mock_warn(msg):
141+
warned.append(msg)
142+
143+
cluster = _Cluster()
144+
with _Monkey(MUT, _WARN=mock_warn):
145+
self._makeOne(cluster=cluster, host=object(),
146+
port=object(), compat=object(),
147+
transport=object(), protocol=object())
67148

68-
def test_constructor_with_protocol(self):
149+
self.assertEqual(len(warned), 1)
150+
self.assertIn('host', warned[0])
151+
self.assertIn('port', warned[0])
152+
self.assertIn('compat', warned[0])
153+
self.assertIn('transport', warned[0])
154+
self.assertIn('protocol', warned[0])
155+
156+
def test_constructor_with_timeout_and_cluster(self):
157+
cluster = _Cluster()
69158
with self.assertRaises(ValueError):
70-
self._makeOne(protocol=object())
159+
self._makeOne(cluster=cluster, timeout=object())
71160

72161
def test_constructor_non_string_prefix(self):
73162
table_prefix = object()
@@ -82,3 +171,42 @@ def test_constructor_non_string_prefix_separator(self):
82171
with self.assertRaises(TypeError):
83172
self._makeOne(autoconnect=False,
84173
table_prefix_separator=table_prefix_separator)
174+
175+
176+
class _Client(object):
177+
178+
def __init__(self, *args, **kwargs):
179+
self.clusters = kwargs.pop('clusters', [])
180+
for cluster in self.clusters:
181+
cluster.client = self
182+
self.failed_zones = kwargs.pop('failed_zones', [])
183+
self.args = args
184+
self.kwargs = kwargs
185+
self.start_calls = 0
186+
self.stop_calls = 0
187+
188+
def start(self):
189+
self.start_calls += 1
190+
191+
def stop(self):
192+
self.stop_calls += 1
193+
194+
def list_clusters(self):
195+
return self.clusters, self.failed_zones
196+
197+
198+
class _Cluster(object):
199+
200+
def __init__(self, copies=(), list_tables_result=()):
201+
self.copies = list(copies)
202+
# Included to support Connection.__del__
203+
self._client = _Client()
204+
self.list_tables_result = list_tables_result
205+
206+
def copy(self):
207+
if self.copies:
208+
result = self.copies[0]
209+
self.copies[:] = self.copies[1:]
210+
return result
211+
else:
212+
return self

0 commit comments

Comments
 (0)