diff --git a/gcloud/bigtable/client.py b/gcloud/bigtable/client.py index 6d32554aa24e..32b37b691bc4 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' @@ -130,6 +143,12 @@ def __init__(self, project=None, credentials=None, self.user_agent = user_agent self.timeout_seconds = timeout_seconds + # These will be set in start(). + self._data_stub = None + self._cluster_stub = None + self._operations_stub = None + self._table_stub = None + @property def credentials(self): """Getter for client's credentials. @@ -139,3 +158,173 @@ def credentials(self): :returns: The credentials stored on the client. """ return self._credentials + + @property + def project_name(self): + """Project name to be used with Cluster Admin API. + + .. note:: + + This property will not change if ``project`` does not, but the + return value is not cached. + + The project name is of the form + + ``"projects/{project_id}"`` + + :rtype: str + :returns: The project name to be used with the Cloud Bigtable Admin + API RPC service. + """ + return 'projects/' + self.project + + @property + def data_stub(self): + """Getter for the gRPC stub used for the Data API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + :raises: :class:`ValueError ` if the current + client has not been :meth:`start`-ed. + """ + if self._data_stub is None: + raise ValueError('Client has not been started.') + return self._data_stub + + @property + def cluster_stub(self): + """Getter for the gRPC stub used for the Cluster Admin API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. + """ + if not self._admin: + raise ValueError('Client is not an admin client.') + if self._cluster_stub is None: + raise ValueError('Client has not been started.') + return self._cluster_stub + + @property + def operations_stub(self): + """Getter for the gRPC stub used for the Operations API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. + """ + if not self._admin: + raise ValueError('Client is not an admin client.') + if self._operations_stub is None: + raise ValueError('Client has not been started.') + return self._operations_stub + + @property + def table_stub(self): + """Getter for the gRPC stub used for the Table Admin API. + + :rtype: :class:`grpc.early_adopter.implementations._Stub` + :returns: A gRPC stub object. + :raises: :class:`ValueError ` if the current + client is not an admin client or if it has not been + :meth:`start`-ed. + """ + if not self._admin: + raise ValueError('Client is not an admin client.') + 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 + + # NOTE: We __enter__ the stubs more-or-less permanently. This is + # because only after entering the context managers is the + # connection created. We don't want to immediately close + # those connections since the client will make many + # requests with it over HTTP/2. + 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 8755048b06a8..71bb31a5ab6e 100644 --- a/gcloud/bigtable/test_client.py +++ b/gcloud/bigtable/test_client.py @@ -46,6 +46,11 @@ def _constructor_test_helper(self, expected_scopes, creds, self.assertEqual(client.project, PROJECT) self.assertEqual(client.timeout_seconds, timeout_seconds) self.assertEqual(client.user_agent, user_agent) + # Check stubs are set (but null) + self.assertEqual(client._data_stub, None) + self.assertEqual(client._cluster_stub, None) + self.assertEqual(client._operations_stub, None) + self.assertEqual(client._table_stub, None) def test_constructor_default_scopes(self): from gcloud.bigtable import client as MUT @@ -105,6 +110,331 @@ def test_credentials_getter(self): client = self._makeOne(project=project, credentials=credentials) self.assertTrue(client.credentials is credentials) + def test_project_name_property(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + project_name = 'projects/' + project + self.assertEqual(client.project_name, project_name) + + def test_data_stub_getter(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + client._data_stub = object() + self.assertTrue(client.data_stub is client._data_stub) + + def test_data_stub_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials) + with self.assertRaises(ValueError): + getattr(client, 'data_stub') + + def test_cluster_stub_getter(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + client._cluster_stub = object() + self.assertTrue(client.cluster_stub is client._cluster_stub) + + def test_cluster_stub_non_admin_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=False) + with self.assertRaises(ValueError): + getattr(client, 'cluster_stub') + + def test_cluster_stub_unset_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + with self.assertRaises(ValueError): + getattr(client, 'cluster_stub') + + def test_operations_stub_getter(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + client._operations_stub = object() + self.assertTrue(client.operations_stub is client._operations_stub) + + def test_operations_stub_non_admin_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=False) + with self.assertRaises(ValueError): + getattr(client, 'operations_stub') + + def test_operations_stub_unset_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + with self.assertRaises(ValueError): + getattr(client, 'operations_stub') + + def test_table_stub_getter(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + client._table_stub = object() + self.assertTrue(client.table_stub is client._table_stub) + + def test_table_stub_non_admin_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=False) + with self.assertRaises(ValueError): + getattr(client, 'table_stub') + + def test_table_stub_unset_failure(self): + credentials = _Credentials() + project = 'PROJECT' + client = self._makeOne(project=project, credentials=credentials, + admin=True) + 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): @@ -113,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