Skip to content

Commit d760c2c

Browse files
Ilya Gurovlarkee
andauthored
feat(db_api): support JSON data type (#627)
Co-authored-by: larkee <31196561+larkee@users.noreply.github.com>
1 parent d769ff8 commit d760c2c

File tree

7 files changed

+57
-25
lines changed

7 files changed

+57
-25
lines changed

google/cloud/spanner_dbapi/cursor.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,7 @@ def execute(self, sql, args=None):
223223
ResultsChecksum(),
224224
classification == parse_utils.STMT_INSERT,
225225
)
226+
226227
(self._result_set, self._checksum,) = self.connection.run_statement(
227228
statement
228229
)

google/cloud/spanner_v1/_helpers.py

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,6 @@
1717
import datetime
1818
import decimal
1919
import math
20-
import json
2120

2221
from google.protobuf.struct_pb2 import ListValue
2322
from google.protobuf.struct_pb2 import Value
@@ -166,9 +165,8 @@ def _make_value_pb(value):
166165
_assert_numeric_precision_and_scale(value)
167166
return Value(string_value=str(value))
168167
if isinstance(value, JsonObject):
169-
return Value(
170-
string_value=json.dumps(value, sort_keys=True, separators=(",", ":"),)
171-
)
168+
return Value(string_value=value.serialize())
169+
172170
raise ValueError("Unknown type: %s" % (value,))
173171

174172

@@ -243,7 +241,7 @@ def _parse_value_pb(value_pb, field_type):
243241
elif type_code == TypeCode.NUMERIC:
244242
return decimal.Decimal(value_pb.string_value)
245243
elif type_code == TypeCode.JSON:
246-
return value_pb.string_value
244+
return JsonObject.from_str(value_pb.string_value)
247245
else:
248246
raise ValueError("Unknown type: %s" % (field_type,))
249247

google/cloud/spanner_v1/data_types.py

Lines changed: 32 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,8 @@
1414

1515
"""Custom data types for spanner."""
1616

17+
import json
18+
1719

1820
class JsonObject(dict):
1921
"""
@@ -22,4 +24,33 @@ class JsonObject(dict):
2224
normal parameters and JSON parameters.
2325
"""
2426

25-
pass
27+
def __init__(self, *args, **kwargs):
28+
self._is_null = (args, kwargs) == ((), {}) or args == (None,)
29+
if not self._is_null:
30+
super(JsonObject, self).__init__(*args, **kwargs)
31+
32+
@classmethod
33+
def from_str(cls, str_repr):
34+
"""Initiate an object from its `str` representation.
35+
36+
Args:
37+
str_repr (str): JSON text representation.
38+
39+
Returns:
40+
JsonObject: JSON object.
41+
"""
42+
if str_repr == "null":
43+
return cls()
44+
45+
return cls(json.loads(str_repr))
46+
47+
def serialize(self):
48+
"""Return the object text representation.
49+
50+
Returns:
51+
str: JSON object text representation.
52+
"""
53+
if self._is_null:
54+
return None
55+
56+
return json.dumps(self, sort_keys=True, separators=(",", ":"))

samples/samples/snippets_test.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,13 @@ def sample_name():
5050

5151
@pytest.fixture(scope="module")
5252
def create_instance_id():
53-
""" Id for the low-cost instance. """
53+
"""Id for the low-cost instance."""
5454
return f"create-instance-{uuid.uuid4().hex[:10]}"
5555

5656

5757
@pytest.fixture(scope="module")
5858
def lci_instance_id():
59-
""" Id for the low-cost instance. """
59+
"""Id for the low-cost instance."""
6060
return f"lci-instance-{uuid.uuid4().hex[:10]}"
6161

6262

@@ -91,7 +91,7 @@ def database_ddl():
9191

9292
@pytest.fixture(scope="module")
9393
def default_leader():
94-
""" Default leader for multi-region instances. """
94+
"""Default leader for multi-region instances."""
9595
return "us-east4"
9696

9797

