Skip to content

Commit 9c32b86

Browse files
committed
Make sure set_type_codec() overrides core codec for all formats.
When a custom codec is set for a type, it should be called regardless of whether its format matches the preferred format requested by the pipeline. Fixes: #140 Fixes: #148
1 parent 6ca1f28 commit 9c32b86

File tree

3 files changed

+72
-10
lines changed

3 files changed

+72
-10
lines changed

asyncpg/protocol/codecs/base.pxd

+2
Original file line numberDiff line numberDiff line change
@@ -161,3 +161,5 @@ cdef class DataCodecConfig:
161161
dict _local_type_codecs
162162

163163
cdef inline Codec get_codec(self, uint32_t oid, CodecFormat format)
164+
cdef inline Codec get_local_codec(
165+
self, uint32_t oid, CodecFormat preferred_format=*)

asyncpg/protocol/codecs/base.pyx

+31-9
Original file line numberDiff line numberDiff line change
@@ -563,17 +563,39 @@ cdef class DataCodecConfig:
563563
cdef inline Codec get_codec(self, uint32_t oid, CodecFormat format):
564564
cdef Codec codec
565565

566-
try:
567-
return self._local_type_codecs[oid, format]
568-
except KeyError:
569-
codec = get_core_codec(oid, format)
570-
if codec is not None:
566+
codec = self.get_local_codec(oid, format)
567+
if codec is not None:
568+
if codec.format != format:
569+
# The codec for this OID has been overridden by
570+
# set_{builtin}_type_codec with a different format.
571+
# We must respect that and not return a core codec.
572+
return None
573+
else:
571574
return codec
575+
576+
codec = get_core_codec(oid, format)
577+
if codec is not None:
578+
return codec
579+
else:
580+
try:
581+
return self._type_codecs_cache[oid, format]
582+
except KeyError:
583+
return None
584+
585+
cdef inline Codec get_local_codec(
586+
self, uint32_t oid, CodecFormat preferred_format=PG_FORMAT_BINARY):
587+
cdef Codec codec
588+
589+
codec = self._local_type_codecs.get((oid, preferred_format))
590+
if codec is None:
591+
if preferred_format == PG_FORMAT_BINARY:
592+
alt_format = PG_FORMAT_TEXT
572593
else:
573-
try:
574-
return self._type_codecs_cache[oid, format]
575-
except KeyError:
576-
return None
594+
alt_format = PG_FORMAT_BINARY
595+
596+
codec = self._local_type_codecs.get((oid, alt_format))
597+
598+
return codec
577599

578600

579601
cdef inline Codec get_core_codec(uint32_t oid, CodecFormat format):

tests/test_codecs.py

+39-1
Original file line numberDiff line numberDiff line change
@@ -924,7 +924,7 @@ def hstore_encoder(obj):
924924
DROP EXTENSION hstore
925925
''')
926926

927-
async def test_custom_codec_override(self):
927+
async def test_custom_codec_override_binary(self):
928928
"""Test overriding core codecs."""
929929
import json
930930

@@ -948,6 +948,44 @@ def _decoder(value):
948948
finally:
949949
await conn.close()
950950

951+
async def test_custom_codec_override_text(self):
952+
"""Test overriding core codecs."""
953+
import json
954+
955+
conn = await self.cluster.connect(database='postgres', loop=self.loop)
956+
try:
957+
def _encoder(value):
958+
return json.dumps(value)
959+
960+
def _decoder(value):
961+
return json.loads(value)
962+
963+
await conn.set_type_codec(
964+
'json', encoder=_encoder, decoder=_decoder,
965+
schema='pg_catalog', binary=False
966+
)
967+
968+
data = {'foo': 'bar', 'spam': 1}
969+
res = await conn.fetchval('SELECT $1::json', data)
970+
self.assertEqual(data, res)
971+
972+
def _encoder(value):
973+
return value
974+
975+
def _decoder(value):
976+
return value
977+
978+
await conn.set_type_codec(
979+
'uuid', encoder=_encoder, decoder=_decoder,
980+
schema='pg_catalog', binary=False
981+
)
982+
983+
data = '14058ad9-0118-4b7e-ac15-01bc13e2ccd1'
984+
res = await conn.fetchval('SELECT $1::uuid', data)
985+
self.assertEqual(res, data)
986+
finally:
987+
await conn.close()
988+
951989
async def test_composites_in_arrays(self):
952990
await self.con.execute('''
953991
CREATE TYPE t AS (a text, b int);

0 commit comments

Comments
 (0)