Skip to content

Commit

Permalink
Add max_receive_message_length for larger rows.
Browse files Browse the repository at this point in the history
  • Loading branch information
daspecster committed Dec 29, 2016
1 parent 455eaf6 commit 878fb58
Show file tree
Hide file tree
Showing 4 changed files with 77 additions and 8 deletions.
11 changes: 10 additions & 1 deletion bigtable/google/cloud/bigtable/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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)
Expand Down
8 changes: 7 additions & 1 deletion bigtable/unit_tests/test_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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, [
(
Expand All @@ -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):
Expand Down
25 changes: 19 additions & 6 deletions core/google/cloud/_helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -480,22 +480,30 @@ 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,
target,
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.
Expand All @@ -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)


Expand Down
41 changes: 41 additions & 0 deletions system_tests/bigtable.py
Original file line number Diff line number Diff line change
Expand Up @@ -356,6 +356,47 @@ 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 = '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]
value = column[0].value
self.assertIsNotNone(value)
self.assertEqual(len(value), number_of_bytes)

def test_read_large_cells_over_limit(self):
row = self._table.row(ROW_KEY)
self.rows_to_delete.append(row)

data = '1' * 51 * 1024 * 1024 # 51MB.
row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, data)
row.commit()
row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, data)
row.commit()

with self.assertRaises(Exception):
partial_row_data = self._table.read_row(ROW_KEY)


def test_read_large_cell_over_limit(self):
row = self._table.row(ROW_KEY)
self.rows_to_delete.append(row)

data = '1' * 101 * 1024 * 1024 # 11MB of 1's.
row.set_cell(COLUMN_FAMILY_ID1, COL_NAME1, data)
with self.assertRaises(Exception):
row.commit()

def test_read_row(self):
row = self._table.row(ROW_KEY)
self.rows_to_delete.append(row)
Expand Down

0 comments on commit 878fb58

Please sign in to comment.