Skip to content

Commit

Permalink
Data type converter: Support converting elements of collections (ARRAY)
Browse files Browse the repository at this point in the history
  • Loading branch information
amotl committed Sep 7, 2022
1 parent d483667 commit b21c38e
Show file tree
Hide file tree
Showing 2 changed files with 54 additions and 4 deletions.
16 changes: 12 additions & 4 deletions src/crate/client/converter.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,8 +27,7 @@
from copy import deepcopy
from datetime import datetime
from enum import Enum
from typing import Any, Callable, Dict, Optional, Union

from typing import Any, Callable, Dict, Optional, Union, List

InputVal = Any

Expand Down Expand Up @@ -128,8 +127,17 @@ def get_mapping_key(type_: Union[CrateDatatypeIdentifier, int]) -> int:
return type_

def convert(self, type_: int, value: Optional[Any]) -> Optional[Any]:
converter = self.get(type_)
return converter(value)
if isinstance(type_, List):
type_, inner_type = type_
assert type_ == 100, f"Type {type_} not implemented as collection type"
if value is None:
result = self.convert(inner_type, None)
else:
result = [self.convert(inner_type, item) for item in value]
else:
converter = self.get(type_)
result = converter(value)
return result


class DefaultTypeConverter(Converter):
Expand Down
42 changes: 42 additions & 0 deletions src/crate/client/test_cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,3 +92,45 @@ def test_execute_with_converter(self):
self.assertEqual(result, ['foo', '10.10.10.1', 1658167836758, "B'0110'"])

conn.close()

def test_execute_array_with_converter(self):
client = ClientMocked()
conn = connect(client=client)
converter = Cursor.get_default_converter()
cursor = conn.cursor(converter=converter)

conn.client.set_next_response({
"col_types": [4, [100, 5]],
"cols": ["name", "address"],
"rows": [["foo", ["10.10.10.1", "10.10.10.2"]]],
"rowcount": 1,
"duration": 123
})

cursor.execute("")
result = cursor.fetchone()
self.assertEqual(result, [
'foo',
[IPv4Address('10.10.10.1'), IPv4Address('10.10.10.2')],
])

def test_execute_nested_array_with_converter(self):
client = ClientMocked()
conn = connect(client=client)
converter = Cursor.get_default_converter()
cursor = conn.cursor(converter=converter)

conn.client.set_next_response({
"col_types": [4, [100, [100, 5]]],
"cols": ["name", "address_buckets"],
"rows": [["foo", [["10.10.10.1", "10.10.10.2"], ["10.10.10.3"], [], None]]],
"rowcount": 1,
"duration": 123
})

cursor.execute("")
result = cursor.fetchone()
self.assertEqual(result, [
'foo',
[[IPv4Address('10.10.10.1'), IPv4Address('10.10.10.2')], [IPv4Address('10.10.10.3')], [], None],
])

0 comments on commit b21c38e

Please sign in to comment.