diff --git a/asyncpg/protocol/codecs/base.pyx b/asyncpg/protocol/codecs/base.pyx index 2da74249..70aa650e 100644 --- a/asyncpg/protocol/codecs/base.pyx +++ b/asyncpg/protocol/codecs/base.pyx @@ -474,10 +474,6 @@ cdef class DataCodecConfig: encoder, decoder, binary): format = PG_FORMAT_BINARY if binary else PG_FORMAT_TEXT - if self.get_codec(typeoid, format) is not None: - raise ValueError('cannot override codec for type {}'.format( - typeoid)) - self._local_type_codecs[typeoid, format] = \ Codec.new_python_codec(typeoid, typename, typeschema, typekind, encoder, decoder, format) @@ -519,17 +515,17 @@ cdef class DataCodecConfig: cdef inline Codec get_codec(self, uint32_t oid, CodecFormat format): cdef Codec codec - codec = get_core_codec(oid, format) - if codec is not None: - return codec - try: - return self._type_codecs_cache[oid, format] + return self._local_type_codecs[oid, format] except KeyError: - try: - return self._local_type_codecs[oid, format] - except KeyError: - return None + codec = get_core_codec(oid, format) + if codec is not None: + return codec + else: + try: + return self._type_codecs_cache[oid, format] + except KeyError: + return None cdef inline Codec get_core_codec(uint32_t oid, CodecFormat format): diff --git a/tests/test_codecs.py b/tests/test_codecs.py index ec8356e5..a810bd00 100644 --- a/tests/test_codecs.py +++ b/tests/test_codecs.py @@ -900,6 +900,30 @@ def hstore_encoder(obj): DROP EXTENSION hstore ''') + async def test_custom_codec_override(self): + """Test overriding core codecs.""" + import json + + conn = await self.cluster.connect(database='postgres', loop=self.loop) + try: + def _encoder(value): + return json.dumps(value).encode('utf-8') + + def _decoder(value): + return json.loads(value.decode('utf-8')) + + await conn.set_type_codec( + 'json', encoder=_encoder, decoder=_decoder, + schema='pg_catalog', binary=True + ) + + data = {'foo': 'bar', 'spam': 1} + res = await conn.fetchval('SELECT $1::json', data) + self.assertEqual(data, res) + + finally: + await conn.close() + async def test_composites_in_arrays(self): await self.con.execute(''' CREATE TYPE t AS (a text, b int); diff --git a/tests/test_connect.py b/tests/test_connect.py index 09c987e7..5d54fac5 100644 --- a/tests/test_connect.py +++ b/tests/test_connect.py @@ -96,9 +96,21 @@ def tearDown(self): async def test_auth_bad_user(self): with self.assertRaises( asyncpg.InvalidAuthorizationSpecificationError): - await self.cluster.connect(user='__nonexistent__', - database='postgres', - loop=self.loop) + for tried in range(3): + try: + await self.cluster.connect(user='__nonexistent__', + database='postgres', + loop=self.loop) + except asyncpg.ConnectionDoesNotExistError: + if _system == 'Windows': + # On Windows the server sometimes just closes + # the connection sooner than we receive the + # actual error. + continue + else: + raise + else: + break async def test_auth_trust(self): conn = await self.cluster.connect(