Added multipolygon support, related to #183 (#185)

Author: bakwc

mysql_ch_replicator/converter.py

@@ -182,6 +182,93 @@ def parse_mysql_polygon(binary):
     return points
 
 
+def parse_mysql_multipolygon(binary):
+    """
+    Parses the binary representation of a MySQL MULTIPOLYGON data type
+    and returns a list of polygons, where each polygon is a list of tuples 
+    [(x1,y1), (x2,y2), ...] representing the polygon vertices.
+
+    :param binary: The binary data representing the MULTIPOLYGON.
+    :return: A list of lists of tuples with the coordinate values.
+    """
+    if binary is None:
+        return []
+
+    # Determine if SRID is present
+    has_srid = len(binary) > 25
+    offset = 4 if has_srid else 0
+
+    # Read byte order
+    byte_order = binary[offset]
+    if byte_order == 0:
+        endian = '>'
+    elif byte_order == 1:
+        endian = '<'
+    else:
+        raise ValueError("Invalid byte order in WKB MULTIPOLYGON")
+
+    # Read WKB Type
+    wkb_type = struct.unpack(endian + 'I', binary[offset + 1:offset + 5])[0]
+    if wkb_type != 6:  # WKB type 6 means MULTIPOLYGON
+        raise ValueError("Not a WKB MULTIPOLYGON type")
+
+    # Read number of polygons
+    num_polygons = struct.unpack(endian + 'I', binary[offset + 5:offset + 9])[0]
+    if num_polygons == 0:
+        return []
+
+    polygons = []
+    current_offset = offset + 9
+
+    for polygon_idx in range(num_polygons):
+        # Each polygon starts with its own WKB header
+        # Read byte order for this polygon
+        polygon_byte_order = binary[current_offset]
+        if polygon_byte_order == 0:
+            polygon_endian = '>'
+        elif polygon_byte_order == 1:
+            polygon_endian = '<'
+        else:
+            raise ValueError("Invalid byte order in WKB POLYGON within MULTIPOLYGON")
+
+        # Read WKB Type for this polygon
+        polygon_wkb_type = struct.unpack(polygon_endian + 'I', binary[current_offset + 1:current_offset + 5])[0]
+        if polygon_wkb_type != 3:  # WKB type 3 means POLYGON
+            raise ValueError("Not a WKB POLYGON type within MULTIPOLYGON")
+
+        # Read number of rings for this polygon
+        num_rings = struct.unpack(polygon_endian + 'I', binary[current_offset + 5:current_offset + 9])[0]
+        if num_rings == 0:
+            polygons.append([])
+            current_offset += 9
+            continue
+
+        # Read the first ring (outer boundary) of this polygon
+        ring_offset = current_offset + 9
+        num_points = struct.unpack(polygon_endian + 'I', binary[ring_offset:ring_offset + 4])[0]
+        points = []
+
+        # Read each point in the ring
+        for i in range(num_points):
+            point_offset = ring_offset + 4 + (i * 16)  # 16 bytes per point (8 for x, 8 for y)
+            x = struct.unpack(polygon_endian + 'd', binary[point_offset:point_offset + 8])[0]
+            y = struct.unpack(polygon_endian + 'd', binary[point_offset + 8:point_offset + 16])[0]
+            points.append((x, y))
+
+        polygons.append(points)
+
+        # Move to next polygon
+        # Skip the current polygon's data: header (9 bytes) + ring header (4 bytes) + points (16 bytes each)
+        current_offset = ring_offset + 4 + (num_points * 16)
+
+        # Skip any additional rings (holes) for this polygon
+        for ring_idx in range(1, num_rings):
+            ring_num_points = struct.unpack(polygon_endian + 'I', binary[current_offset:current_offset + 4])[0]
+            current_offset += 4 + (ring_num_points * 16)
+
+    return polygons
+
+
 def strip_sql_name(name):
     name = name.strip()
     if name.startswith('`'):
@@ -261,6 +348,9 @@ def convert_type(self, mysql_type, parameters):
         if mysql_type == 'polygon':
             return 'Array(Tuple(x Float32, y Float32))'
 
+        if mysql_type == 'multipolygon':
+            return 'Array(Array(Tuple(x Float32, y Float32)))'
+
         # Correctly handle numeric types
         if mysql_type.startswith('numeric'):
             # Determine if parameters are specified via parentheses:
