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)
 
 

system_tests/bigtable.py

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