@@ -582,7 +582,7 @@ def test_update_data_with_json(capsys, instance_id, sample_database):
582582
def test_query_data_with_json_parameter(capsys, instance_id, sample_database):
583583
snippets.query_data_with_json_parameter(instance_id, sample_database.database_id)
584584
out, _ = capsys.readouterr()
585-
assert "VenueId: 19, VenueDetails: {\"open\":true,\"rating\":9}" in out
585+
assert "VenueId: 19, VenueDetails: {'open': True, 'rating': 9}" in out
586586

587587

588588
@pytest.mark.dependency(depends=["insert_datatypes_data"])

tests/system/test_dbapi.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -364,7 +364,7 @@ def test_autocommit_with_json_data(shared_instance, dbapi_database):
364364
# Assert the response
365365
assert len(got_rows) == 1
366366
assert got_rows[0][0] == 123
367-
assert got_rows[0][1] == '{"age":"26","name":"Jakob"}'
367+
assert got_rows[0][1] == {"age": "26", "name": "Jakob"}
368368

369369
# Drop the table
370370
cur.execute("DROP TABLE JsonDetails")

tests/system/test_session_api.py

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,6 @@
1919
import struct
2020
import threading
2121
import time
22-
import json
2322
import pytest
2423

2524
import grpc
@@ -28,6 +27,7 @@
2827
from google.api_core import exceptions
2928
from google.cloud import spanner_v1
3029
from google.cloud._helpers import UTC
30+
from google.cloud.spanner_v1.data_types import JsonObject
3131
from tests import _helpers as ot_helpers
3232
from . import _helpers
3333
from . import _sample_data
@@ -43,23 +43,17 @@
4343
BYTES_2 = b"Ym9vdHM="
4444
NUMERIC_1 = decimal.Decimal("0.123456789")
4545
NUMERIC_2 = decimal.Decimal("1234567890")
46-
JSON_1 = json.dumps(
46+
JSON_1 = JsonObject(
4747
{
4848
"sample_boolean": True,
4949
"sample_int": 872163,
5050
"sample float": 7871.298,
5151
"sample_null": None,
5252
"sample_string": "abcdef",
5353
"sample_array": [23, 76, 19],
54-
},
55-
sort_keys=True,
56-
separators=(",", ":"),
57-
)
58-
JSON_2 = json.dumps(
59-
{"sample_object": {"name": "Anamika", "id": 2635}},
60-
sort_keys=True,
61-
separators=(",", ":"),
54+
}
6255
)
56+
JSON_2 = JsonObject({"sample_object": {"name": "Anamika", "id": 2635}},)
6357

6458
COUNTERS_TABLE = "counters"
6559
COUNTERS_COLUMNS = ("name", "value")

tests/unit/test__helpers.py

Lines changed: 12 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -567,14 +567,22 @@ def test_w_json(self):
567567
from google.cloud.spanner_v1 import Type
568568
from google.cloud.spanner_v1 import TypeCode
569569

570-
VALUE = json.dumps(
571-
{"id": 27863, "Name": "Anamika"}, sort_keys=True, separators=(",", ":")
572-
)
570+
VALUE = {"id": 27863, "Name": "Anamika"}
571+
str_repr = json.dumps(VALUE, sort_keys=True, separators=(",", ":"))
572+
573573
field_type = Type(code=TypeCode.JSON)
574-
value_pb = Value(string_value=VALUE)
574+
value_pb = Value(string_value=str_repr)
575575

576576
self.assertEqual(self._callFUT(value_pb, field_type), VALUE)
577577

578+
VALUE = None
579+
str_repr = json.dumps(VALUE, sort_keys=True, separators=(",", ":"))
580+
581+
field_type = Type(code=TypeCode.JSON)
582+
value_pb = Value(string_value=str_repr)
583+
584+
self.assertEqual(self._callFUT(value_pb, field_type), {})
585+
578586
def test_w_unknown_type(self):
579587
from google.protobuf.struct_pb2 import Value
580588
from google.cloud.spanner_v1 import Type

0 commit comments

Comments
 (0)