Skip to content

Commit 8bee26d

Browse files
committed
Handle custom encoders in nested data types
The root of the problem was that nested data types would use the default encoders for subitems. When the encoders were customized, they would not be used for those nested items. This fix moves the encoder functions into a class so that collections, tuples, and UDTs will use the customized mapping when encoding subitems. Fixes PYTHON-100.
1 parent 9121cb1 commit 8bee26d

9 files changed

Lines changed: 233 additions & 244 deletions

File tree

cassandra/cluster.py

Lines changed: 18 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@
4444
from cassandra import (ConsistencyLevel, AuthenticationFailed,
4545
InvalidRequest, OperationTimedOut, UnsupportedOperation)
4646
from cassandra.connection import ConnectionException, ConnectionShutdown
47-
from cassandra.encoder import cql_encode_all_types, cql_encoders
47+
from cassandra.encoder import Encoder
4848
from cassandra.protocol import (QueryMessage, ResultMessage,
4949
ErrorMessage, ReadTimeoutErrorMessage,
5050
WriteTimeoutErrorMessage,
@@ -1162,25 +1162,25 @@ class Session(object):
11621162
.. versionadded:: 2.1.0
11631163
"""
11641164

1165-
encoders = None
1165+
encoder = None
11661166
"""
1167-
A map of python types to CQL encoder functions that will be used when
1168-
formatting query parameters for non-prepared statements. This mapping
1169-
is not used for prepared statements (because prepared statements
1170-
give the driver more information about what CQL types are expected, allowing
1171-
it to accept a wider range of python types).
1167+
A :class:`~cassandra.encoder.Encoder` instance that will be used when
1168+
formatting query parameters for non-prepared statements. This is not used
1169+
for prepared statements (because prepared statements give the driver more
1170+
information about what CQL types are expected, allowing it to accept a
1171+
wider range of python types).
11721172
1173-
This mapping can be be modified by users as they see fit. Functions from
1174-
:mod:`cassandra.encoder` should be used, if possible, because they take
1175-
precautions to avoid injections and properly sanitize data.
1173+
The encoder uses a mapping from python types to encoder methods (for
1174+
specific CQL types). This mapping can be be modified by users as they see
1175+
fit. Methods of :class:`~cassandra.encoder.Encoder` should be used for mapping
1176+
values if possible, because they take precautions to avoid injections and
1177+
properly sanitize data.
11761178
11771179
Example::
11781180
1179-
from cassandra.encoder import cql_encode_tuple
1180-
11811181
cluster = Cluster()
11821182
session = cluster.connect("mykeyspace")
1183-
session.encoders[tuple] = cql_encode_tuple
1183+
session.encoder.mapping[tuple] = session.encoder.cql_encode_tuple
11841184
11851185
session.execute("CREATE TABLE mytable (k int PRIMARY KEY, col tuple<int, ascii>)")
11861186
session.execute("INSERT INTO mytable (k, col) VALUES (%s, %s)", [0, (123, 'abc')])
@@ -1202,7 +1202,7 @@ def __init__(self, cluster, hosts):
12021202
self._metrics = cluster.metrics
12031203
self._protocol_version = self.cluster.protocol_version
12041204

1205-
self.encoders = cql_encoders.copy()
1205+
self.encoder = Encoder()
12061206

12071207
# create connection pools in parallel
12081208
futures = []
@@ -1328,7 +1328,7 @@ def _create_response_future(self, query, parameters, trace):
13281328
if six.PY2 and isinstance(query_string, six.text_type):
13291329
query_string = query_string.encode('utf-8')
13301330
if parameters:
1331-
query_string = bind_params(query_string, parameters, self.encoders)
1331+
query_string = bind_params(query_string, parameters, self.encoder)
13321332
message = QueryMessage(
13331333
query_string, cl, query.serial_consistency_level,
13341334
fetch_size, timestamp=timestamp)
@@ -1585,13 +1585,13 @@ def user_type_registered(self, keyspace, user_type, klass):
15851585
raise UserTypeDoesNotExist(
15861586
'User type %s does not exist in keyspace %s' % (user_type, keyspace))
15871587

1588-
def encode(val):
1588+
def encode(encoder_self, val):
15891589
return '{ %s }' % ' , '.join('%s : %s' % (
15901590
field_name,
1591-
cql_encode_all_types(getattr(val, field_name))
1591+
encoder_self.cql_encode_all_types(getattr(val, field_name))
15921592
) for field_name in type_meta.field_names)
15931593

1594-
self.encoders[klass] = encode
1594+
self.encoder.mapping[klass] = encode
15951595

15961596
def submit(self, fn, *args, **kwargs):
15971597
""" Internal """

cassandra/encoder.py

Lines changed: 149 additions & 147 deletions
Original file line numberDiff line numberDiff line change
@@ -49,154 +49,156 @@ def cql_quote(term):
4949
return str(term)
5050

5151

52-
def cql_encode_none(val):
53-
"""
54-
Converts :const:`None` to the string 'NULL'.
55-
"""
56-
return 'NULL'
57-
58-
59-
def cql_encode_unicode(val):
60-
"""
61-
Converts :class:`unicode` objects to UTF-8 encoded strings with quote escaping.
62-
"""
63-
return cql_quote(val.encode('utf-8'))
64-
65-
66-
def cql_encode_str(val):
67-
"""
68-
Escapes quotes in :class:`str` objects.
69-
"""
70-
return cql_quote(val)
71-
72-
73-
if six.PY3:
74-
def cql_encode_bytes(val):
75-
return (b'0x' + hexlify(val)).decode('utf-8')
76-
elif sys.version_info >= (2, 7):
77-
def cql_encode_bytes(val): # noqa
78-
return b'0x' + hexlify(val)
79-
else:
80-
# python 2.6 requires string or read-only buffer for hexlify
81-
def cql_encode_bytes(val): # noqa
82-
return b'0x' + hexlify(buffer(val))
83-
84-
85-
def cql_encode_object(val):
86-
"""
87-
Default encoder for all objects that do not have a specific encoder function
88-
registered. This function simply calls :meth:`str()` on the object.
89-
"""
90-
return str(val)
91-
92-
93-
def cql_encode_datetime(val):
94-
"""
95-
Converts a :class:`datetime.datetime` object to a (string) integer timestamp
96-
with millisecond precision.
97-
"""
98-
timestamp = calendar.timegm(val.utctimetuple())
99-
return str(long(timestamp * 1e3 + getattr(val, 'microsecond', 0) / 1e3))
100-
101-
102-
def cql_encode_date(val):
103-
"""
104-
Converts a :class:`datetime.date` object to a string with format
105-
``YYYY-MM-DD-0000``.
106-
"""
107-
return "'%s'" % val.strftime('%Y-%m-%d-0000')
52+
class ValueSequence(list):
53+
pass
10854

10955

110-
def cql_encode_sequence(val):
56+
class Encoder(object):
57+
"""
58+
A container for mapping python types to CQL string literals when working
59+
with non-prepared statements. The type :attr:`~.Encoder.mapping` can be
60+
directly customized by users.
61+
"""
62+
63+
mapping = None
64+
"""
65+
A map of python types to encoder functions.
66+
"""
67+
68+
def __init__(self):
69+
self.mapping = {
70+
float: self.cql_encode_object,
71+
bytearray: self.cql_encode_bytes,
72+
str: self.cql_encode_str,
73+
int: self.cql_encode_object,
74+
UUID: self.cql_encode_object,
75+
datetime.datetime: self.cql_encode_datetime,
76+
datetime.date: self.cql_encode_date,
77+
dict: self.cql_encode_map_collection,
78+
OrderedDict: self.cql_encode_map_collection,
79+
list: self.cql_encode_list_collection,
80+
tuple: self.cql_encode_list_collection,
81+
set: self.cql_encode_set_collection,
82+
frozenset: self.cql_encode_set_collection,
83+
types.GeneratorType: self.cql_encode_list_collection,
84+
ValueSequence: self.cql_encode_sequence
85+
}
86+
87+
if six.PY2:
88+
self.mapping.update({
89+
unicode: self.cql_encode_unicode,
90+
buffer: self.cql_encode_bytes,
91+
long: self.cql_encode_object,
92+
types.NoneType: self.cql_encode_none,
93+
})
94+
else:
95+
self.mapping.update({
96+
memoryview: self.cql_encode_bytes,
97+
bytes: self.cql_encode_bytes,
98+
type(None): self.cql_encode_none,
99+
})
100+
101+
# sortedset is optional
102+
try:
103+
from blist import sortedset
104+
self.mapping.update({
105+
sortedset: self.cql_encode_set_collection
106+
})
107+
except ImportError:
108+
pass
109+
110+
def cql_encode_none(self, val):
111+
"""
112+
Converts :const:`None` to the string 'NULL'.
113+
"""
114+
return 'NULL'
115+
116+
def cql_encode_unicode(self, val):
117+
"""
118+
Converts :class:`unicode` objects to UTF-8 encoded strings with quote escaping.
119+
"""
120+
return cql_quote(val.encode('utf-8'))
121+
122+
def cql_encode_str(self, val):
123+
"""
124+
Escapes quotes in :class:`str` objects.
125+
"""
126+
return cql_quote(val)
127+
128+
if six.PY3:
129+
def cql_encode_bytes(self, val):
130+
return (b'0x' + hexlify(val)).decode('utf-8')
131+
elif sys.version_info >= (2, 7):
132+
def cql_encode_bytes(self, val): # noqa
133+
return b'0x' + hexlify(val)
134+
else:
135+
# python 2.6 requires string or read-only buffer for hexlify
136+
def cql_encode_bytes(self, val): # noqa
137+
return b'0x' + hexlify(buffer(val))
138+
139+
def cql_encode_object(self, val):
140+
"""
141+
Default encoder for all objects that do not have a specific encoder function
142+
registered. This function simply calls :meth:`str()` on the object.
143+
"""
144+
return str(val)
145+
146+
def cql_encode_datetime(self, val):
147+
"""
148+
Converts a :class:`datetime.datetime` object to a (string) integer timestamp
149+
with millisecond precision.
150+
"""
151+
timestamp = calendar.timegm(val.utctimetuple())
152+
return str(long(timestamp * 1e3 + getattr(val, 'microsecond', 0) / 1e3))
153+
154+
def cql_encode_date(self, val):
155+
"""
156+
Converts a :class:`datetime.date` object to a string with format
157+
``YYYY-MM-DD-0000``.
158+
"""
159+
return "'%s'" % val.strftime('%Y-%m-%d-0000')
160+
161+
def cql_encode_sequence(self, val):
162+
"""
163+
Converts a sequence to a string of the form ``(item1, item2, ...)``. This
164+
is suitable for ``IN`` value lists.
165+
"""
166+
return '( %s )' % ' , '.join(self.mapping.get(type(v), self.cql_encode_object)(v)
167+
for v in val)
168+
169+
cql_encode_tuple = cql_encode_sequence
111170
"""
112171
Converts a sequence to a string of the form ``(item1, item2, ...)``. This
113-
is suitable for ``IN`` value lists.
114-
"""
115-
return '( %s )' % ' , '.join(cql_encoders.get(type(v), cql_encode_object)(v)
116-
for v in val)
117-
118-
119-
cql_encode_tuple = cql_encode_sequence
120-
"""
121-
Converts a sequence to a string of the form ``(item1, item2, ...)``. This
122-
is suitable for ``tuple`` type columns.
123-
"""
124-
125-
126-
def cql_encode_map_collection(val):
127-
"""
128-
Converts a dict into a string of the form ``{key1: val1, key2: val2, ...}``.
129-
This is suitable for ``map`` type columns.
130-
"""
131-
return '{ %s }' % ' , '.join('%s : %s' % (
132-
cql_encode_all_types(k),
133-
cql_encode_all_types(v)
134-
) for k, v in six.iteritems(val))
135-
136-
137-
def cql_encode_list_collection(val):
138-
"""
139-
Converts a sequence to a string of the form ``[item1, item2, ...]``. This
140-
is suitable for ``list`` type columns.
141-
"""
142-
return '[ %s ]' % ' , '.join(map(cql_encode_all_types, val))
143-
144-
145-
def cql_encode_set_collection(val):
146-
"""
147-
Converts a sequence to a string of the form ``{item1, item2, ...}``. This
148-
is suitable for ``set`` type columns.
149-
"""
150-
return '{ %s }' % ' , '.join(map(cql_encode_all_types, val))
151-
152-
153-
def cql_encode_all_types(val):
154-
"""
155-
Converts any type into a CQL string, defaulting to ``cql_encode_object``
156-
if :attr:`~.cql_encoders` does not contain an entry for the type.
157-
"""
158-
return cql_encoders.get(type(val), cql_encode_object)(val)
159-
160-
161-
cql_encoders = {
162-
float: cql_encode_object,
163-
bytearray: cql_encode_bytes,
164-
str: cql_encode_str,
165-
int: cql_encode_object,
166-
UUID: cql_encode_object,
167-
datetime.datetime: cql_encode_datetime,
168-
datetime.date: cql_encode_date,
169-
dict: cql_encode_map_collection,
170-
OrderedDict: cql_encode_map_collection,
171-
list: cql_encode_list_collection,
172-
tuple: cql_encode_list_collection,
173-
set: cql_encode_set_collection,
174-
frozenset: cql_encode_set_collection,
175-
types.GeneratorType: cql_encode_list_collection
176-
}
177-
"""
178-
A map of python types to encoder functions.
179-
"""
180-
181-
if six.PY2:
182-
cql_encoders.update({
183-
unicode: cql_encode_unicode,
184-
buffer: cql_encode_bytes,
185-
long: cql_encode_object,
186-
types.NoneType: cql_encode_none,
187-
})
188-
else:
189-
cql_encoders.update({
190-
memoryview: cql_encode_bytes,
191-
bytes: cql_encode_bytes,
192-
type(None): cql_encode_none,
193-
})
194-
195-
# sortedset is optional
196-
try:
197-
from blist import sortedset
198-
cql_encoders.update({
199-
sortedset: cql_encode_set_collection
200-
})
201-
except ImportError:
202-
pass
172+
is suitable for ``tuple`` type columns.
173+
"""
174+
175+
def cql_encode_map_collection(self, val):
176+
"""
177+
Converts a dict into a string of the form ``{key1: val1, key2: val2, ...}``.
178+
This is suitable for ``map`` type columns.
179+
"""
180+
return '{ %s }' % ' , '.join('%s : %s' % (
181+
self.mapping.get(type(k), self.cql_encode_object)(k),
182+
self.mapping.get(type(v), self.cql_encode_object)(v)
183+
) for k, v in six.iteritems(val))
184+
185+
def cql_encode_list_collection(self, val):
186+
"""
187+
Converts a sequence to a string of the form ``[item1, item2, ...]``. This
188+
is suitable for ``list`` type columns.
189+
"""
190+
return '[ %s ]' % ' , '.join(self.mapping.get(type(v), self.cql_encode_object)(v) for v in val)
191+
192+
def cql_encode_set_collection(self, val):
193+
"""
194+
Converts a sequence to a string of the form ``{item1, item2, ...}``. This
195+
is suitable for ``set`` type columns.
196+
"""
197+
return '{ %s }' % ' , '.join(self.mapping.get(type(v), self.cql_encode_object)(v) for v in val)
198+
199+
def cql_encode_all_types(self, val):
200+
"""
201+
Converts any type into a CQL string, defaulting to ``cql_encode_object``
202+
if :attr:`~Encoder.mapping` does not contain an entry for the type.
203+
"""
204+
return self.mapping.get(type(val), self.cql_encode_object)(val)

0 commit comments

Comments
 (0)