diff options
| -rw-r--r-- | doc/build/changelog/unreleased_14/7284.rst | 13 | ||||
| -rw-r--r-- | lib/sqlalchemy/dialects/postgresql/asyncpg.py | 84 | ||||
| -rw-r--r-- | test/dialect/postgresql/test_async_pg_py3k.py | 20 |
3 files changed, 88 insertions, 29 deletions
diff --git a/doc/build/changelog/unreleased_14/7284.rst b/doc/build/changelog/unreleased_14/7284.rst new file mode 100644 index 000000000..b5d23739c --- /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 2225a7278..fedc0b495 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 62c8f5dde..12917e976 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)]) |