Fix handling of JSON data

Introduced in 5.7.8

Fixes julien-duponchelle#181

Signed-off-by: Arthur Gautier <[email protected]>

Author: Arthur Gautier

pymysqlreplication/column.py

@@ -47,6 +47,8 @@ def __parse_column_definition(self, column_type, column_schema, packet):
             self.length_size = packet.read_uint8()
         elif self.type == FIELD_TYPE.GEOMETRY:
             self.length_size = packet.read_uint8()
+        elif self.type == FIELD_TYPE.JSON:
+            self.length_size = packet.read_uint8()
         elif self.type == FIELD_TYPE.NEWDECIMAL:
             self.precision = packet.read_uint8()
             self.decimals = packet.read_uint8()

pymysqlreplication/constants/FIELD_TYPE.py

@@ -41,6 +41,7 @@
 TIMESTAMP2 = 17
 DATETIME2 = 18
 TIME2 = 19
+JSON = 245 # Introduced in 5.7.8
 NEWDECIMAL = 246
 ENUM = 247
 SET = 248

pymysqlreplication/packet.py

@@ -18,6 +18,23 @@
 UNSIGNED_INT64_LENGTH = 8
 
 
+JSONB_TYPE_SMALL_OBJECT = 0x0
+JSONB_TYPE_LARGE_OBJECT = 0x1
+JSONB_TYPE_SMALL_ARRAY = 0x2
+JSONB_TYPE_LARGE_ARRAY = 0x3
+JSONB_TYPE_LITERAL = 0x4
+JSONB_TYPE_INT16 = 0x5
+JSONB_TYPE_UINT16 = 0x6
+JSONB_TYPE_INT32 = 0x7
+JSONB_TYPE_UINT32 = 0x8
+JSONB_TYPE_INT64 = 0x9
+JSONB_TYPE_UINT64 = 0xA
+JSONB_TYPE_DOUBLE = 0xB
+JSONB_TYPE_STRING = 0xC
+JSONB_TYPE_OPAQUE = 0xF
+
+
+
 class BinLogPacketWrapper(object):
     """
     Bin Log Packet Wrapper. It uses an existing packet object, and wraps
@@ -281,3 +298,63 @@ def unpack_int32(self, n):
                 + (struct.unpack('B', n[3])[0] << 24)
         except TypeError:
             return n[0] + (n[1] << 8) + (n[2] << 16) + (n[3] << 24)
+
+    def read_binary_json(self, size):
+        length = self.read_uint_by_size(size)
+        t = self.read_uint8()
+
+        return self.read_binary_json_type(t, length)
+
+    def read_binary_json_type(self, t, length):
+        if t in (JSONB_TYPE_SMALL_OBJECT, JSONB_TYPE_LARGE_OBJECT):
+            return self.read_binary_json_object(length - 1,
+                large=(t == JSONB_TYPE_LARGE_OBJECT))
+        elif t in (JSONB_TYPE_SMALL_ARRAY, JSONB_TYPE_LARGE_ARRAY):
+            return self.read_binary_json_array(length - 1,
+                large=(t == JSONB_TYPE_LARGE_ARRAY))
+        elif t in (JSONB_TYPE_STRING,):
+            return self.read_length_coded_pascal_string(1)
+
+    def read_binary_json_object(self, length, large):
+        if large:
+            elements = self.read_uint32()
+            size = self.read_uint32()
+        else:
+            elements = self.read_uint16()
+            size = self.read_uint16()
+
+        if size > length:
+            raise ValueError('Json length is larger than packet length')
+
+        if large:
+            key_offset_lengths = [(
+               self.read_uint32(),  # offset (we don't actually need that)
+               self.read_uint32()   # size of the key
+               ) for _ in range(elements)]
+
+            value_type_lengths = [(
+               self.read_uint8(),  # type
+               self.read_uint32()   # offset
+               ) for _ in range(elements)]
+        else:
+            key_offset_lengths = [(
+               self.read_uint16(),  # offset (we don't actually need that)
+               self.read_uint16()   # size of key
+               ) for _ in range(elements)]
+
+            value_type_lengths = [(
+               self.read_uint8(),  # type
+               self.read_uint16()   # offset
+               ) for _ in range(elements)]
+
+        keys = [self.read(x[1]) for x in key_offset_lengths]
+
+        out = {}
+        for i in range(elements):
+            t = value_type_lengths[i][0]
+            data = self.read_binary_json_type(t, length)
+            out[keys[i]] = data
+
+        return out
+
+

pymysqlreplication/row_event.py

@@ -3,6 +3,7 @@
 import struct
 import decimal
 import datetime
+import json
 
 from pymysql.util import byte2int
 from pymysql.charset import charset_to_encoding
@@ -167,6 +168,8 @@ def _read_column_data(self,  cols_bitmap):
             elif column.type == FIELD_TYPE.GEOMETRY:
                 values[name] = self.packet.read_length_coded_pascal_string(
                     column.length_size)
+            elif column.type == FIELD_TYPE.JSON:
+                values[name] = self.packet.read_binary_json(column.length_size)
             else:
                 raise NotImplementedError("Unknown MySQL column type: %d" %
                                           (column.type))

pymysqlreplication/tests/test_data_type.py

@@ -413,6 +413,23 @@ def test_geometry(self):
         event = self.create_and_insert_value(create_query, insert_query)
         self.assertEqual(event.rows[0]["values"]["test"], b'\x00\x00\x00\x00\x01\x01\x00\x00\x00\x00\x00\x00\x00\x00\x00\xf0?\x00\x00\x00\x00\x00\x00\xf0?')
 
+    def test_json(self):
+        if not self.isMySQL57():
+            self.skipTest("Json is only supported in mysql 5.7")
+        create_query = "CREATE TABLE test (id int, value json);"
+        insert_query = """INSERT INTO test (id, value) VALUES (1, '{"my_key": "my_val", "my_key2": "my_val2"}');"""
+        event = self.create_and_insert_value(create_query, insert_query)
+        self.assertEqual(event.rows[0]["values"]["value"], {b"my_key": b"my_val", b"my_key2": b"my_val2"})
+
+    def test_json_unicode(self):
+        if not self.isMySQL57():
+            self.skipTest("Json is only supported in mysql 5.7")
+        create_query = "CREATE TABLE test (id int, value json);"
+        insert_query = u"""INSERT INTO test (id, value) VALUES (1, '{"miam": "🍔"}');"""
+        event = self.create_and_insert_value(create_query, insert_query)
+        print(event.rows[0]["values"])
+        self.assertEqual(event.rows[0]["values"]["value"][b"miam"], u'🍔'.encode())
+
     def test_null(self):
         create_query = "CREATE TABLE test ( \
             test TINYINT NULL DEFAULT NULL, \