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

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)])