From c6cba9fbe4c717f1f8e2a97e3f76bfe6b956e55b Mon Sep 17 00:00:00 2001 From: larkee <31196561+larkee@users.noreply.github.com> Date: Wed, 24 Mar 2021 13:21:04 +1100 Subject: [PATCH] fix: avoid consuming pending null values when merging (#286) * 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 --- google/cloud/spanner_v1/streamed.py | 24 ++++++++++++++++-------- tests/unit/test_streamed.py | 22 ++++++++++++++++++---- 2 files changed, 34 insertions(+), 12 deletions(-) diff --git a/google/cloud/spanner_v1/streamed.py b/google/cloud/spanner_v1/streamed.py index 20d814ebed..fbcca77795 100644 --- 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 index 63f3bf81fe..7b12f6a94b 100644 --- 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)