googleapis · dhermes · Feb 12, 2016 · Jan 21, 2016 · Feb 12, 2016
diff --git 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 <exceptions.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 <oauth2client.client.Credentials>` stored on the
        client.

-    :type host: :data:`NoneType <types.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 <types.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 <types.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 <types.NoneType>`
-    :param transport: Unused parameter. Provided for compatibility with
-                      HappyBase, but irrelevant for Cloud Bigtable since the
-                      transport is fixed.
-
-    :type protocol: :data:`NoneType <types.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 <exceptions.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 <exceptions.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
@@ -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