From 43c190bc694d56e0c57d96dbaa7fc48117f3c971 Mon Sep 17 00:00:00 2001
From: Sri Harsha CH <57220027+harshachinta@users.noreply.github.com>
Date: Mon, 28 Oct 2024 19:22:25 +0530
Subject: [PATCH] fix: add PROTO in streaming chunks (#1213)

b/372956316
When the row size exceeds a certain limit, the rows are divided into chunks and sent to the client in multiple parts. The client is responsible for merging these chunks to reconstruct the full row.
However, for PROTO and ENUM types, this chunk-merging logic was not implemented, causing a KeyError: 13 when attempting to merge proto chunks.


#### Sample to reproduce the test case
[Python file](https://gist.github.com/harshachinta/95a81eeda81c422814353a5995d01e20)
[proto file
](https://gist.github.com/harshachinta/fd15bf558bd4f40443411ddd164638cc)

#### Steps to generate descriptors.pb and code file from proto
```
protoc --proto_path=testdata/ --include_imports --descriptor_set_out=testdata/descriptors.pb --python_out=testdata/ testdata/wrapper.proto
```
---
 google/cloud/spanner_v1/streamed.py |  2 ++
 tests/unit/test_streamed.py         | 40 +++++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/google/cloud/spanner_v1/streamed.py b/google/cloud/spanner_v1/streamed.py
index 03acc9010a..89bde0e334 100644
--- a/google/cloud/spanner_v1/streamed.py
+++ b/google/cloud/spanner_v1/streamed.py
@@ -345,6 +345,8 @@ def _merge_struct(lhs, rhs, type_):
     TypeCode.TIMESTAMP: _merge_string,
     TypeCode.NUMERIC: _merge_string,
     TypeCode.JSON: _merge_string,
+    TypeCode.PROTO: _merge_string,
+    TypeCode.ENUM: _merge_string,
 }
 
 
diff --git a/tests/unit/test_streamed.py b/tests/unit/test_streamed.py
index 85dcb40026..83aa25a9d1 100644
--- a/tests/unit/test_streamed.py
+++ b/tests/unit/test_streamed.py
@@ -272,6 +272,46 @@ def test__merge_chunk_string_w_bytes(self):
         )
         self.assertIsNone(streamed._pending_chunk)
 
+    def test__merge_chunk_proto(self):
+        from google.cloud.spanner_v1 import TypeCode
+
+        iterator = _MockCancellableIterator()
+        streamed = self._make_one(iterator)
+        FIELDS = [self._make_scalar_field("proto", TypeCode.PROTO)]
+        streamed._metadata = self._make_result_set_metadata(FIELDS)
+        streamed._pending_chunk = self._make_value(
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAAAAAA"
+            "6fptVAAAACXBIWXMAAAsTAAALEwEAmpwYAAAA\n"
+        )
+        chunk = self._make_value(
+            "B3RJTUUH4QQGFwsBTL3HMwAAABJpVFh0Q29tbWVudAAAAAAAU0FNUExF"
+            "MG3E+AAAAApJREFUCNdj\nYAAAAAIAAeIhvDMAAAAASUVORK5CYII=\n"
+        )
+
+        merged = streamed._merge_chunk(chunk)
+
+        self.assertEqual(
+            merged.string_value,
+            "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAAAAAA6fptVAAAACXBIWXMAAAsTAAAL"
+            "EwEAmpwYAAAA\nB3RJTUUH4QQGFwsBTL3HMwAAABJpVFh0Q29tbWVudAAAAAAAU0"
+            "FNUExFMG3E+AAAAApJREFUCNdj\nYAAAAAIAAeIhvDMAAAAASUVORK5CYII=\n",
+        )
+        self.assertIsNone(streamed._pending_chunk)
+
+    def test__merge_chunk_enum(self):
+        from google.cloud.spanner_v1 import TypeCode
+
+        iterator = _MockCancellableIterator()
+        streamed = self._make_one(iterator)
+        FIELDS = [self._make_scalar_field("age", TypeCode.ENUM)]
+        streamed._metadata = self._make_result_set_metadata(FIELDS)
+        streamed._pending_chunk = self._make_value(42)
+        chunk = self._make_value(13)
+
+        merged = streamed._merge_chunk(chunk)
+        self.assertEqual(merged.string_value, "4213")
+        self.assertIsNone(streamed._pending_chunk)
+
     def test__merge_chunk_array_of_bool(self):
         from google.cloud.spanner_v1 import TypeCode