perf: Skip gRPC trailers for StreamingRead & ExecuteStreamingSql

Author: sakthivelmanii

google/cloud/spanner_v1/streamed.py

@@ -54,6 +54,7 @@ def __init__(
         self._column_info = column_info  # Column information
         self._field_decoders = None
         self._lazy_decode = lazy_decode  # Return protobuf values
+        self._done = False
 
     @property
     def fields(self):
@@ -159,11 +160,16 @@ def _consume_next(self):
 
         self._merge_values(values)
 
+        if response_pb.last:
+            self._done = True
+
     def __iter__(self):
         while True:
             iter_rows, self._rows[:] = self._rows[:], ()
             while iter_rows:
                 yield iter_rows.pop(0)
+            if self._done:
+                return
             try:
                 self._consume_next()
             except StopIteration:

tests/unit/test_streamed.py

@@ -124,12 +124,12 @@ def _make_result_set_stats(query_plan=None, **kw):
 
     @staticmethod
     def _make_partial_result_set(
-        values, metadata=None, stats=None, chunked_value=False
+        values, metadata=None, stats=None, chunked_value=False, last=False
     ):
         from google.cloud.spanner_v1 import PartialResultSet
 
         results = PartialResultSet(
-            metadata=metadata, stats=stats, chunked_value=chunked_value
+            metadata=metadata, stats=stats, chunked_value=chunked_value, last=last
         )
         for v in values:
             results.values.append(v)
@@ -164,6 +164,35 @@ def test__merge_chunk_bool(self):
         with self.assertRaises(Unmergeable):
             streamed._merge_chunk(chunk)
 
+    def test__PartialResultSetWithLastFlag(self):
+        from google.cloud.spanner_v1 import TypeCode
+
+        fields = [
+            self._make_scalar_field("ID", TypeCode.INT64),
+            self._make_scalar_field("NAME", TypeCode.STRING),
+        ]
+        for length in range(4, 6):
+            metadata = self._make_result_set_metadata(fields)
+            result_sets = [
+                self._make_partial_result_set([self._make_value(0), "google_0"],
+                                              metadata=metadata)]
+            for i in range(1, 5):
+                bares = [i]
+                values = [[self._make_value(bare), "google_" + str(bare)] for
+                          bare in bares]
+                result_sets.append(self._make_partial_result_set(*values,
+                                                                 metadata=metadata,
+                                                                 last=(i == length - 1)))
+
+            iterator = _MockCancellableIterator(*result_sets)
+            streamed = self._make_one(iterator)
+            count = 0
+            for row in streamed:
+                self.assertEqual(row[0], count)
+                self.assertEqual(row[1], "google_" + str(count))
+                count += 1
+            self.assertEqual(count, length)
+
     def test__merge_chunk_numeric(self):
         from google.cloud.spanner_v1 import TypeCode