From: Mike Bayer Date: Wed, 3 Nov 2021 15:32:51 +0000 (-0400) Subject: simplify and publicize the asyncpg JSON(B) codec registrsation X-Git-Tag: rel_2_0_0b1~669^2 X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=8bd8f6c5aa1e85907b1517a57a91997532f3ebd7;p=thirdparty%2Fsqlalchemy%2Fsqlalchemy.git simplify and publicize the asyncpg JSON(B) codec registrsation Added overridable methods ``PGDialect_asyncpg.setup_asyncpg_json_codec`` and ``PGDialect_asyncpg.setup_asyncpg_jsonb_codec`` codec, which handle the required task of registering JSON/JSONB codecs for these datatypes when using asyncpg. The change is that methods are broken out as individual, overridable methods to support third party dialects that need to alter or disable how these particular codecs are set up. Fixes: #7284 Change-Id: I3eac258fea61f3975bd03c428747f788813ce45e --- diff --git a/doc/build/changelog/unreleased_14/7284.rst b/doc/build/changelog/unreleased_14/7284.rst new file mode 100644 index 0000000000..b5d23739c8 --- /dev/null +++ b/doc/build/changelog/unreleased_14/7284.rst @@ -0,0 +1,13 @@ +.. change:: + :tags: postgresql, usecase, asyncpg + :tickets: 7284 + :versions: 2.0.0b1 + + Added overridable methods ``PGDialect_asyncpg.setup_asyncpg_json_codec`` + and ``PGDialect_asyncpg.setup_asyncpg_jsonb_codec`` codec, which handle the + required task of registering JSON/JSONB codecs for these datatypes when + using asyncpg. The change is that methods are broken out as individual, + overridable methods to support third party dialects that need to alter or + disable how these particular codecs are set up. + + diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py index 2225a72786..fedc0b495b 100644 --- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py +++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py @@ -1003,8 +1003,42 @@ class PGDialect_asyncpg(PGDialect): } ) - def on_connect(self): - super_connect = super(PGDialect_asyncpg, self).on_connect() + async def setup_asyncpg_json_codec(self, conn): + """set up JSON codec for asyncpg. + + This occurs for all new connections and + can be overridden by third party dialects. + + .. versionadded:: 1.4.27 + + """ + + asyncpg_connection = conn._connection + deserializer = self._json_deserializer or _py_json.loads + + def _json_decoder(bin_value): + return deserializer(bin_value.decode()) + + await asyncpg_connection.set_type_codec( + "json", + encoder=str.encode, + decoder=_json_decoder, + schema="pg_catalog", + format="binary", + ) + + async def setup_asyncpg_jsonb_codec(self, conn): + """set up JSONB codec for asyncpg. + + This occurs for all new connections and + can be overridden by third party dialects. + + .. versionadded:: 1.4.27 + + """ + + asyncpg_connection = conn._connection + deserializer = self._json_deserializer or _py_json.loads def _jsonb_encoder(str_value): # \x01 is the prefix for jsonb used by PostgreSQL. @@ -1013,43 +1047,35 @@ class PGDialect_asyncpg(PGDialect): deserializer = self._json_deserializer or _py_json.loads - def _json_decoder(bin_value): - return deserializer(bin_value.decode()) - def _jsonb_decoder(bin_value): # the byte is the \x01 prefix for jsonb used by PostgreSQL. # asyncpg returns it when format='binary' return deserializer(bin_value[1:].decode()) - async def _setup_type_codecs(conn): - """set up type decoders at the asyncpg level. + await asyncpg_connection.set_type_codec( + "jsonb", + encoder=_jsonb_encoder, + decoder=_jsonb_decoder, + schema="pg_catalog", + format="binary", + ) - these are set_type_codec() calls to normalize - There was a tentative decoder for the "char" datatype here - to have it return strings however this type is actually a binary - type that other drivers are likely mis-interpreting. + def on_connect(self): + """on_connect for asyncpg - See https://github.com/MagicStack/asyncpg/issues/623 for reference - on why it's set up this way. + A major component of this for asyncpg is to set up type decoders at the + asyncpg level. - """ - await conn._connection.set_type_codec( - "json", - encoder=str.encode, - decoder=_json_decoder, - schema="pg_catalog", - format="binary", - ) - await conn._connection.set_type_codec( - "jsonb", - encoder=_jsonb_encoder, - decoder=_jsonb_decoder, - schema="pg_catalog", - format="binary", - ) + See https://github.com/MagicStack/asyncpg/issues/623 for + notes on JSON/JSONB implementation. + + """ + + super_connect = super(PGDialect_asyncpg, self).on_connect() def connect(conn): - conn.await_(_setup_type_codecs(conn)) + conn.await_(self.setup_asyncpg_json_codec(conn)) + conn.await_(self.setup_asyncpg_jsonb_codec(conn)) if super_connect is not None: super_connect(conn) diff --git a/test/dialect/postgresql/test_async_pg_py3k.py b/test/dialect/postgresql/test_async_pg_py3k.py index 62c8f5dde9..12917e9766 100644 --- a/test/dialect/postgresql/test_async_pg_py3k.py +++ b/test/dialect/postgresql/test_async_pg_py3k.py @@ -13,6 +13,7 @@ from sqlalchemy.dialects.postgresql import ENUM from sqlalchemy.testing import async_test from sqlalchemy.testing import eq_ from sqlalchemy.testing import fixtures +from sqlalchemy.testing import mock class AsyncPgTest(fixtures.TestBase): @@ -251,3 +252,22 @@ class AsyncPgTest(fixtures.TestBase): await conn.begin() await conn.rollback() + + @testing.combinations( + "setup_asyncpg_json_codec", + "setup_asyncpg_jsonb_codec", + argnames="methname", + ) + @async_test + async def test_codec_registration( + self, metadata, async_testing_engine, methname + ): + """test new hooks added for #7284""" + + engine = async_testing_engine() + with mock.patch.object(engine.dialect, methname) as codec_meth: + conn = await engine.connect() + adapted_conn = (await conn.get_raw_connection()).connection + await conn.close() + + eq_(codec_meth.mock_calls, [mock.call(adapted_conn)])