fix: avoid consuming pending null values when merging (#286)

larkee · web-flow · commit c6cba9fbe4c7 · 2021-03-24T13:21:04.000+11:00
* test: add test for string array with pending null

* fix: avoid consuming pending null values when merging

* test: match implementation to test names

Co-authored-by: larkee &lt;larkee@users.noreply.github.com&gt;
diff --git a/google/cloud/spanner_v1/streamed.py b/google/cloud/spanner_v1/streamed.py
@@ -258,13 +258,17 @@ def _merge_array(lhs, rhs, type_):
         lhs.append(first)
     else:
         last = lhs.pop()
-        try:
-            merged = _merge_by_type(last, first, element_type)
-        except Unmergeable:
+        if last.HasField("null_value"):
             lhs.append(last)
             lhs.append(first)
         else:
-            lhs.append(merged)
+            try:
+                merged = _merge_by_type(last, first, element_type)
+            except Unmergeable:
+                lhs.append(last)
+                lhs.append(first)
+            else:
+                lhs.append(merged)
     return Value(list_value=ListValue(values=(lhs + rhs)))
 
 
@@ -284,13 +288,17 @@ def _merge_struct(lhs, rhs, type_):
         lhs.append(first)
     else:
         last = lhs.pop()
-        try:
-            merged = _merge_by_type(last, first, candidate_type)
-        except Unmergeable:
+        if last.HasField("null_value"):
             lhs.append(last)
             lhs.append(first)
         else:
-            lhs.append(merged)
+            try:
+                merged = _merge_by_type(last, first, candidate_type)
+            except Unmergeable:
+                lhs.append(last)
+                lhs.append(first)
+            else:
+                lhs.append(merged)
     return Value(list_value=ListValue(values=lhs + rhs))
 
 
diff --git a/tests/unit/test_streamed.py b/tests/unit/test_streamed.py
@@ -336,11 +336,11 @@ def test__merge_chunk_array_of_string(self):
         FIELDS = [self._make_array_field("name", element_type_code=TypeCode.STRING)]
         streamed._metadata = self._make_result_set_metadata(FIELDS)
         streamed._pending_chunk = self._make_list_value([u"A", u"B", u"C"])
-        chunk = self._make_list_value([None, u"D", u"E"])
+        chunk = self._make_list_value([u"D", u"E"])
 
         merged = streamed._merge_chunk(chunk)
 
-        expected = self._make_list_value([u"A", u"B", u"C", None, u"D", u"E"])
+        expected = self._make_list_value([u"A", u"B", u"CD", u"E"])
         self.assertEqual(merged, expected)
         self.assertIsNone(streamed._pending_chunk)
 
@@ -352,11 +352,25 @@ def test__merge_chunk_array_of_string_with_null(self):
         FIELDS = [self._make_array_field("name", element_type_code=TypeCode.STRING)]
         streamed._metadata = self._make_result_set_metadata(FIELDS)
         streamed._pending_chunk = self._make_list_value([u"A", u"B", u"C"])
-        chunk = self._make_list_value([u"D", u"E"])
+        chunk = self._make_list_value([None, u"D", u"E"])
 
         merged = streamed._merge_chunk(chunk)
 
-        expected = self._make_list_value([u"A", u"B", u"CD", u"E"])
+        expected = self._make_list_value([u"A", u"B", u"C", None, u"D", u"E"])
+        self.assertEqual(merged, expected)
+        self.assertIsNone(streamed._pending_chunk)
+
+    def test__merge_chunk_array_of_string_with_null_pending(self):
+        from google.cloud.spanner_v1 import TypeCode
+
+        iterator = _MockCancellableIterator()
+        streamed = self._make_one(iterator)
+        FIELDS = [self._make_array_field("name", element_type_code=TypeCode.STRING)]
+        streamed._metadata = self._make_result_set_metadata(FIELDS)
+        streamed._pending_chunk = self._make_list_value([u"A", u"B", u"C", None])
+        chunk = self._make_list_value([u"D", u"E"])
+        merged = streamed._merge_chunk(chunk)
+        expected = self._make_list_value([u"A", u"B", u"C", None, u"D", u"E"])
         self.assertEqual(merged, expected)
         self.assertIsNone(streamed._pending_chunk)
 
