Add preliminary support for asyncpg

Author: gordthompson

README.asyncpg.md

@@ -0,0 +1,34 @@
+## asyncpg support (experimental)
+
+Preliminary support for asyncpg has been added. The connection URI is of the form::
+
+    cockroachdb+asyncpg://root@localhost:26257/defaultdb
+
+The example code at
+
+https://docs.sqlalchemy.org/en/14/_modules/examples/asyncio/async_orm.html
+
+runs without error and seems to be producing reasonable results.
+
+There is also a customized version of the FastAPI SQL database tutorial for
+`cockroachdb+asyncpg` available at
+
+https://github.com/gordthompson/fastapi-tutorial-cockroachdb-async
+
+### Testing
+
+asyncpg support has not yet been added to CI testing. There is at least one outstanding issue
+that needs to be resolved before the Alembic test suite can be successfully run on an asyncpg
+connection:
+
+https://github.com/cockroachdb/cockroach/issues/71908
+
+So, if you want to run all of the tests *except* the Alembic tests then invoke pytest
+using a command like
+
+    pytest --db=asyncpg --ignore-glob='*test_suite_alembic.py'
+
+assuming that you have an entry in test.cfg that looks something like
+
+    [db]
+    asyncpg=cockroachdb+asyncpg://root@localhost:26257/defaultdb

README.md

@@ -2,31 +2,49 @@
 
 ## Prerequisites
 
