Skip to content
Permalink

Comparing changes

Choose two branches to see what’s changed or to start a new pull request. If you need to, you can also or learn more about diff comparisons.

Open a pull request

Create a new pull request by comparing changes across two branches. If you need to, you can also . Learn more about diff comparisons here.
base repository: sqlalchemy/sqlalchemy
Failed to load repositories. Confirm that selected base ref is valid, then try again.
Loading
base: main
Choose a base ref
...
head repository: CobaltCause/sqlalchemy
Failed to load repositories. Confirm that selected head ref is valid, then try again.
Loading
compare: master
Choose a head ref
Can’t automatically merge. Don’t worry, you can still create the pull request.
  • 1 commit
  • 2 files changed
  • 1 contributor

Commits on Jul 29, 2021

  1. fix cockroachdb support via asyncpg 

    Fixes #6825
    @CobaltCause
    CobaltCause committed Jul 29, 2021
    Copy the full SHA
    9f208d4 View commit details
Showing with 52 additions and 18 deletions.
  1. +50 −16 lib/sqlalchemy/dialects/postgresql/asyncpg.py
  2. +2 −2 lib/sqlalchemy/dialects/postgresql/base.py
66 changes: 50 additions & 16 deletions lib/sqlalchemy/dialects/postgresql/asyncpg.py
View file
Original file line number Diff line number Diff line change
@@ -1012,7 +1012,7 @@ def _jsonb_decoder(bin_value):
# asyncpg returns it when format='binary'
return deserializer(bin_value[1:].decode())

async def _setup_type_codecs(conn):
async def _setup_type_codecs(db_name, db_major_ver, conn):
"""set up type decoders at the asyncpg level.
these are set_type_codec() calls to normalize
@@ -1023,23 +1023,57 @@ async def _setup_type_codecs(conn):
See https://github.com/MagicStack/asyncpg/issues/623 for reference
on why it's set up this way.
"""
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",
)

# CockroachDB only has `jsonb`
if db_name != "CockroachDB CCL":
await conn._connection.set_type_codec(
"json",
encoder=str.encode,
decoder=_json_decoder,
schema="pg_catalog",
format="binary",
)

try:
# This only works on DBs other than CockroachDB, other than
# versions greater than 21.1.0. This check is good enough
# because there also aren't any 21.0.* versions.
jsonb_support = (
db_name != "CockroachDB CCL"
or (
db_name == "CockroachDB CCL"
and int(db_major_ver) > 21
)
)
except ValueError:
# Short-circuit evaluation means this will only happen if the
# database is CockroachDB and the version is unparseable, in
# which case it's probably safe to assume that jsonb won't work
jsonb_support = False

if jsonb_support:
await conn._connection.set_type_codec(
"jsonb",
encoder=_jsonb_encoder,
decoder=_jsonb_decoder,
schema="pg_catalog",
format="binary",
)

def connect(conn):
conn.await_(_setup_type_codecs(conn))
f = conn._connection.fetch("select version()")
v = conn.await_(f)[0]["version"]
m = re.match(
r".*(PostgreSQL|EnterpriseDB|CockroachDB CCL) "
r"v?(\d+)\.?(?:\d+)?(?:\.(?:\d+))?(?:\.\d+)?(?:devel|beta)?",
v
)
if not m:
raise AssertionError(
"Could not determine database kind from string '%s'" % v
)

conn.await_(_setup_type_codecs(m.group(1), m.group(2), conn))
if super_connect is not None:
super_connect(conn)

4 changes: 2 additions & 2 deletions lib/sqlalchemy/dialects/postgresql/base.py
View file
Original file line number Diff line number Diff line change
@@ -3427,8 +3427,8 @@ def has_type(self, connection, type_name, schema=None):
def _get_server_version_info(self, connection):
v = connection.exec_driver_sql("select version()").scalar()
m = re.match(
r".*(?:PostgreSQL|EnterpriseDB) "
r"(\d+)\.?(\d+)?(?:\.(\d+))?(?:\.\d+)?(?:devel|beta)?",
r".*(?:PostgreSQL|EnterpriseDB|CockroachDB CCL) "
r"v?(\d+)\.?(\d+)?(?:\.(\d+))?(?:\.\d+)?(?:devel|beta)?",
v,
)
if not m: