From e33fe3d5c1093785b9f26c97e4f13c0043e8e40d Mon Sep 17 00:00:00 2001 From: Danny Hermes Date: Thu, 24 Sep 2015 10:53:26 -0700 Subject: [PATCH] Adding ability to start/stop clients. This creates gRPC stubs for talking to each Bigtable service needed by the client. --- gcloud/bigtable/client.py | 98 +++++++++++++ gcloud/bigtable/test_client.py | 247 +++++++++++++++++++++++++++++++++ setup.py | 2 +- 3 files changed, 346 insertions(+), 1 deletion(-) diff --git a/gcloud/bigtable/client.py b/gcloud/bigtable/client.py index 6ebe07698b068..d10cda2aded8e 100644 --- a/gcloud/bigtable/client.py +++ b/gcloud/bigtable/client.py @@ -27,26 +27,39 @@ """ +from gcloud.bigtable._generated import bigtable_cluster_service_pb2 +from gcloud.bigtable._generated import bigtable_service_pb2 +from gcloud.bigtable._generated import bigtable_table_service_pb2 +from gcloud.bigtable._generated import operations_pb2 +from gcloud.bigtable._helpers import make_stub from gcloud.client import _ClientFactoryMixin from gcloud.client import _ClientProjectMixin from gcloud.credentials import get_credentials +TABLE_STUB_FACTORY = (bigtable_table_service_pb2. + early_adopter_create_BigtableTableService_stub) TABLE_ADMIN_HOST = 'bigtabletableadmin.googleapis.com' """Table Admin API request host.""" TABLE_ADMIN_PORT = 443 """Table Admin API request port.""" +CLUSTER_STUB_FACTORY = (bigtable_cluster_service_pb2. + early_adopter_create_BigtableClusterService_stub) CLUSTER_ADMIN_HOST = 'bigtableclusteradmin.googleapis.com' """Cluster Admin API request host.""" CLUSTER_ADMIN_PORT = 443 """Cluster Admin API request port.""" +DATA_STUB_FACTORY = (bigtable_service_pb2. + early_adopter_create_BigtableService_stub) DATA_API_HOST = 'bigtable.googleapis.com' """Data API request host.""" DATA_API_PORT = 443 """Data API request port.""" +OPERATIONS_STUB_FACTORY = operations_pb2.early_adopter_create_Operations_stub + ADMIN_SCOPE = 'https://www.googleapis.com/auth/cloud-bigtable.admin' """Scope for interacting with the Cluster Admin and Table Admin APIs.""" DATA_SCOPE = 'https://www.googleapis.com/auth/cloud-bigtable.data' @@ -225,3 +238,88 @@ def table_stub(self): if self._table_stub is None: raise ValueError('Client has not been started.') return self._table_stub + + def _make_data_stub(self): + """Creates gRPC stub to make requests to the Data API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + """ + return make_stub(self, DATA_STUB_FACTORY, + DATA_API_HOST, DATA_API_PORT) + + def _make_cluster_stub(self): + """Creates gRPC stub to make requests to the Cluster Admin API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + """ + return make_stub(self, CLUSTER_STUB_FACTORY, + CLUSTER_ADMIN_HOST, CLUSTER_ADMIN_PORT) + + def _make_operations_stub(self): + """Creates gRPC stub to make requests to the Operations API. + + These are for long-running operations of the Cluster Admin API, + hence the host and port matching. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + """ + return make_stub(self, OPERATIONS_STUB_FACTORY, + CLUSTER_ADMIN_HOST, CLUSTER_ADMIN_PORT) + + def _make_table_stub(self): + """Creates gRPC stub to make requests to the Table Admin API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + """ + return make_stub(self, TABLE_STUB_FACTORY, + TABLE_ADMIN_HOST, TABLE_ADMIN_PORT) + + def is_started(self): + """Check if the client has been started. + + :rtype: bool + :returns: Boolean indicating if the client has been started. + """ + return self._data_stub is not None + + def start(self): + """Prepare the client to make requests. + + Activates gRPC contexts for making requests to the Bigtable + Service(s). + """ + if self.is_started(): + return + + self._data_stub = self._make_data_stub() + self._data_stub.__enter__() + if self._admin: + self._cluster_stub = self._make_cluster_stub() + self._operations_stub = self._make_operations_stub() + self._table_stub = self._make_table_stub() + + self._cluster_stub.__enter__() + self._operations_stub.__enter__() + self._table_stub.__enter__() + + def stop(self): + """Closes all the open gRPC clients.""" + if not self.is_started(): + return + + # When exit-ing, we pass None as the exception type, value and + # traceback to __exit__. + self._data_stub.__exit__(None, None, None) + if self._admin: + self._cluster_stub.__exit__(None, None, None) + self._operations_stub.__exit__(None, None, None) + self._table_stub.__exit__(None, None, None) + + self._data_stub = None + self._cluster_stub = None + self._operations_stub = None + self._table_stub = None diff --git a/gcloud/bigtable/test_client.py b/gcloud/bigtable/test_client.py index 2f1e1c3be1dec..71bb31a5ab6e3 100644 --- a/gcloud/bigtable/test_client.py +++ b/gcloud/bigtable/test_client.py @@ -203,6 +203,238 @@ def test_table_stub_unset_failure(self): with self.assertRaises(ValueError): getattr(client, 'table_stub') + def test__make_data_stub(self): + from gcloud._testing import _Monkey + from gcloud.bigtable import client as MUT + from gcloud.bigtable.client import DATA_API_HOST + from gcloud.bigtable.client import DATA_API_PORT + from gcloud.bigtable.client import DATA_STUB_FACTORY + + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + + fake_stub = object() + make_stub_args = [] + + def mock_make_stub(*args): + make_stub_args.append(args) + return fake_stub + + with _Monkey(MUT, make_stub=mock_make_stub): + result = client._make_data_stub() + + self.assertTrue(result is fake_stub) + self.assertEqual(make_stub_args, [ + ( + client, + DATA_STUB_FACTORY, + DATA_API_HOST, + DATA_API_PORT, + ), + ]) + + def test__make_cluster_stub(self): + from gcloud._testing import _Monkey + from gcloud.bigtable import client as MUT + from gcloud.bigtable.client import CLUSTER_ADMIN_HOST + from gcloud.bigtable.client import CLUSTER_ADMIN_PORT + from gcloud.bigtable.client import CLUSTER_STUB_FACTORY + + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + + fake_stub = object() + make_stub_args = [] + + def mock_make_stub(*args): + make_stub_args.append(args) + return fake_stub + + with _Monkey(MUT, make_stub=mock_make_stub): + result = client._make_cluster_stub() + + self.assertTrue(result is fake_stub) + self.assertEqual(make_stub_args, [ + ( + client, + CLUSTER_STUB_FACTORY, + CLUSTER_ADMIN_HOST, + CLUSTER_ADMIN_PORT, + ), + ]) + + def test__make_operations_stub(self): + from gcloud._testing import _Monkey + from gcloud.bigtable import client as MUT + from gcloud.bigtable.client import CLUSTER_ADMIN_HOST + from gcloud.bigtable.client import CLUSTER_ADMIN_PORT + from gcloud.bigtable.client import OPERATIONS_STUB_FACTORY + + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + + fake_stub = object() + make_stub_args = [] + + def mock_make_stub(*args): + make_stub_args.append(args) + return fake_stub + + with _Monkey(MUT, make_stub=mock_make_stub): + result = client._make_operations_stub() + + self.assertTrue(result is fake_stub) + self.assertEqual(make_stub_args, [ + ( + client, + OPERATIONS_STUB_FACTORY, + CLUSTER_ADMIN_HOST, + CLUSTER_ADMIN_PORT, + ), + ]) + + def test__make_table_stub(self): + from gcloud._testing import _Monkey + from gcloud.bigtable import client as MUT + from gcloud.bigtable.client import TABLE_ADMIN_HOST + from gcloud.bigtable.client import TABLE_ADMIN_PORT + from gcloud.bigtable.client import TABLE_STUB_FACTORY + + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + + fake_stub = object() + make_stub_args = [] + + def mock_make_stub(*args): + make_stub_args.append(args) + return fake_stub + + with _Monkey(MUT, make_stub=mock_make_stub): + result = client._make_table_stub() + + self.assertTrue(result is fake_stub) + self.assertEqual(make_stub_args, [ + ( + client, + TABLE_STUB_FACTORY, + TABLE_ADMIN_HOST, + TABLE_ADMIN_PORT, + ), + ]) + + def test_is_started(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + + self.assertFalse(client.is_started()) + client._data_stub = object() + self.assertTrue(client.is_started()) + client._data_stub = None + self.assertFalse(client.is_started()) + + def _start_method_helper(self, admin): + from gcloud._testing import _Monkey + from gcloud.bigtable import client as MUT + + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=admin) + + stub = _FakeStub() + make_stub_args = [] + + def mock_make_stub(*args): + make_stub_args.append(args) + return stub + + with _Monkey(MUT, make_stub=mock_make_stub): + client.start() + + self.assertTrue(client._data_stub is stub) + if admin: + self.assertTrue(client._cluster_stub is stub) + self.assertTrue(client._operations_stub is stub) + self.assertTrue(client._table_stub is stub) + self.assertEqual(stub._entered, 4) + self.assertEqual(len(make_stub_args), 4) + else: + self.assertTrue(client._cluster_stub is None) + self.assertTrue(client._operations_stub is None) + self.assertTrue(client._table_stub is None) + self.assertEqual(stub._entered, 1) + self.assertEqual(len(make_stub_args), 1) + self.assertEqual(stub._exited, []) + + def test_start_non_admin(self): + self._start_method_helper(admin=False) + + def test_start_with_admin(self): + self._start_method_helper(admin=True) + + def test_start_while_started(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + client._data_stub = data_stub = object() + self.assertTrue(client.is_started()) + client.start() + + # Make sure the stub did not change. + self.assertEqual(client._data_stub, data_stub) + + def _stop_method_helper(self, admin): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=admin) + + stub1 = _FakeStub() + stub2 = _FakeStub() + client._data_stub = stub1 + client._cluster_stub = stub2 + client._operations_stub = stub2 + client._table_stub = stub2 + client.stop() + self.assertTrue(client._data_stub is None) + self.assertTrue(client._cluster_stub is None) + self.assertTrue(client._operations_stub is None) + self.assertTrue(client._table_stub is None) + self.assertEqual(stub1._entered, 0) + self.assertEqual(stub2._entered, 0) + exc_none_triple = (None, None, None) + self.assertEqual(stub1._exited, [exc_none_triple]) + if admin: + self.assertEqual(stub2._exited, [exc_none_triple] * 3) + else: + self.assertEqual(stub2._exited, []) + + def test_stop_non_admin(self): + self._stop_method_helper(admin=False) + + def test_stop_with_admin(self): + self._stop_method_helper(admin=True) + + def test_stop_while_stopped(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + self.assertFalse(client.is_started()) + + # This is a bit hacky. We set the cluster stub protected value + # since it isn't used in is_started() and make sure that stop + # doesn't reset this value to None. + client._cluster_stub = cluster_stub = object() + client.stop() + # Make sure the cluster stub did not change. + self.assertEqual(client._cluster_stub, cluster_stub) + class _Credentials(object): @@ -211,3 +443,18 @@ class _Credentials(object): def create_scoped(self, scope): self._scopes = scope return self + + +class _FakeStub(object): + + def __init__(self): + self._entered = 0 + self._exited = [] + + def __enter__(self): + self._entered += 1 + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self._exited.append((exc_type, exc_val, exc_tb)) + return True diff --git a/setup.py b/setup.py index 0cccb9846624c..2c2067af2c659 100644 --- a/setup.py +++ b/setup.py @@ -15,7 +15,7 @@ 'google-apitools', 'httplib2 >= 0.9.1', 'oauth2client >= 1.4.6', - 'protobuf == 3.0.0-alpha-1', + 'protobuf >= 3.0.0a3', 'pycrypto', 'six', ]