From 96c294da8a50d692b3f0b8e508dbbca5d9c22f1b Mon Sep 17 00:00:00 2001
From: Gord Thompson <gord@gordthompson.com>
Date: Tue, 2 Nov 2021 16:16:50 -0600
Subject: [PATCH] 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
---
 doc/build/changelog/unreleased_14/7284.rst    | 11 +++++
 lib/sqlalchemy/dialects/postgresql/asyncpg.py | 43 +++++++++++++------
 2 files changed, 40 insertions(+), 14 deletions(-)
 create mode 100644 doc/build/changelog/unreleased_14/7284.rst

diff --git a/doc/build/changelog/unreleased_14/7284.rst b/doc/build/changelog/unreleased_14/7284.rst
new file mode 100644
index 0000000000..fb05488f50
--- /dev/null
+++ b/doc/build/changelog/unreleased_14/7284.rst
@@ -0,0 +1,11 @@
+.. change::
+    :tags: asyncpg, usecase
+    :tickets: 7284
+    :versions: 2.0.0b1
+
+    Modified the asyncpg dialect to gracefully degrade types that are
+    not supported by PostgreSQL work-alike databases. For example,
+    CockroachDB supports JSONB but not JSON. Third-party dialects that
+    are derived from ``PGDialect_asyncpg`` will no longer have to
+    work around ValueError being raised by trying to register a codec
+    for an unsupported type.
diff --git a/lib/sqlalchemy/dialects/postgresql/asyncpg.py b/lib/sqlalchemy/dialects/postgresql/asyncpg.py
index 3d195e691a..4e968a9b62 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))
-- 
2.47.3