diff --git a/gcloud/bigtable/happybase/connection.py b/gcloud/bigtable/happybase/connection.py index 7ef6b483bfce..2b38949ef10e 100644 --- a/gcloud/bigtable/happybase/connection.py +++ b/gcloud/bigtable/happybase/connection.py @@ -15,8 +15,12 @@ """Google Cloud Bigtable HappyBase connection module.""" +import warnings + import six +from gcloud.bigtable.client import Client + # Constants reproduced here for HappyBase compatibility, though values # are all null. @@ -29,6 +33,50 @@ DEFAULT_COMPAT = None DEFAULT_PROTOCOL = None +_LEGACY_ARGS = frozenset(('host', 'port', 'compat', 'transport', 'protocol')) +_WARN = warnings.warn + + +def _get_cluster(timeout=None): + """Gets cluster for the default project. + + Creates a client with the inferred credentials and project ID from + the local environment. Then uses :meth:`.Client.list_clusters` to + get the unique cluster owned by the project. + + If the request fails for any reason, or if there isn't exactly one cluster + owned by the project, then this function will fail. + + :type timeout: int + :param timeout: (Optional) The socket timeout in milliseconds. + + :rtype: :class:`gcloud.bigtable.cluster.Cluster` + :returns: The unique cluster owned by the project inferred from + the environment. + :raises: :class:`ValueError ` if their is a failed + zone or any number of clusters other than one. + """ + client_kwargs = {'admin': True} + if timeout is not None: + client_kwargs['timeout_seconds'] = timeout / 1000.0 + client = Client(**client_kwargs) + try: + client.start() + clusters, failed_zones = client.list_clusters() + finally: + client.stop() + + if len(failed_zones) != 0: + raise ValueError('Determining cluster via ListClusters encountered ' + 'failed zones.') + if len(clusters) == 0: + raise ValueError('This client doesn\'t have access to any clusters.') + if len(clusters) > 1: + raise ValueError('This client has access to more than one cluster. ' + 'Please directly pass the cluster you\'d ' + 'like to use.') + return clusters[0] + class Connection(object): """Connection to Cloud Bigtable backend. @@ -41,13 +89,10 @@ class Connection(object): :class:`Credentials ` stored on the client. - :type host: :data:`NoneType ` - :param host: Unused parameter. Provided for compatibility with HappyBase, - but irrelevant for Cloud Bigtable since it has a fixed host. - - :type port: :data:`NoneType ` - :param port: Unused parameter. Provided for compatibility with HappyBase, - but irrelevant for Cloud Bigtable since it has a fixed host. + The arguments ``host``, ``port``, ``compat``, ``transport`` and + ``protocol`` are allowed (as keyword arguments) for compatibility with + HappyBase. However, they will not be used in anyway, and will cause a + warning if passed. :type timeout: int :param timeout: (Optional) The socket timeout in milliseconds. @@ -63,43 +108,28 @@ class Connection(object): :param table_prefix_separator: (Optional) Separator used with ``table_prefix``. Defaults to ``_``. - :type compat: :data:`NoneType ` - :param compat: Unused parameter. Provided for compatibility with - HappyBase, but irrelevant for Cloud Bigtable since there - is only one version. - - :type transport: :data:`NoneType ` - :param transport: Unused parameter. Provided for compatibility with - HappyBase, but irrelevant for Cloud Bigtable since the - transport is fixed. - - :type protocol: :data:`NoneType ` - :param protocol: Unused parameter. Provided for compatibility with - HappyBase, but irrelevant for Cloud Bigtable since the - protocol is fixed. + :type cluster: :class:`gcloud.bigtable.cluster.Cluster` + :param cluster: (Optional) A Cloud Bigtable cluster. The instance also + owns a client for making gRPC requests to the Cloud + Bigtable API. If not passed in, defaults to creating client + with ``admin=True`` and using the ``timeout`` here for the + ``timeout_seconds`` argument to the :class:`.Client`` + constructor. The credentials for the client + will be the implicit ones loaded from the environment. + Then that client is used to retrieve all the clusters + owned by the client's project. + + :type kwargs: dict + :param kwargs: Remaining keyword arguments. Provided for HappyBase + compatibility. :raises: :class:`ValueError ` if any of the unused parameters are specified with a value other than the defaults. """ - def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None, - autoconnect=True, table_prefix=None, - table_prefix_separator='_', compat=DEFAULT_COMPAT, - transport=DEFAULT_TRANSPORT, protocol=DEFAULT_PROTOCOL): - if host is not DEFAULT_HOST: - raise ValueError('Host cannot be set for gcloud HappyBase module') - if port is not DEFAULT_PORT: - raise ValueError('Port cannot be set for gcloud HappyBase module') - if compat is not DEFAULT_COMPAT: - raise ValueError('Compat cannot be set for gcloud ' - 'HappyBase module') - if transport is not DEFAULT_TRANSPORT: - raise ValueError('Transport cannot be set for gcloud ' - 'HappyBase module') - if protocol is not DEFAULT_PROTOCOL: - raise ValueError('Protocol cannot be set for gcloud ' - 'HappyBase module') - + def __init__(self, timeout=None, autoconnect=True, table_prefix=None, + table_prefix_separator='_', cluster=None, **kwargs): + self._handle_legacy_args(kwargs) if table_prefix is not None: if not isinstance(table_prefix, six.string_types): raise TypeError('table_prefix must be a string', 'received', @@ -110,7 +140,37 @@ def __init__(self, host=DEFAULT_HOST, port=DEFAULT_PORT, timeout=None, 'received', table_prefix_separator, type(table_prefix_separator)) - self.timeout = timeout self.autoconnect = autoconnect self.table_prefix = table_prefix self.table_prefix_separator = table_prefix_separator + + if cluster is None: + self._cluster = _get_cluster(timeout=timeout) + else: + if timeout is not None: + raise ValueError('Timeout cannot be used when an existing ' + 'cluster is passed') + self._cluster = cluster.copy() + + @staticmethod + def _handle_legacy_args(arguments_dict): + """Check legacy HappyBase arguments and warn if set. + + :type arguments_dict: dict + :param arguments_dict: Unused keyword arguments. + + :raises: :class:`TypeError ` if a keyword other + than ``host``, ``port``, ``compat``, ``transport`` or + ``protocol`` is used. + """ + common_args = _LEGACY_ARGS.intersection(six.iterkeys(arguments_dict)) + if common_args: + all_args = ', '.join(common_args) + message = ('The HappyBase legacy arguments %s were used. These ' + 'arguments are unused by gcloud.' % (all_args,)) + _WARN(message) + for arg_name in common_args: + arguments_dict.pop(arg_name) + if arguments_dict: + unexpected_names = arguments_dict.keys() + raise TypeError('Received unexpected arguments', unexpected_names) diff --git a/gcloud/bigtable/happybase/test_connection.py b/gcloud/bigtable/happybase/test_connection.py index 5913ab2219c9..33003992045d 100644 --- a/gcloud/bigtable/happybase/test_connection.py +++ b/gcloud/bigtable/happybase/test_connection.py @@ -16,6 +16,60 @@ import unittest2 +class Test__get_cluster(unittest2.TestCase): + + def _callFUT(self, timeout=None): + from gcloud.bigtable.happybase.connection import _get_cluster + return _get_cluster(timeout=timeout) + + def _helper(self, timeout=None, clusters=(), failed_zones=()): + from functools import partial + from gcloud._testing import _Monkey + from gcloud.bigtable.happybase import connection as MUT + + client_with_clusters = partial(_Client, clusters=clusters, + failed_zones=failed_zones) + with _Monkey(MUT, Client=client_with_clusters): + result = self._callFUT(timeout=timeout) + + # If we've reached this point, then _callFUT didn't fail, so we know + # there is exactly one cluster. + cluster, = clusters + self.assertEqual(result, cluster) + client = cluster.client + self.assertEqual(client.args, ()) + expected_kwargs = {'admin': True} + if timeout is not None: + expected_kwargs['timeout_seconds'] = timeout / 1000.0 + self.assertEqual(client.kwargs, expected_kwargs) + self.assertEqual(client.start_calls, 1) + self.assertEqual(client.stop_calls, 1) + + def test_default(self): + cluster = _Cluster() + self._helper(clusters=[cluster]) + + def test_with_timeout(self): + cluster = _Cluster() + self._helper(timeout=2103, clusters=[cluster]) + + def test_with_no_clusters(self): + with self.assertRaises(ValueError): + self._helper() + + def test_with_too_many_clusters(self): + clusters = [_Cluster(), _Cluster()] + with self.assertRaises(ValueError): + self._helper(clusters=clusters) + + def test_with_failed_zones(self): + cluster = _Cluster() + failed_zone = 'us-central1-c' + with self.assertRaises(ValueError): + self._helper(clusters=[cluster], + failed_zones=[failed_zone]) + + class TestConnection(unittest2.TestCase): def _getTargetClass(self): @@ -26,48 +80,83 @@ def _makeOne(self, *args, **kwargs): return self._getTargetClass()(*args, **kwargs) def test_constructor_defaults(self): - connection = self._makeOne() - self.assertEqual(connection.timeout, None) + cluster = _Cluster() # Avoid implicit environ check. + connection = self._makeOne(cluster=cluster) + self.assertTrue(connection.autoconnect) + self.assertEqual(connection._cluster, cluster) self.assertEqual(connection.table_prefix, None) self.assertEqual(connection.table_prefix_separator, '_') + def test_constructor_missing_cluster(self): + from gcloud._testing import _Monkey + from gcloud.bigtable.happybase import connection as MUT + + cluster = _Cluster() + timeout = object() + get_cluster_called = [] + + def mock_get_cluster(timeout): + get_cluster_called.append(timeout) + return cluster + + with _Monkey(MUT, _get_cluster=mock_get_cluster): + connection = self._makeOne(autoconnect=False, cluster=None, + timeout=timeout) + self.assertEqual(connection.table_prefix, None) + self.assertEqual(connection.table_prefix_separator, '_') + self.assertEqual(connection._cluster, cluster) + + self.assertEqual(get_cluster_called, [timeout]) + def test_constructor_explicit(self): - timeout = 12345 autoconnect = False table_prefix = 'table-prefix' table_prefix_separator = 'sep' + cluster_copy = _Cluster() + cluster = _Cluster(copies=[cluster_copy]) connection = self._makeOne( - timeout=timeout, autoconnect=autoconnect, table_prefix=table_prefix, - table_prefix_separator=table_prefix_separator) - self.assertEqual(connection.timeout, timeout) + table_prefix_separator=table_prefix_separator, + cluster=cluster) self.assertFalse(connection.autoconnect) self.assertEqual(connection.table_prefix, table_prefix) self.assertEqual(connection.table_prefix_separator, table_prefix_separator) - def test_constructor_with_host(self): - with self.assertRaises(ValueError): - self._makeOne(host=object()) + def test_constructor_with_unknown_argument(self): + cluster = _Cluster() + with self.assertRaises(TypeError): + self._makeOne(cluster=cluster, unknown='foo') - def test_constructor_with_port(self): - with self.assertRaises(ValueError): - self._makeOne(port=object()) + def test_constructor_with_legacy_args(self): + from gcloud._testing import _Monkey + from gcloud.bigtable.happybase import connection as MUT - def test_constructor_with_compat(self): - with self.assertRaises(ValueError): - self._makeOne(compat=object()) + warned = [] - def test_constructor_with_transport(self): - with self.assertRaises(ValueError): - self._makeOne(transport=object()) + def mock_warn(msg): + warned.append(msg) + + cluster = _Cluster() + with _Monkey(MUT, _WARN=mock_warn): + self._makeOne(cluster=cluster, host=object(), + port=object(), compat=object(), + transport=object(), protocol=object()) - def test_constructor_with_protocol(self): + self.assertEqual(len(warned), 1) + self.assertIn('host', warned[0]) + self.assertIn('port', warned[0]) + self.assertIn('compat', warned[0]) + self.assertIn('transport', warned[0]) + self.assertIn('protocol', warned[0]) + + def test_constructor_with_timeout_and_cluster(self): + cluster = _Cluster() with self.assertRaises(ValueError): - self._makeOne(protocol=object()) + self._makeOne(cluster=cluster, timeout=object()) def test_constructor_non_string_prefix(self): table_prefix = object() @@ -82,3 +171,42 @@ def test_constructor_non_string_prefix_separator(self): with self.assertRaises(TypeError): self._makeOne(autoconnect=False, table_prefix_separator=table_prefix_separator) + + +class _Client(object): + + def __init__(self, *args, **kwargs): + self.clusters = kwargs.pop('clusters', []) + for cluster in self.clusters: + cluster.client = self + self.failed_zones = kwargs.pop('failed_zones', []) + self.args = args + self.kwargs = kwargs + self.start_calls = 0 + self.stop_calls = 0 + + def start(self): + self.start_calls += 1 + + def stop(self): + self.stop_calls += 1 + + def list_clusters(self): + return self.clusters, self.failed_zones + + +class _Cluster(object): + + def __init__(self, copies=(), list_tables_result=()): + self.copies = list(copies) + # Included to support Connection.__del__ + self._client = _Client() + self.list_tables_result = list_tables_result + + def copy(self): + if self.copies: + result = self.copies[0] + self.copies[:] = self.copies[1:] + return result + else: + return self