"""
CREATE TABLE Test (id STRING(36) NOT NULL, megafield BYTES(MAX)) PRIMARY KEY (id)
"""

import base64
from google.cloud import spanner
from google.auth.credentials import AnonymousCredentials

###################################
# HOTFIX
###################################
from google.cloud.spanner_v1.streamed import StreamedResultSet, _merge_by_type

def _merge_chunk(self, value):
    """Merge pending chunk with next value.

    :type value: :class:`~google.protobuf.struct_pb2.Value`
    :param value: continuation of chunked value from previous
                  partial result set.

    :rtype: :class:`~google.protobuf.struct_pb2.Value`
    :returns: the merged value
    """
    current_column = len(self._current_row)
    field = self.fields[current_column]
    merged = _merge_by_type(self._pending_chunk, value, field.type_)
    self._pending_chunk = None
    # Bug fix:
    return merged  #_parse_value(merged, field.type_)

# Uncomment this to fix the bug:
# StreamedResultSet._merge_chunk = _merge_chunk
###################################
# END OF HOTFIX
###################################

instance_id = 'test'
database_id = 'test-db'

spanner_client = spanner.Client(
    project='test',
    client_options={"api_endpoint": 'localhost:9010'},
    credentials=AnonymousCredentials()
)

instance = spanner_client.instance(instance_id)
database = instance.database(database_id)

# This must be large enough that the SDK will split the megafield payload across two query chunks
# and try to recombine them, causing the error:
data = base64.standard_b64encode(("a" * 1000000).encode("utf8"))

with database.batch() as batch:
    batch.insert(
        table="Test",
        columns=("id", "megafield"),
        values=[
            (1, data),
        ],
    )

with database.snapshot() as snapshot:
    results = snapshot.execute_sql(
        "SELECT * FROM Test"
    )

    for row in results:
        print("Id: ", row[0])
        print("Megafield: ", row[1][:100])