Skip to content

Commit

Permalink
Adding ability to start/stop clients.
Browse files Browse the repository at this point in the history
This creates gRPC stubs for talking to each Bigtable
service needed by the client.
  • Loading branch information
dhermes committed Sep 28, 2015
1 parent 8217cb2 commit e33fe3d
Show file tree
Hide file tree
Showing 3 changed files with 346 additions and 1 deletion.
98 changes: 98 additions & 0 deletions gcloud/bigtable/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand Down Expand Up @@ -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
247 changes: 247 additions & 0 deletions gcloud/bigtable/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):

Expand All @@ -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
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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',
]
Expand Down

0 comments on commit e33fe3d

Please sign in to comment.