Gracefully degrade unsupported types with asyncpg

Fixes: #7284

Modify the on_connect() method of PGDialect_asyncpg to
gracefully degrade unsupported types instead of throwing a
ValueError. Useful for third-party dialects that derive
from PGDialect_asyncpg but whose databases do not support
all types (e.g., CockroachDB supports JSONB but not JSON).

Change-Id: Ibb7cc8c3de632d27b9716a93d83956a590b2a2b0

1 files changed, 29 insertions, 14 deletions

diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
index 3d195e691..4e968a9b6 100644
--- a/lib/sqlalchemy/dialects/postgresql/asyncpg.py
+++ b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
@@ -870,6 +870,8 @@ class PGDialect_asyncpg(PGDialect):
 
     use_native_uuid = True
 
+    _exclude_type_codecs = util.EMPTY_SET
+
     colspecs = util.update_copy(
         PGDialect.colspecs,
         {
@@ -1022,21 +1024,34 @@ class PGDialect_asyncpg(PGDialect):
 
             See https://github.com/MagicStack/asyncpg/issues/623 for reference
             on why it's set up this way.
+
+            Also, see https://github.com/sqlalchemy/sqlalchemy/issues/7284 for
+            the rationale behind adding self._exclude_type_codecs
             """
-            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",
-            )
+
+            if "json" not in self._exclude_type_codecs:
+                try:
+                    await conn._connection.set_type_codec(
+                        "json",
+                        encoder=str.encode,
+                        decoder=_json_decoder,
+                        schema="pg_catalog",
+                        format="binary",
+                    )
+                except ValueError:
+                    self._exclude_type_codecs |= {"json"}
+
+            if "jsonb" not in self._exclude_type_codecs:
+                try:
+                    await conn._connection.set_type_codec(
+                        "jsonb",
+                        encoder=_jsonb_encoder,
+                        decoder=_jsonb_decoder,
+                        schema="pg_catalog",
+                        format="binary",
+                    )
+                except ValueError:
+                    self._exclude_type_codecs |= {"jsonb"}
 
         def connect(conn):
             conn.await_(_setup_type_codecs(conn))