support types with comma inside map (#61)

timeplus-io · Dec 2, 2024 · 70a1c63 · 70a1c63
1 parent 28ca4c9
commit 70a1c63
Show file tree

Hide file tree

Showing 3 changed files with 42 additions and 2 deletions.
diff --git a/proton_driver/__init__.py b/proton_driver/__init__.py
@@ -3,7 +3,7 @@
 from .dbapi import connect
 
 
-VERSION = (0, 2, 12)
+VERSION = (0, 2, 13)
 __version__ = '.'.join(str(x) for x in VERSION)
 
 __all__ = ['Client', 'connect']
diff --git a/proton_driver/columns/mapcolumn.py b/proton_driver/columns/mapcolumn.py
@@ -1,6 +1,7 @@
 from .base import Column
 from .intcolumn import UInt64Column
 from ..util.helpers import pairwise
+from .util import get_inner_columns
 
 
 class MapColumn(Column):
@@ -51,7 +52,7 @@ def write_items(self, items, buf):
 
 
 def create_map_column(spec, column_by_spec_getter):
-    key, value = spec[4:-1].split(',')
+    key, value = get_inner_columns('map', spec)
     key_column = column_by_spec_getter(key.strip())
     value_column = column_by_spec_getter(value.strip())
 

diff --git a/tests/columns/test_map.py b/tests/columns/test_map.py
@@ -1,4 +1,5 @@
 from tests.testcase import BaseTestCase
+from decimal import Decimal
 
 
 class MapTestCase(BaseTestCase):
@@ -100,3 +101,41 @@ def test_array(self):
             )
             inserted = self.client.execute(query)
             self.assertEqual(inserted, data)
+
+    def test_decimal(self):
+        with self.create_stream('a map(string, Decimal(9, 2))'):
+            data = [
+                ({'key1': Decimal('123.45')}, ),
+                ({'key2': Decimal('234.56')}, ),
+                ({'key3': Decimal('345.67')}, )
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{'key1':123.45}\n"
+                "{'key2':234.56}\n"
+                "{'key3':345.67}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)
+
+    def test_nested_map(self):
+        with self.create_stream('a map(string, map(string, map(string, int)))'): # noqa
+            data = [
+                ({'key1': {'key2': {'key3': 1}}}, ),
+                ({'key1': {'key2': {'key3': 2}}}, ),
+                ({'key1': {'key2': {'key3': 3}}}, ),
+            ]
+            self.client.execute('INSERT INTO test (a) VALUES', data)
+            query = 'SELECT * FROM test'
+            inserted = self.emit_cli(query)
+            self.assertEqual(
+                inserted,
+                "{'key1':{'key2':{'key3':1}}}\n"
+                "{'key1':{'key2':{'key3':2}}}\n"
+                "{'key1':{'key2':{'key3':3}}}\n"
+            )
+            inserted = self.client.execute(query)
+            self.assertEqual(inserted, data)