@@ -501,6 +591,9 @@ def convert_record(
             if mysql_field_type.startswith('polygon'):
                 clickhouse_field_value = parse_mysql_polygon(clickhouse_field_value)
 
+            if mysql_field_type.startswith('multipolygon'):
+                clickhouse_field_value = parse_mysql_multipolygon(clickhouse_field_value)
+
             if mysql_field_type.startswith('enum('):
                 enum_values = mysql_structure.fields[idx].additional_data
                 field_name = mysql_structure.fields[idx].name if idx < len(mysql_structure.fields) else "unknown"

test_mysql_ch_replicator.py

@@ -1734,24 +1734,25 @@ def test_polygon_type():
 
     prepare_env(cfg, mysql, ch)
 
-    # Create a table with polygon type
+    # Create a table with polygon and multipolygon types
     mysql.execute(f'''
     CREATE TABLE `{TEST_TABLE_NAME}` (
         id INT NOT NULL AUTO_INCREMENT,
         name VARCHAR(50) NOT NULL,
         area POLYGON NOT NULL,
         nullable_area POLYGON,
+        multi_area MULTIPOLYGON,
         PRIMARY KEY (id)
     )
     ''')
 
     # Insert test data with polygons
     # Using ST_GeomFromText to create polygons from WKT (Well-Known Text) format
     mysql.execute(f'''
-    INSERT INTO `{TEST_TABLE_NAME}` (name, area, nullable_area) VALUES 
-    ('Square', ST_GeomFromText('POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))'), ST_GeomFromText('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))')),
-    ('Triangle', ST_GeomFromText('POLYGON((0 0, 1 0, 0.5 1, 0 0))'), NULL),
-    ('Complex', ST_GeomFromText('POLYGON((0 0, 0 3, 3 3, 3 0, 0 0))'), ST_GeomFromText('POLYGON((1 1, 1 2, 2 2, 2 1, 1 1))'));
+    INSERT INTO `{TEST_TABLE_NAME}` (name, area, nullable_area, multi_area) VALUES 
+    ('Square', ST_GeomFromText('POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))'), ST_GeomFromText('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))'), NULL),
+    ('Triangle', ST_GeomFromText('POLYGON((0 0, 1 0, 0.5 1, 0 0))'), NULL, NULL),
+    ('Complex', ST_GeomFromText('POLYGON((0 0, 0 3, 3 3, 3 0, 0 0))'), ST_GeomFromText('POLYGON((1 1, 1 2, 2 2, 2 1, 1 1))'), NULL);
     ''', commit=True)
 
     run_all_runner = RunAllRunner(cfg_file=config_file)
@@ -1772,6 +1773,7 @@ def test_polygon_type():
     assert results[0]['name'] == 'Square'
     assert len(results[0]['area']) == 5  # Square has 5 points (including closing point)
     assert len(results[0]['nullable_area']) == 5
+    assert results[0]['multi_area'] == []  # NULL multipolygon values are returned as empty list
     # Verify some specific points
     assert results[0]['area'][0] == {'x': 0.0, 'y': 0.0}
     assert results[0]['area'][1] == {'x': 0.0, 'y': 1.0}
@@ -1783,6 +1785,7 @@ def test_polygon_type():
     assert results[1]['name'] == 'Triangle'
     assert len(results[1]['area']) == 4  # Triangle has 4 points (including closing point)
     assert results[1]['nullable_area'] == []  # NULL values are returned as empty list
+    assert results[1]['multi_area'] == []  # NULL multipolygon values are returned as empty list
     # Verify some specific points
     assert results[1]['area'][0] == {'x': 0.0, 'y': 0.0}
     assert results[1]['area'][1] == {'x': 1.0, 'y': 0.0}
@@ -1793,6 +1796,7 @@ def test_polygon_type():
     assert results[2]['name'] == 'Complex'
     assert len(results[2]['area']) == 5  # Outer square
     assert len(results[2]['nullable_area']) == 5  # Inner square
+    assert results[2]['multi_area'] == []  # NULL multipolygon values are returned as empty list
     # Verify some specific points
     assert results[2]['area'][0] == {'x': 0.0, 'y': 0.0}
     assert results[2]['area'][2] == {'x': 3.0, 'y': 3.0}
@@ -1801,10 +1805,10 @@ def test_polygon_type():
 
     # Test realtime replication by adding more records
     mysql.execute(f'''
-    INSERT INTO `{TEST_TABLE_NAME}` (name, area, nullable_area) VALUES 
-    ('Pentagon', ST_GeomFromText('POLYGON((0 0, 1 0, 1.5 1, 0.5 1.5, 0 0))'), ST_GeomFromText('POLYGON((0.2 0.2, 0.8 0.2, 1 0.8, 0.5 1, 0.2 0.2))')),
-    ('Hexagon', ST_GeomFromText('POLYGON((0 0, 1 0, 1.5 0.5, 1 1, 0.5 1, 0 0))'), NULL),
-    ('Circle', ST_GeomFromText('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))'), ST_GeomFromText('POLYGON((0.5 0.5, 0.5 1.5, 1.5 1.5, 1.5 0.5, 0.5 0.5))'));
+    INSERT INTO `{TEST_TABLE_NAME}` (name, area, nullable_area, multi_area) VALUES 
+    ('Pentagon', ST_GeomFromText('POLYGON((0 0, 1 0, 1.5 1, 0.5 1.5, 0 0))'), ST_GeomFromText('POLYGON((0.2 0.2, 0.8 0.2, 1 0.8, 0.5 1, 0.2 0.2))'), NULL),
+    ('Hexagon', ST_GeomFromText('POLYGON((0 0, 1 0, 1.5 0.5, 1 1, 0.5 1, 0 0))'), NULL, NULL),
+    ('Circle', ST_GeomFromText('POLYGON((0 0, 0 2, 2 2, 2 0, 0 0))'), ST_GeomFromText('POLYGON((0.5 0.5, 0.5 1.5, 1.5 1.5, 1.5 0.5, 0.5 0.5))'), NULL);
     ''', commit=True)
 
     # Wait for new records to be replicated
@@ -1816,6 +1820,7 @@ def test_polygon_type():
     assert pentagon['name'] == 'Pentagon'
     assert len(pentagon['area']) == 5  # Pentagon has 5 points
     assert len(pentagon['nullable_area']) == 5  # Inner pentagon
+    assert pentagon['multi_area'] == []  # NULL multipolygon values are returned as empty list
     assert abs(pentagon['area'][0]['x'] - 0.0) < 1e-6
     assert abs(pentagon['area'][0]['y'] - 0.0) < 1e-6
     assert abs(pentagon['area'][2]['x'] - 1.5) < 1e-6
@@ -1830,6 +1835,7 @@ def test_polygon_type():
     assert hexagon['name'] == 'Hexagon'
     assert len(hexagon['area']) == 6  # Hexagon has 6 points
     assert hexagon['nullable_area'] == []  # NULL values are returned as empty list
+    assert hexagon['multi_area'] == []  # NULL multipolygon values are returned as empty list
     assert abs(hexagon['area'][0]['x'] - 0.0) < 1e-6
     assert abs(hexagon['area'][0]['y'] - 0.0) < 1e-6
     assert abs(hexagon['area'][2]['x'] - 1.5) < 1e-6
@@ -1842,6 +1848,7 @@ def test_polygon_type():
     assert circle['name'] == 'Circle'
     assert len(circle['area']) == 5  # Outer square
     assert len(circle['nullable_area']) == 5  # Inner square
+    assert circle['multi_area'] == []  # NULL multipolygon values are returned as empty list
     assert abs(circle['area'][0]['x'] - 0.0) < 1e-6
     assert abs(circle['area'][0]['y'] - 0.0) < 1e-6
     assert abs(circle['area'][2]['x'] - 2.0) < 1e-6
@@ -1851,6 +1858,46 @@ def test_polygon_type():
     assert abs(circle['nullable_area'][2]['x'] - 1.5) < 1e-6
     assert abs(circle['nullable_area'][2]['y'] - 1.5) < 1e-6
 
+    # Test multipolygon type - insert a record with multipolygon data
+    mysql.execute(f'''
+    INSERT INTO `{TEST_TABLE_NAME}` (name, area, nullable_area, multi_area) VALUES 
+    ('MultiSquares', 
+     ST_GeomFromText('POLYGON((0 0, 0 1, 1 1, 1 0, 0 0))'), 
+     NULL,
+     ST_GeomFromText('MULTIPOLYGON(((0 0, 0 1, 1 1, 1 0, 0 0)), ((2 2, 2 3, 3 3, 3 2, 2 2)))')
+    );
+    ''', commit=True)
+
+    # Wait for the new record with multipolygon to be replicated
+    assert_wait(lambda: len(ch.select(TEST_TABLE_NAME)) == 7)
+    
+    # Verify the multipolygon data
+    multi_squares = ch.select(TEST_TABLE_NAME, where="name='MultiSquares'")[0]
+    assert multi_squares['name'] == 'MultiSquares'
+    
+    # Check that multi_area contains multiple polygons
+    # The multipolygon should be represented as an array of polygon arrays
+    assert isinstance(multi_squares['multi_area'], list)
+    assert len(multi_squares['multi_area']) == 2  # Two polygons in the multipolygon
+    
+    # Check first polygon in multipolygon
+    first_polygon = multi_squares['multi_area'][0]
+    assert len(first_polygon) == 5  # Square has 5 points (including closing point)
+    assert first_polygon[0] == {'x': 0.0, 'y': 0.0}
+    assert first_polygon[1] == {'x': 0.0, 'y': 1.0}
+    assert first_polygon[2] == {'x': 1.0, 'y': 1.0}
+    assert first_polygon[3] == {'x': 1.0, 'y': 0.0}
+    assert first_polygon[4] == {'x': 0.0, 'y': 0.0}  # Closing point
+    
+    # Check second polygon in multipolygon
+    second_polygon = multi_squares['multi_area'][1]
+    assert len(second_polygon) == 5  # Square has 5 points (including closing point)
+    assert second_polygon[0] == {'x': 2.0, 'y': 2.0}
+    assert second_polygon[1] == {'x': 2.0, 'y': 3.0}
+    assert second_polygon[2] == {'x': 3.0, 'y': 3.0}
+    assert second_polygon[3] == {'x': 3.0, 'y': 2.0}
+    assert second_polygon[4] == {'x': 2.0, 'y': 2.0}  # Closing point
+
     run_all_runner.stop()
     assert_wait(lambda: 'stopping db_replicator' in read_logs(TEST_DB_NAME))
     assert('Traceback' not in read_logs(TEST_DB_NAME))