-You must install either:
+For psycopg2 support you must install either:
 
 * [psycopg2](https://pypi.org/project/psycopg2/), which has some
   [prerequisites](https://www.psycopg.org/docs/install.html#prerequisites) of
   its own.
 
 * [psycopg2-binary](https://pypi.org/project/psycopg2-binary/)
 
-The binary package is a practical choice for development and testing but in
-production it is advised to use the package built from sources.
+(The binary package is a practical choice for development and testing but in
+production it is advised to use the package built from sources.)
 
+Or, for asyncpg support (experimental) you must install
+
+* [asyncpg](https://pypi.org/project/asyncpg/)
+ 
 ## Install and usage
 
 Use `pip` to install the latest version.
 
 `pip install sqlalchemy-cockroachdb`
 
 Use a `cockroachdb` connection string when creating the `Engine`. For example,
-to connect to an insecure, local CockroachDB cluster:
+to connect to an insecure, local CockroachDB cluster using psycopg2:
 
 ```
 from sqlalchemy import create_engine
 engine = create_engine('cockroachdb://root@localhost:26257/defaultdb?sslmode=disable')
 ```
 
+or
+
+```
+from sqlalchemy import create_engine
+engine = create_engine('cockroachdb+psycopg2://root@localhost:26257/defaultdb?sslmode=disable')
+```
+
+To connect using asyncpg:
+
+```
+from sqlalchemy import create_async_engine
+engine = create_async_engine('cockroachdb+asyncpg://root@localhost:26257/defaultdb')
+```
+
 ## Changelog
 
 See [CHANGES.md](https://github.com/cockroachdb/sqlalchemy-cockroachdb/blob/master/CHANGES.md)

dev-requirements.txt

@@ -1,58 +1,64 @@
-appdirs==1.4.4
+backports.entry-points-selectable==1.1.1
     # via virtualenv
-bleach==3.3.0
+bleach==4.1.0
     # via readme-renderer
-certifi==2021.5.30
+certifi==2021.10.8
     # via requests
-cffi==1.14.5
+cffi==1.15.0
     # via cryptography
-chardet==4.0.0
+charset-normalizer==2.0.7
     # via requests
 colorama==0.4.4
     # via twine
-cryptography==3.4.7
+cryptography==35.0.0
     # via secretstorage
-distlib==0.3.2
+distlib==0.3.3
     # via virtualenv
-docutils==0.17.1
+docutils==0.18
     # via readme-renderer
-filelock==3.0.12
+filelock==3.3.2
     # via
     #   tox
     #   virtualenv
-idna==2.10
+idna==3.3
     # via requests
-importlib-metadata==4.5.0
+importlib-metadata==4.8.2
     # via
+    #   backports.entry-points-selectable
     #   keyring
+    #   pluggy
+    #   tox
     #   twine
-jeepney==0.6.0
+    #   virtualenv
+jeepney==0.7.1
     # via
     #   keyring
     #   secretstorage
-keyring==23.0.1
+keyring==23.2.1
     # via twine
-packaging==20.9
+packaging==21.2
     # via
     #   bleach
     #   tox
-pkginfo==1.7.0
+pkginfo==1.7.1
     # via twine
-pluggy==0.13.1
+platformdirs==2.4.0
+    # via virtualenv
+pluggy==1.0.0
     # via tox
-py==1.10.0
+py==1.11.0
     # via tox
-pycparser==2.20
+pycparser==2.21
     # via cffi
-pygments==2.9.0
+pygments==2.10.0
     # via readme-renderer
 pyparsing==2.4.7
     # via packaging
-readme-renderer==29.0
+readme-renderer==30.0
     # via twine
 requests-toolbelt==0.9.1
     # via twine
-requests==2.25.1
+requests==2.26.0
     # via
     #   requests-toolbelt
     #   twine
@@ -63,22 +69,23 @@ secretstorage==3.3.1
 six==1.16.0
     # via
     #   bleach
-    #   readme-renderer
     #   tox
     #   virtualenv
 toml==0.10.2
     # via tox
-tox==3.23.1
+tox==3.24.4
     # via -r dev-requirements.in
-tqdm==4.61.1
+tqdm==4.62.3
     # via twine
-twine==3.4.1
+twine==3.6.0
     # via -r dev-requirements.in
-urllib3==1.26.5
+typing-extensions==3.10.0.2
+    # via importlib-metadata
+urllib3==1.26.7
     # via requests
-virtualenv==20.4.7
+virtualenv==20.10.0
     # via tox
 webencodings==0.5.1
     # via bleach
-zipp==3.4.1
+zipp==3.6.0
     # via importlib-metadata

setup.py

@@ -40,8 +40,9 @@
     zip_safe=False,
     entry_points={
         "sqlalchemy.dialects": [
-            "cockroachdb.psycopg2 = sqlalchemy_cockroachdb.psycopg2:CockroachDBDialect_psycopg2",
             "cockroachdb = sqlalchemy_cockroachdb.psycopg2:CockroachDBDialect_psycopg2",
+            "cockroachdb.psycopg2 = sqlalchemy_cockroachdb.psycopg2:CockroachDBDialect_psycopg2",
+            "cockroachdb.asyncpg = sqlalchemy_cockroachdb.asyncpg:CockroachDBDialect_asyncpg",
         ],
     },
 )

sqlalchemy_cockroachdb/__init__.py

@@ -8,3 +8,8 @@
     "sqlalchemy_cockroachdb.psycopg2",
     "CockroachDBDialect_psycopg2",
 )
+_registry.register(
+    "cockroachdb.asyncpg",
+    "sqlalchemy_cockroachdb.asyncpg",
+    "CockroachDBDialect_asyncpg",
+)

sqlalchemy_cockroachdb/asyncpg.py

@@ -0,0 +1,18 @@
+import json as _py_json
+
+from sqlalchemy.dialects.postgresql.asyncpg import PGDialect_asyncpg
+from .base import CockroachDBDialect
+from .ddl_compiler import CockroachDDLCompiler
+from .stmt_compiler import CockroachIdentifierPreparer
+
+
+class CockroachDBDialect_asyncpg(PGDialect_asyncpg, CockroachDBDialect):
+    driver = "asyncpg"  # driver name
+    preparer = CockroachIdentifierPreparer
+    ddl_compiler = CockroachDDLCompiler
+
+    supports_statement_cache = True
+
+    async def setup_asyncpg_json_codec(self, conn):
+        # https://github.com/cockroachdb/cockroach/issues/9990#issuecomment-579202144
+        pass

sqlalchemy_cockroachdb/base.py

@@ -6,6 +6,7 @@
 from sqlalchemy.dialects.postgresql import INET
 from sqlalchemy.dialects.postgresql import UUID
 from sqlalchemy.ext.compiler import compiles
+from sqlalchemy.util import asbool
 from sqlalchemy.util import warn
 import sqlalchemy.sql as sql
 
@@ -96,16 +97,11 @@ class CockroachDBDialect(PGDialect):
     # Override connect so we can take disable_cockroachdb_telemetry as a connect_arg to sqlalchemy.
     def connect(
         self,
-        dsn=None,
-        connection_factory=None,
-        cursor_factory=None,
         disable_cockroachdb_telemetry=False,
         **kwargs,
     ):
-        self.disable_cockroachdb_telemetry = disable_cockroachdb_telemetry
-        return super().connect(
-            dsn, connection_factory,
-            cursor_factory, **kwargs)
+        self.disable_cockroachdb_telemetry = asbool(disable_cockroachdb_telemetry)
+        return super().connect(**kwargs)
 
     def __init__(self, *args, **kwargs):
         if kwargs.get("use_native_hstore", False):

sqlalchemy_cockroachdb/psycopg2.py

@@ -5,7 +5,6 @@
 
 
 class CockroachDBDialect_psycopg2(PGDialect_psycopg2, CockroachDBDialect):
-    name = "cockroachdb"  # dialect name
     driver = "psycopg2"  # driver name
     preparer = CockroachIdentifierPreparer
     ddl_compiler = CockroachDDLCompiler

sqlalchemy_cockroachdb/requirements.py

@@ -70,10 +70,10 @@ class Requirements(SuiteRequirementsSQLA, SuiteRequirementsAlembic):
         lambda config: not config.db.dialect._is_v202plus,
         "older versions don't support this correctly.",
     )
-    # TODO: enable after 20.2 beta comes out
-    # check_constraint_reflection = \
-    #        exclusions.skip_if(lambda config: not config.db.dialect._is_v202plus,
-    #                           "older versions don't support this correctly.")
+    check_constraint_reflection = exclusions.skip_if(
+        lambda config: not config.db.dialect._is_v202plus,
+        "older versions don't support this correctly.",
+    )
     cross_schema_fk_reflection = exclusions.closed()
     non_updating_cascade = exclusions.open()
     deferrable_fks = exclusions.closed()
@@ -155,10 +155,16 @@ class Requirements(SuiteRequirementsSQLA, SuiteRequirementsAlembic):
     implicitly_named_constraints = exclusions.open()
     supports_distinct_on = exclusions.open()
 
+    @property
+    def sync_driver(self):
+        return exclusions.skip_if(
+            lambda config: config.db.dialect.is_async
+        )
+
     def get_isolation_levels(self, config):
         return {"default": "SERIALIZABLE", "supported": ["SERIALIZABLE"]}
 
-# non-default requirements for Alembic test suite
+    # non-default requirements for Alembic test suite
 
     @property
     def autoincrement_on_composite_pk(self):

test-requirements.in

@@ -5,6 +5,7 @@
 # generated test-requirements.txt) and run make update-requirements
 
 alembic
+asyncpg
 futures
 mock
 more-itertools