From 0ed6d3a29b18fe00c1e73383f20ceacdcd570a58 Mon Sep 17 00:00:00 2001 From: Thomas Schultz Date: Fri, 23 Dec 2016 07:58:49 -0500 Subject: [PATCH] Add max_receive_message_length for larger rows. --- bigtable/google/cloud/bigtable/client.py | 11 +++++++- bigtable/unit_tests/test_client.py | 8 +++++- core/google/cloud/_helpers.py | 25 +++++++++++++----- core/unit_tests/test__helpers.py | 33 ++++++++++++++++++++++-- system_tests/bigtable.py | 17 ++++++++++++ 5 files changed, 84 insertions(+), 10 deletions(-) diff --git a/bigtable/google/cloud/bigtable/client.py b/bigtable/google/cloud/bigtable/client.py index 037b6efa15ea..38a103c7c005 100644 --- a/bigtable/google/cloud/bigtable/client.py +++ b/bigtable/google/cloud/bigtable/client.py @@ -65,6 +65,14 @@ READ_ONLY_SCOPE = 'https://www.googleapis.com/auth/bigtable.data.readonly' """Scope for reading table data.""" +# NOTE: 'grpc.max_message_length' will no longer be recognized in +# grpcio 1.1 and later. +_MAX_MSG_LENGTH_100MB = 100 * 1024 * 1024 +_GRPC_MAX_LENGTH_OPTIONS = ( + ('grpc.max_message_length', _MAX_MSG_LENGTH_100MB), + ('grpc.max_receive_message_length', _MAX_MSG_LENGTH_100MB), +) + def _make_data_stub(client): """Creates gRPC stub to make requests to the Data API. @@ -77,7 +85,8 @@ def _make_data_stub(client): """ if client.emulator_host is None: return make_secure_stub(client.credentials, client.user_agent, - bigtable_pb2.BigtableStub, DATA_API_HOST) + bigtable_pb2.BigtableStub, DATA_API_HOST, + extra_options=_GRPC_MAX_LENGTH_OPTIONS) else: return make_insecure_stub(bigtable_pb2.BigtableStub, client.emulator_host) diff --git a/bigtable/unit_tests/test_client.py b/bigtable/unit_tests/test_client.py index 4c912b4794d3..24fee441a6c9 100644 --- a/bigtable/unit_tests/test_client.py +++ b/bigtable/unit_tests/test_client.py @@ -46,13 +46,18 @@ def test_without_emulator(self): fake_stub = object() make_secure_stub_args = [] - def mock_make_secure_stub(*args): + def mock_make_secure_stub(*args, **kwargs): make_secure_stub_args.append(args) + make_secure_stub_args.append(kwargs) return fake_stub with _Monkey(MUT, make_secure_stub=mock_make_secure_stub): result = self._call_fut(client) + extra_options = {'extra_options': ( + ('grpc.max_message_length', 104857600), + ('grpc.max_receive_message_length', 104857600) + )} self.assertIs(result, fake_stub) self.assertEqual(make_secure_stub_args, [ ( @@ -61,6 +66,7 @@ def mock_make_secure_stub(*args): MUT.bigtable_pb2.BigtableStub, MUT.DATA_API_HOST, ), + extra_options, ]) def test_with_emulator(self): diff --git a/core/google/cloud/_helpers.py b/core/google/cloud/_helpers.py index f8d535b47099..03f6d2fdb60c 100644 --- a/core/google/cloud/_helpers.py +++ b/core/google/cloud/_helpers.py @@ -465,7 +465,7 @@ def _name_from_project_path(path, project, template): return match.group('name') -def make_secure_channel(credentials, user_agent, host): +def make_secure_channel(credentials, user_agent, host, extra_options=None): """Makes a secure channel for an RPC service. Uses / depends on gRPC. @@ -480,14 +480,21 @@ def make_secure_channel(credentials, user_agent, host): :type host: str :param host: The host for the service. + :type extra_options: tuple + :param extra_options: (Optional) Extra gRPC options used when creating the + channel. + :rtype: :class:`grpc._channel.Channel` :returns: gRPC secure channel with credentials attached. """ target = '%s:%d' % (host, http_client.HTTPS_PORT) http_request = google_auth_httplib2.Request(http=httplib2.Http()) - options = ( - ('grpc.primary_user_agent', user_agent), - ) + + user_agent_option = ('grpc.primary_user_agent', user_agent) + if extra_options is not None: + options = (user_agent_option,) + extra_options + else: + options = (user_agent_option,) return google.auth.transport.grpc.secure_authorized_channel( credentials, http_request, @@ -495,7 +502,8 @@ def make_secure_channel(credentials, user_agent, host): options=options) -def make_secure_stub(credentials, user_agent, stub_class, host): +def make_secure_stub(credentials, user_agent, stub_class, host, + extra_options=None): """Makes a secure stub for an RPC service. Uses / depends on gRPC. @@ -513,10 +521,15 @@ def make_secure_stub(credentials, user_agent, stub_class, host): :type host: str :param host: The host for the service. + :type extra_options: tuple + :param extra_options: (Optional) Extra gRPC options passed when creating + the channel. + :rtype: object, instance of ``stub_class`` :returns: The stub object used to make gRPC requests to a given API. """ - channel = make_secure_channel(credentials, user_agent, host) + channel = make_secure_channel(credentials, user_agent, host, + extra_options=extra_options) return stub_class(channel) diff --git a/core/unit_tests/test__helpers.py b/core/unit_tests/test__helpers.py index 59752a6e594f..bc1a4dafb6b7 100644 --- a/core/unit_tests/test__helpers.py +++ b/core/unit_tests/test__helpers.py @@ -644,6 +644,33 @@ def test_it(self): secure_authorized_channel.assert_called_once_with( credentials, mock.ANY, expected_target, options=expected_options) + def test_extra_options(self): + from six.moves import http_client + + credentials = object() + host = 'HOST' + user_agent = 'USER_AGENT' + extra_options = (('some', 'option'),) + + secure_authorized_channel_patch = mock.patch( + 'google.auth.transport.grpc.secure_authorized_channel', + autospec=True) + + with secure_authorized_channel_patch as secure_authorized_channel: + result = self._call_fut(credentials, user_agent, host, + extra_options) + + self.assertIs(result, secure_authorized_channel.return_value) + + expected_target = '%s:%d' % (host, http_client.HTTPS_PORT) + expected_options = ( + ('grpc.primary_user_agent', user_agent), + extra_options[0], + ) + + secure_authorized_channel.assert_called_once_with( + credentials, mock.ANY, expected_target, options=expected_options) + class Test_make_secure_stub(unittest.TestCase): @@ -664,13 +691,15 @@ def stub_class(channel): channels.append(channel) return result - def mock_channel(*args): + def mock_channel(*args, **kwargs): channel_args.append(args) + channel_args.append(kwargs) return channel_obj credentials = object() user_agent = 'you-sir-age-int' host = 'localhost' + extra_options = {'extra_options': None} with _Monkey(MUT, make_secure_channel=mock_channel): stub = self._call_fut(credentials, user_agent, stub_class, host) @@ -678,7 +707,7 @@ def mock_channel(*args): self.assertIs(stub, result) self.assertEqual(channels, [channel_obj]) self.assertEqual(channel_args, - [(credentials, user_agent, host)]) + [(credentials, user_agent, host), extra_options]) class Test_make_insecure_stub(unittest.TestCase): diff --git a/system_tests/bigtable.py b/system_tests/bigtable.py index bdf48fe64701..a57c1d781d93 100644 --- a/system_tests/bigtable.py +++ b/system_tests/bigtable.py @@ -356,6 +356,23 @@ def _write_to_row(self, row1=None, row2=None, row3=None, row4=None): cell4 = Cell(CELL_VAL4, timestamp4) return cell1, cell2, cell3, cell4 + def test_read_large_cell_limit(self): + row = self._table.row(ROW_KEY) + self.rows_to_delete.append(row) + + number_of_bytes = 10 * 1024 * 1024 + data = b'1' * number_of_bytes # 10MB of 1's. + row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, data) + row.commit() + + # Read back the contents of the row. + partial_row_data = self._table.read_row(ROW_KEY) + self.assertEqual(partial_row_data.row_key, ROW_KEY) + cell = partial_row_data.cells[COLUMN_FAMILY_ID1] + column = cell[COL_NAME1] + self.assertEqual(len(column), 1) + self.assertEqual(column[0].value, data) + def test_read_row(self): row = self._table.row(ROW_KEY) self.rows_to_delete.append(row)