diff --git a/google/cloud/spanner_v1/_helpers.py b/google/cloud/spanner_v1/_helpers.py index 27e53200ed..352f72ccaa 100644 --- a/google/cloud/spanner_v1/_helpers.py +++ b/google/cloud/spanner_v1/_helpers.py @@ -20,6 +20,7 @@ import time import base64 import threading +import uuid from google.protobuf.struct_pb2 import ListValue from google.protobuf.struct_pb2 import Value @@ -219,6 +220,8 @@ def _make_value_pb(value): return Value(null_value="NULL_VALUE") else: return Value(string_value=base64.b64encode(value)) + if isinstance(value, uuid.UUID): + return Value(string_value=str(value)) raise ValueError("Unknown type: %s" % (value,)) @@ -320,6 +323,8 @@ def _get_type_decoder(field_type, field_name, column_info=None): return _parse_numeric elif type_code == TypeCode.JSON: return _parse_json + elif type_code == TypeCode.UUID: + return _parse_uuid elif type_code == TypeCode.PROTO: return lambda value_pb: _parse_proto(value_pb, column_info, field_name) elif type_code == TypeCode.ENUM: @@ -400,6 +405,10 @@ def _parse_json(value_pb): return JsonObject.from_str(value_pb.string_value) +def _parse_uuid(value_pb): + return uuid.UUID(value_pb.string_value) + + def _parse_proto(value_pb, column_info, field_name): bytes_value = base64.b64decode(value_pb.string_value) if column_info is not None and column_info.get(field_name) is not None: diff --git a/google/cloud/spanner_v1/param_types.py b/google/cloud/spanner_v1/param_types.py index 5416a26d61..c983531b06 100644 --- a/google/cloud/spanner_v1/param_types.py +++ b/google/cloud/spanner_v1/param_types.py @@ -33,6 +33,7 @@ TIMESTAMP = Type(code=TypeCode.TIMESTAMP) NUMERIC = Type(code=TypeCode.NUMERIC) JSON = Type(code=TypeCode.JSON) +UUID = Type(code=TypeCode.UUID) PG_NUMERIC = Type(code=TypeCode.NUMERIC, type_annotation=TypeAnnotationCode.PG_NUMERIC) PG_JSONB = Type(code=TypeCode.JSON, type_annotation=TypeAnnotationCode.PG_JSONB) PG_OID = Type(code=TypeCode.INT64, type_annotation=TypeAnnotationCode.PG_OID) diff --git a/google/cloud/spanner_v1/streamed.py b/google/cloud/spanner_v1/streamed.py index 7c067e97b6..145c636c12 100644 --- a/google/cloud/spanner_v1/streamed.py +++ b/google/cloud/spanner_v1/streamed.py @@ -392,6 +392,7 @@ def _merge_struct(lhs, rhs, type_): TypeCode.JSON: _merge_string, TypeCode.PROTO: _merge_string, TypeCode.ENUM: _merge_string, + TypeCode.UUID: _merge_string, } diff --git a/tests/system/test_session_api.py b/tests/system/test_session_api.py index 4de0e681f6..8d2796a3e2 100644 --- a/tests/system/test_session_api.py +++ b/tests/system/test_session_api.py @@ -19,6 +19,7 @@ import struct import threading import time +import uuid import pytest import grpc @@ -2825,6 +2826,16 @@ def test_execute_sql_returning_transfinite_floats(sessions_database, not_postgre assert math.isnan(float_array[2]) +def test_execute_sql_w_uuid_bindings(sessions_database, database_dialect): + _bind_test_helper( + sessions_database, + database_dialect, + spanner_v1.param_types.UUID, + uuid.uuid4(), + [uuid.uuid4(), uuid.uuid4()], + ) + + def test_partition_query(sessions_database, not_emulator): row_count = 40 sql = f"SELECT * FROM {_sample_data.TABLE}" diff --git a/tests/unit/test__helpers.py b/tests/unit/test__helpers.py index ecc8018648..dd560b73f4 100644 --- a/tests/unit/test__helpers.py +++ b/tests/unit/test__helpers.py @@ -14,6 +14,7 @@ import unittest +import uuid import mock @@ -738,6 +739,18 @@ def test_w_proto_enum(self): self._callFUT(value_pb, field_type, field_name, column_info), VALUE ) + def test_w_uuid(self): + from google.protobuf.struct_pb2 import Value + from google.cloud.spanner_v1 import Type + from google.cloud.spanner_v1 import TypeCode + + VALUE = uuid.uuid4() + field_type = Type(code=TypeCode.UUID) + field_name = "uuid_column" + value_pb = Value(string_value=str(VALUE)) + + self.assertEqual(self._callFUT(value_pb, field_type, field_name), VALUE) + class Test_parse_list_value_pbs(unittest.TestCase): def _callFUT(self, *args, **kw):