Add max_receive_message_length for larger rows.

Author: daspecster

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)

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):

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,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.
@@ -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)
 
 

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,21 +691,23 @@ 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)
 
         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):

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)